From f2521c4a5c228deea80b1e1c000fa661754719a7 Mon Sep 17 00:00:00 2001
From: Myk Taylor <myk002@yahoo.com>
Date: Fri, 3 Feb 2023 12:52:50 -0800
Subject: [PATCH] protect against bad tree data

---
 plugins/autochop.cpp | 116 +++++++++++++++++++++++++++----------------
 1 file changed, 73 insertions(+), 43 deletions(-)

diff --git a/plugins/autochop.cpp b/plugins/autochop.cpp
index 62433a1f0..5b1cdf97e 100644
--- a/plugins/autochop.cpp
+++ b/plugins/autochop.cpp
@@ -327,16 +327,23 @@ static bool is_protected(const df::plant * plant, PersistentDataItem &c) {
 }
 
 static int32_t estimate_logs(const df::plant *plant) {
+    if (!plant->tree_info)
+        return 0;
+
     //adapted from code by aljohnston112 @ github
     df::plant_tree_tile** tiles = plant->tree_info->body;
-    df::plant_tree_tile* tilesRow;
 
-    int trunks = 0;
+    if (!tiles)
+        return 0;
+
+    int32_t trunks = 0;
+    const int32_t area = plant->tree_info->dim_y * plant->tree_info->dim_x;
     for (int i = 0; i < plant->tree_info->body_height; i++) {
-        tilesRow = tiles[i];
-        for (int j = 0; j < plant->tree_info->dim_y*plant->tree_info->dim_x; j++) {
+        df::plant_tree_tile* tilesRow = tiles[i];
+        if (!tilesRow)
+            return 0; // tree data is corrupt; let's not touch it
+        for (int j = 0; j < area; j++)
             trunks += tilesRow[j].bits.trunk;
-        }
     }
 
     return trunks;
@@ -386,6 +393,55 @@ static void bucket_watched_burrows(color_ostream & out,
 
 typedef multimap<int, df::plant *, std::greater<int>> TreesBySize;
 
+static int32_t scan_tree(color_ostream & out, df::plant *plant, int32_t *expected_yield,
+        TreesBySize *designatable_trees_by_size, bool designate_clearcut,
+        const vector<df::unit *> &citizens, int32_t *accessible_trees,
+        int32_t *inaccessible_trees, int32_t *designated_trees, int32_t *accessible_yield,
+        map<int32_t, int32_t> *tree_counts,
+        map<int32_t, int32_t> *designated_tree_counts,
+        map<int, PersistentDataItem *> &clearcut_burrows,
+        map<int, PersistentDataItem *> &chop_burrows) {
+    TRACE(cycle,out).print("  scanning tree at %d,%d,%d\n",
+                            plant->pos.x, plant->pos.y, plant->pos.z);
+
+    if (!is_valid_tree(plant))
+        return 0;
+
+    bool accessible = is_accessible_tree(plant->pos, citizens);
+    int32_t yield = estimate_logs(plant);
+
+    if (accessible) {
+        if (accessible_trees)
+            ++*accessible_trees;
+        if (accessible_yield)
+            *accessible_yield += yield;
+    } else {
+        if (inaccessible_trees)
+            ++*inaccessible_trees;
+    }
+
+    bool can_chop = false;
+    bool designated = Designations::isPlantMarked(plant);
+    bool was_designated = designated;
+    bucket_tree(plant, designate_clearcut, &designated, &can_chop, tree_counts,
+            designated_tree_counts, clearcut_burrows, chop_burrows);
+
+    int32_t ret = 0;
+    if (designated) {
+        if (!was_designated)
+            ret = 1;
+        if (designated_trees)
+            ++*designated_trees;
+        if (expected_yield)
+            *expected_yield += yield;
+    } else if (can_chop && accessible) {
+        if (designatable_trees_by_size)
+            designatable_trees_by_size->emplace(yield, plant);
+    }
+
+    return ret;
+}
+
 // returns the number of trees that were newly marked
 static int32_t scan_trees(color_ostream & out, int32_t *expected_yield,
         TreesBySize *designatable_trees_by_size, bool designate_clearcut,
@@ -415,44 +471,18 @@ static int32_t scan_trees(color_ostream & out, int32_t *expected_yield,
     map<int, PersistentDataItem *> clearcut_burrows, chop_burrows;
     bucket_watched_burrows(out, clearcut_burrows, chop_burrows);
 
-    for (auto plant : world->plants.all) {
-        TRACE(cycle,out).print("  scanning tree at %d,%d,%d\n",
-                               plant->pos.x, plant->pos.y, plant->pos.z);
-
-        if (!is_valid_tree(plant))
-            continue;
-
-        bool accessible = is_accessible_tree(plant->pos, citizens);
-        int32_t yield = estimate_logs(plant);
-
-        if (accessible) {
-            if (accessible_trees)
-                ++*accessible_trees;
-            if (accessible_yield)
-                *accessible_yield += yield;
-        } else {
-            if (inaccessible_trees)
-                ++*inaccessible_trees;
-        }
-
-        bool can_chop = false;
-        bool designated = Designations::isPlantMarked(plant);
-        bool was_designated = designated;
-        bucket_tree(plant, designate_clearcut, &designated, &can_chop, tree_counts,
-                designated_tree_counts, clearcut_burrows, chop_burrows);
-
-        if (designated) {
-            if (!was_designated)
-                ++newly_marked;
-            if (designated_trees)
-                ++*designated_trees;
-            if (expected_yield)
-                *expected_yield += yield;
-        } else if (can_chop && accessible) {
-            if (designatable_trees_by_size)
-                designatable_trees_by_size->emplace(yield, plant);
-        }
-    }
+    for (auto plant : world->plants.tree_dry)
+        newly_marked += scan_tree(out, plant, expected_yield, designatable_trees_by_size,
+                                  designate_clearcut, citizens, accessible_trees,
+                                  inaccessible_trees, designated_trees, accessible_yield,
+                                  tree_counts, designated_tree_counts,
+                                  clearcut_burrows, chop_burrows);
+    for (auto plant : world->plants.tree_wet)
+        newly_marked += scan_tree(out, plant, expected_yield, designatable_trees_by_size,
+                                  designate_clearcut, citizens, accessible_trees,
+                                  inaccessible_trees, designated_trees, accessible_yield,
+                                  tree_counts, designated_tree_counts,
+                                  clearcut_burrows, chop_burrows);
 
     return newly_marked;
 }