Merge pull request #5256 from ong-yinggao98/squad-features

Add function to remove unit from squad

Author: myk002

docs/changelog.txt

@@ -35,9 +35,11 @@ Template for new versions:
 ## Documentation
 
 ## API
+- ``Military::removeFromSquad``: removes unit from any squad assignments
 
 ## Lua
 - ``dfhack.units.setAutomaticProfessions``: sets unit labors according to current work detail settings
+- ``dfhack.military.removeFromSquad``: Lua API for ``Military::removeFromSquad``
 
 ## Removed
 

docs/dev/Lua API.rst

@@ -1986,6 +1986,14 @@ Military module
 
   Returns the name of a squad as a string.
 
+* ``dfhack.military.removeFromSquad(unit_id)``
+
+  Removes a unit from its squad. Unsets the unit's
+  military information (i.e., ``unit.military.squad_id`` and
+  ``unit.military.squad_pos``), the squad's position information (i.e.,
+  ``squad.positions[squad_pos].occupant``), and modifies the unit's entity links
+  to indicate former squad membership or command.
+
 Items module
 ------------
 

library/LuaApi.cpp

@@ -2341,6 +2341,7 @@ static const LuaWrapper::FunctionReg dfhack_military_module[] = {
     WRAPM(Military, makeSquad),
     WRAPM(Military, updateRoomAssignments),
     WRAPM(Military, getSquadName),
+    WRAPM(Military, removeFromSquad),
     { NULL, NULL }
 };
 

library/include/modules/Military.h

@@ -17,6 +17,7 @@ namespace Military
 DFHACK_EXPORT std::string getSquadName(int32_t squad_id);
 DFHACK_EXPORT df::squad* makeSquad(int32_t assignment_id);
 DFHACK_EXPORT void updateRoomAssignments(int32_t squad_id, int32_t civzone_id, df::squad_use_flags flags);
+DFHACK_EXPORT bool removeFromSquad(int32_t unit_id);
 
 }
 }

library/modules/Military.cpp

@@ -4,8 +4,13 @@
 #include "MiscUtils.h"
 #include "modules/Military.h"
 #include "modules/Translation.h"
+#include "modules/Units.h"
 #include "df/building.h"
 #include "df/building_civzonest.h"
+#include "df/histfig_entity_link_former_positionst.h"
+#include "df/histfig_entity_link_former_squadst.h"
+#include "df/histfig_entity_link_positionst.h"
+#include "df/histfig_entity_link_squadst.h"
 #include "df/historical_figure.h"
 #include "df/historical_entity.h"
 #include "df/entity_position.h"
@@ -16,6 +21,7 @@
 #include "df/squad_schedule_order.h"
 #include "df/squad_order.h"
 #include "df/squad_order_trainst.h"
+#include "df/unit.h"
 #include "df/world.h"
 
 using namespace DFHack;
@@ -289,3 +295,125 @@ void Military::updateRoomAssignments(int32_t squad_id, int32_t civzone_id, df::s
         }
     }
 }
+
+static void remove_soldier_entity_link(df::historical_figure* hf, df::squad* squad)
+{
+    int32_t start_year = -1;
+    for (size_t i = 0; i < hf->entity_links.size(); i++)
+    {
+        df::histfig_entity_link* link = hf->entity_links[i];
+        if (link->getType() != df::enums::histfig_entity_link_type::SQUAD)
+            continue;
+
+        auto squad_link = strict_virtual_cast<df::histfig_entity_link_squadst>(link);
+        if (squad_link == nullptr || squad_link->squad_id != squad->id)
+            continue;
+
+        hf->entity_links.erase(hf->entity_links.begin() + i);
+        start_year = squad_link->start_year;
+
+        delete squad_link;
+        break;
+    }
+
+    if (start_year == -1)
+        return;
+
+    auto former_squad = df::allocate<df::histfig_entity_link_former_squadst>();
+    former_squad->squad_id = squad->id;
+    former_squad->entity_id = squad->entity_id;
+    former_squad->start_year = start_year;
+    former_squad->end_year = *df::global::cur_year;
+    former_squad->link_strength = 100;
+
+    hf->entity_links.push_back(former_squad);
+}
+
+static void remove_officer_entity_link(df::historical_figure* hf, df::squad* squad)
+{
+    std::vector<Units::NoblePosition> nps;
+    if (!Units::getNoblePositions(&nps, hf))
+        return;
+
+    int32_t assignment_id = -1;
+    for (auto& np : nps)
+    {
+        if (np.entity->id != squad->entity_id || np.assignment->squad_id != squad->id)
+            continue;
+
+        np.assignment->histfig = -1;
+        np.assignment->histfig2 = -1;
+
+        assignment_id = np.assignment->id;
+        break;
+    }
+
+    if (assignment_id == -1)
+        return;
+
+    int32_t start_year = -1;
+    for (size_t i = 0; i < hf->entity_links.size(); i++)
+    {
+        df::histfig_entity_link* link = hf->entity_links[i];
+        if (link->getType() != df::enums::histfig_entity_link_type::POSITION)
+            continue;
+
+        auto pos_link = strict_virtual_cast<df::histfig_entity_link_positionst>(link);
+        if (pos_link == nullptr)
+            continue;
+        if (pos_link->assignment_id != assignment_id && pos_link->entity_id != squad->entity_id)
+            continue;
+
+        hf->entity_links.erase(hf->entity_links.begin() + i);
+        start_year = pos_link->start_year;
+
+        delete pos_link;
+        break;
+    }
+
+    if (start_year == -1)
+        return;
+
+    auto former_pos = df::allocate<df::histfig_entity_link_former_positionst>();
+    former_pos->assignment_id = assignment_id;
+    former_pos->entity_id = squad->entity_id;
+    former_pos->start_year = start_year;
+    former_pos->end_year = *df::global::cur_year;
+    former_pos->link_strength = 100;
+
+    hf->entity_links.push_back(former_pos);
+}
+
+bool Military::removeFromSquad(int32_t unit_id)
+{
+    df::unit *unit = df::unit::find(unit_id);
+    if (unit == nullptr || unit->military.squad_id == -1 || unit->military.squad_position == -1)
+        return false;
+
+    int32_t squad_id = unit->military.squad_id;
+    df::squad* squad = df::squad::find(squad_id);
+    if (squad == nullptr)
+        return false;
+
+    int32_t squad_pos = unit->military.squad_position;
+    df::squad_position* pos = vector_get(squad->positions, squad_pos);
+    if (pos == nullptr)
+        return false;
+
+    df::historical_figure* hf = df::historical_figure::find(unit->hist_figure_id);
+    if (hf == nullptr)
+        return false;
+
+    // remove from squad information
+    pos->occupant = -1;
+    // remove from unit information
+    unit->military.squad_id = -1;
+    unit->military.squad_position = -1;
+
+    if (squad_pos == 0) // is unit a commander?
+        remove_officer_entity_link(hf, squad);
+    else
+        remove_soldier_entity_link(hf, squad);
+
+    return true;
+}