add Military::addToSquad

docs/changelog.txt

@@ -62,8 +62,10 @@ Template for new versions:
 ## Documentation
 
 ## API
+- ``Military`` module: added ``addToSquad`` function
 
 ## Lua
+- ``dfhack.military.addToSquad``: expose Military API function
 
 ## Removed
 

docs/dev/Lua API.rst

@@ -1995,6 +1995,23 @@ Military module
   to indicate former squad membership or command, and creates a corresponding
   world history event.
 
+  * ``dfhack.military.addToSquad(unit_id, squad_id, squad_pos)``
+
+  Adds a unit to a squad. Sets 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``), adds a unit's entity links to
+  indicate squad membership. Does not currently add world history events.
+  If ``squad_pos`` is -1, the unit will be added to the first open slot in
+  the squad.
+
+  This API cannot be used to set or change the leader of a squad and will fail
+  if ``squad_pos`` is specified as 0 or if ``squad_pos`` is specified as -1 and
+  the squad leader position is currently vacant. It will also fail if
+  the requested squad position is already occupied, the squad does not exist,
+  the unit does not exist, or the requested unit is already a member of another
+  squad.
+
 Items module
 ------------
 

library/LuaApi.cpp

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

library/include/modules/Military.h

@@ -18,6 +18,7 @@ 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);
+DFHACK_EXPORT bool addToSquad(int32_t unit_id, int32_t squad_id, int32_t squad_pos = -1);
 
 }
 }

library/modules/Military.cpp

@@ -408,6 +408,61 @@ static void remove_officer_entity_link(df::historical_figure* hf, df::squad* squ
     df::global::world->history.events.push_back(former_pos_event);
 }
 
+static void add_soldier_entity_link(df::historical_figure* hf, df::squad* squad, int32_t squad_pos)
+{
+    auto squad_link = df::allocate<df::histfig_entity_link_squadst>();
+    squad_link->squad_id = squad->id;
+    squad_link->squad_position = squad_pos;
+    squad_link->entity_id = squad->entity_id;
+    squad_link->start_year = *df::global::cur_year;
+    squad_link->link_strength = 100;
+
+    hf->entity_links.push_back(squad_link);
+}
+
+bool Military::addToSquad(int32_t unit_id, int32_t squad_id, int32_t squad_pos)
+{
+    df::unit* unit = df::unit::find(unit_id);
+    if (unit == nullptr || unit->military.squad_id != -1) return false;
+
+    df::historical_figure* hf = df::historical_figure::find(unit->hist_figure_id);
+    if (hf == nullptr)
+        return false;
+
+    df::squad* squad = df::squad::find(squad_id);
+    if (squad == nullptr) return false;
+
+    if (squad_pos == -1)
+    {
+        for (int p = 0; p < 10; p++)
+        {
+            auto pp = vector_get(squad->positions, p);
+            if (pp == nullptr || pp->occupant == -1)
+            {
+                squad_pos = p;
+                break;
+            }
+        }
+    }
+    if (squad_pos == -1) return false;
+
+    // this function cannot (currently) change the squad commander
+    if (squad_pos == 0) return false;
+
+    df::squad_position* pos = vector_get(squad->positions, squad_pos);
+    if (pos == nullptr)
+        pos = squad->positions[squad_pos] = df::allocate<df::squad_position>();
+
+    pos->occupant = hf->id;
+    // does anything else need to be set here?
+
+    unit->military.squad_id = squad->id;
+    unit->military.squad_position = squad_pos;
+
+    add_soldier_entity_link(hf, squad, squad_pos);
+    return true;
+}
+
 bool Military::removeFromSquad(int32_t unit_id)
 {
     df::unit *unit = df::unit::find(unit_id);