Designate the correct tile for trees with multi-tile trunks

DF always designates the southeast trunk tile, while plant.pos is the center, which is sometimes inaccessible (see DFHack#1014)
maxthyme · May 6, 2017 · 5593170 · 5593170
1 parent fbae598
commit 5593170
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 10 deletions.
diff --git a/library/LuaApi.cpp b/library/LuaApi.cpp
@@ -2328,6 +2328,16 @@ static const LuaWrapper::FunctionReg dfhack_designations_module[] = {
     {NULL, NULL}
 };
 
+static int designations_getPlantDesignationTile(lua_State *state)
+{
+    return Lua::PushPosXYZ(state, Designations::getPlantDesignationTile(Lua::CheckDFObject<df::plant>(state, 1)));
+}
+
+static const luaL_Reg dfhack_designations_funcs[] = {
+    {"getPlantDesignationTile", designations_getPlantDesignationTile},
+    {NULL, NULL}
+};
+
 /***** Internal module *****/
 
 static void *checkaddr(lua_State *L, int idx, bool allow_null = false)
@@ -2806,6 +2816,6 @@ void OpenDFHackApi(lua_State *state)
     OpenModule(state, "constructions", dfhack_constructions_module);
     OpenModule(state, "screen", dfhack_screen_module, dfhack_screen_funcs);
     OpenModule(state, "filesystem", dfhack_filesystem_module, dfhack_filesystem_funcs);
-    OpenModule(state, "designations", dfhack_designations_module);
+    OpenModule(state, "designations", dfhack_designations_module, dfhack_designations_funcs);
     OpenModule(state, "internal", dfhack_internal_module, dfhack_internal_funcs);
 }
diff --git a/library/include/modules/Designations.h b/library/include/modules/Designations.h
@@ -14,5 +14,8 @@ namespace DFHack {
         DFHACK_EXPORT bool canMarkPlant(const df::plant *plant);
         DFHACK_EXPORT bool canUnmarkPlant(const df::plant *plant);
         DFHACK_EXPORT bool isPlantMarked(const df::plant *plant);
+
+        // Return the tile that should be designated for this plant
+        DFHACK_EXPORT df::coord getPlantDesignationTile(const df::plant *plant);
     }
 }
diff --git a/library/modules/Designations.cpp b/library/modules/Designations.cpp
@@ -8,6 +8,8 @@
 #include "df/job.h"
 #include "df/map_block.h"
 #include "df/plant.h"
+#include "df/plant_tree_info.h"
+#include "df/plant_tree_tile.h"
 #include "df/tile_dig_designation.h"
 #include "df/world.h"
 
@@ -20,18 +22,48 @@ static df::map_block *getPlantBlock(const df::plant *plant)
 {
     if (!world)
         return nullptr;
-    return Maps::getTileBlock(plant->pos);
+    return Maps::getTileBlock(Designations::getPlantDesignationTile(plant));
+}
+
+df::coord Designations::getPlantDesignationTile(const df::plant *plant)
+{
+    CHECK_NULL_POINTER(plant);
+
+    if (!plant->tree_info)
+        return plant->pos;
+
+    int dimx = plant->tree_info->dim_x;
+    int dimy = plant->tree_info->dim_y;
+    int cx = dimx / 2;
+    int cy = dimy / 2;
+
+    // Find the southeast trunk tile
+    int x = cx;
+    int y = cy;
+
+    while (x + 1 < dimx && y + 1 < dimy)
+    {
+        if (plant->tree_info->body[0][(y * dimx) + (x + 1)].bits.trunk)
+            ++x;
+        else if (plant->tree_info->body[0][((y + 1) * dimx) + x].bits.trunk)
+            ++y;
+        else
+            break;
+    }
+
+    return df::coord(plant->pos.x - cx + x, plant->pos.y - cy + y, plant->pos.z);
 }
 
 bool Designations::isPlantMarked(const df::plant *plant)
 {
     CHECK_NULL_POINTER(plant);
 
-    df::map_block *block = getPlantBlock(plant);
+    df::coord des_pos = getPlantDesignationTile(plant);
+    df::map_block *block = Maps::getTileBlock(des_pos);
     if (!block)
         return false;
 
-    if (block->designation[plant->pos.x % 16][plant->pos.y % 16].bits.dig == tile_dig_designation::Default)
+    if (block->designation[des_pos.x % 16][des_pos.y % 16].bits.dig == tile_dig_designation::Default)
         return true;
 
     for (auto *link = world->job_list.next; link; link = link->next)
@@ -41,7 +73,7 @@ bool Designations::isPlantMarked(const df::plant *plant)
             continue;
         if (job->job_type != job_type::FellTree && job->job_type != job_type::GatherPlants)
             continue;
-        if (job->pos == plant->pos)
+        if (job->pos == des_pos)
             return true;
     }
     return false;
@@ -63,8 +95,9 @@ bool Designations::markPlant(const df::plant *plant)
 
     if (canMarkPlant(plant))
     {
-        df::map_block *block = getPlantBlock(plant);
-        block->designation[plant->pos.x % 16][plant->pos.y % 16].bits.dig = tile_dig_designation::Default;
+        df::coord des_pos = getPlantDesignationTile(plant);
+        df::map_block *block = Maps::getTileBlock(des_pos);
+        block->designation[des_pos.x % 16][des_pos.y % 16].bits.dig = tile_dig_designation::Default;
         block->flags.bits.designated = true;
         return true;
     }
@@ -90,8 +123,9 @@ bool Designations::unmarkPlant(const df::plant *plant)
 
     if (canUnmarkPlant(plant))
     {
-        df::map_block *block = getPlantBlock(plant);
-        block->designation[plant->pos.x % 16][plant->pos.y % 16].bits.dig = tile_dig_designation::No;
+        df::coord des_pos = getPlantDesignationTile(plant);
+        df::map_block *block = Maps::getTileBlock(des_pos);
+        block->designation[des_pos.x % 16][des_pos.y % 16].bits.dig = tile_dig_designation::No;
         block->flags.bits.designated = true;
 
         auto *link = world->job_list.next;
@@ -102,7 +136,7 @@ bool Designations::unmarkPlant(const df::plant *plant)
 
             if (job &&
                 (job->job_type == job_type::FellTree || job->job_type == job_type::GatherPlants) &&
-                job->pos == plant->pos)
+                job->pos == des_pos)
             {
                 Job::removeJob(job);
             }