Merge pull request #1490 from Jarkami/uniform-unstick-fix

ab9rf · web-flow · commit f585c746dc39 · 2025-08-18T09:29:37.000-05:00
Fix state inconsistency in uniform-unstick
diff --git a/changelog.txt b/changelog.txt
@@ -44,6 +44,7 @@ Template for new versions:
 - `confirm`: when removing a manager order, show specific item/job type for ammo, shield, helm, gloves, shoes, trap component, and meal orders
 - `confirm`: the pause option now pauses individual confirmation types, allowing multiple different confirmations to be paused independently
 - `immortal-cravings`: prioritize high-value meals, properly split of portions, and don't go eating or drinking on a full stomach
+- `uniform-unstick`: no longer causes units to equip multiples of assigned items
 
 ## Misc Improvements
 - `devel/hello-world`: updated to show off the new Slider widget
diff --git a/uniform-unstick.lua b/uniform-unstick.lua
@@ -15,11 +15,15 @@ local validArgs = utils.invert({
 
 -- Functions
 
+-- @param item df.item
+-- @return string
 local function item_description(item)
-    return dfhack.df2console(dfhack.items.getDescription(item, 0, true))
+    return "item #" .. item.id .. " '" .. dfhack.df2console(dfhack.items.getDescription(item, 0, true)) .. "'"
 end
 
-local function get_item_pos(item)
+-- @param item df.item
+-- @return df.coord|nil
+local function get_visible_item_pos(item)
     local x, y, z = dfhack.items.getPosition(item)
     if not x or not y or not z then
         return
@@ -30,24 +34,30 @@ local function get_item_pos(item)
     end
 end
 
-local function get_squad_position(unit, unit_name)
+-- @param unit df.unit
+-- @return df.squad_position|nil
+local function get_squad_position(unit)
     local squad = df.squad.find(unit.military.squad_id)
-    if squad then
-        if squad.entity_id ~= df.global.plotinfo.group_id then
-            print("WARNING: Unit " .. unit_name .. " is a member of a squad from another site!" ..
-                " This may be preventing them from doing any useful work." ..
-                " You can fix this by assigning them to a local squad and then unassigning them.")
-            print()
-            return
-        end
-    else
+    if not squad then
+        return
+    end
+
+    if squad.entity_id ~= df.global.plotinfo.group_id then
+        print("WARNING: Unit " .. dfhack.df2console(dfhack.units.getReadableName(unit)) .. " is a member of a squad from another site!" ..
+            " This may be preventing them from doing any useful work." ..
+            " You can fix this by assigning them to a local squad and then unassigning them.")
+        print()
         return
     end
+
     if #squad.positions > unit.military.squad_position then
         return squad.positions[unit.military.squad_position]
     end
 end
 
+-- @param unit df.unit
+-- @param item df.item
+-- @return number[] list of body part ids
 local function bodyparts_that_can_wear(unit, item)
     local bodyparts = {}
     local unitparts = dfhack.units.getCasteRaw(unit).body_info.body_parts
@@ -89,47 +99,61 @@ local function bodyparts_that_can_wear(unit, item)
     return bodyparts
 end
 
--- returns new value of need_newline
-local function print_line(text, need_newline)
-    if need_newline then
-        print()
-    end
-    print(text)
-    return false
+-- @param unit_name string
+-- @param labor_name string
+local function print_bad_labor(unit_name, labor_name)
+    return print("WARNING: Unit " .. unit_name .. " has the " .. labor_name ..
+        " labor enabled, which conflicts with military uniforms.")
 end
 
-local function print_bad_labor(unit_name, labor_name, need_newline)
-    return print_line("WARNING: Unit " .. unit_name .. " has the " .. labor_name ..
-        " labor enabled, which conflicts with military uniforms.", need_newline)
+-- @param squad_position df.squad_position
+-- @param item_id number
+local function remove_item_from_position(squad_position, item_id)
+    for _, uniform_slot_specs in ipairs(squad_position.equipment.uniform) do
+        for _, uniform_spec in ipairs(uniform_slot_specs) do
+            for idx, assigned_item_id in ipairs(uniform_spec.assigned) do
+                if assigned_item_id == item_id then
+                    uniform_spec.assigned:erase(idx)
+                    return
+                end
+            end
+        end
+    end
 end
 
 -- Will figure out which items need to be moved to the floor, returns an item_id:item map
-local function process(unit, args, need_newline)
+local function process(unit, args)
     local silent = args.all -- Don't print details if we're iterating through all dwarves
     local unit_name = dfhack.df2console(dfhack.units.getReadableName(unit))
+    local printed = false
 
     if not silent then
-        need_newline = print_line("Processing unit " .. unit_name, need_newline)
+        print("Processing unit " .. unit_name)
+        printed = true
     end
 
     -- The return value
     local to_drop = {} -- item id to item object
 
     -- First get squad position for an early-out for non-military dwarves
-    local squad_position = get_squad_position(unit, unit_name)
+    local squad_position = get_squad_position(unit)
     if not squad_position then
         if not silent then
-            need_newline = print_line(unit_name .. " does not have a military uniform.", need_newline)
+            print(unit_name .. " does not have a military uniform.")
+            print()
         end
         return
     end
 
     if unit.status.labors.MINE then
-        need_newline = print_bad_labor(unit_name, "mining", need_newline)
+        print_bad_labor(unit_name, "mining")
+        printed = true
     elseif unit.status.labors.CUTWOOD then
-        need_newline = print_bad_labor(unit_name, "woodcutting", need_newline)
+        print_bad_labor(unit_name, "woodcutting")
+        printed = true
     elseif unit.status.labors.HUNT then
-        need_newline = print_bad_labor(unit_name, "hunting", need_newline)
+        print_bad_labor(unit_name, "hunting")
+        printed = true
     end
 
     -- Find all worn items which may be at issue.
@@ -148,12 +172,12 @@ local function process(unit, args, need_newline)
     end
 
     -- Now get info about which items have been assigned as part of the uniform
-    local assigned_items = {} -- assigned item ids mapped to item objects
-    for _, specs in ipairs(squad_position.equipment.uniform) do
-        for _, spec in ipairs(specs) do
-            for _, assigned in ipairs(spec.assigned) do
+    local uniform_assigned_items = {} -- assigned item ids mapped to item objects
+    for _, uniform_slot_specs in ipairs(squad_position.equipment.uniform) do
+        for _, uniform_spec in ipairs(uniform_slot_specs) do
+            for _, assigned_item_id in ipairs(uniform_spec.assigned) do
                 -- Include weapon and shield so we can avoid dropping them, or pull them out of container/inventory later
-                assigned_items[assigned] = df.item.find(assigned)
+                uniform_assigned_items[assigned_item_id] = df.item.find(assigned_item_id)
             end
         end
     end
@@ -163,36 +187,49 @@ local function process(unit, args, need_newline)
 
     local present_ids = {} -- map of item ID to item object
     local missing_ids = {} -- map of item ID to item object
-    for u_id, item in pairs(assigned_items) do
-        if not worn_items[u_id] then
+    for item_id, item in pairs(uniform_assigned_items) do
+        if not worn_items[item_id] then
             if not silent then
-                need_newline = print_line(unit_name .. " is missing an assigned item, object #" .. u_id .. " '" ..
-                    item_description(item) .. "'", need_newline)
+                print(unit_name .. " is missing an assigned item, " .. item_description(item))
+                printed = true
             end
             if dfhack.items.getGeneralRef(item, df.general_ref_type.UNIT_HOLDER) then
-                need_newline = print_line(unit_name .. " cannot equip item: another unit has a claim on object #" .. u_id .. " '" .. item_description(item) .. "'", need_newline)
+                print(unit_name .. " cannot equip item: another unit has a claim on " .. item_description(item))
+                printed = true
                 if args.free then
                     print("  Removing from uniform")
-                    assigned_items[u_id] = nil
-                    for _, specs in ipairs(squad_position.equipment.uniform) do
-                        for _, spec in ipairs(specs) do
-                            for idx, assigned in ipairs(spec.assigned) do
-                                if assigned == u_id then
-                                    spec.assigned:erase(idx)
-                                    break
-                                end
-                            end
-                        end
-                    end
+                    uniform_assigned_items[item_id] = nil
+                    remove_item_from_position(squad_position, item_id)
                 end
             else
-                missing_ids[u_id] = item
+                missing_ids[item_id] = item
                 if args.free then
-                    to_drop[u_id] = item
+                    to_drop[item_id] = item
                 end
             end
         else
-            present_ids[u_id] = item
+            present_ids[item_id] = item
+        end
+    end
+
+    -- Make the equipment.assigned_items list consistent with what is present in equipment.uniform
+    for i=#(squad_position.equipment.assigned_items)-1,0,-1 do
+        local assigned_item_id = squad_position.equipment.assigned_items[i]
+        -- Quiver, backpack, and flask are assigned in their own locations rather than in equipment.uniform, and thus need their own checks
+        -- If more separately-assigned items are added in the future, this handling will need to be updated accordingly
+        if uniform_assigned_items[assigned_item_id] == nil and
+            assigned_item_id ~= squad_position.equipment.quiver and
+            assigned_item_id ~= squad_position.equipment.backpack and
+            assigned_item_id ~= squad_position.equipment.flask
+        then
+            local item = df.item.find(assigned_item_id)
+            if item ~= nil then
+                print(unit_name .. " has an improperly assigned item, " .. item_description(item) .. "; removing it")
+            else
+                print(unit_name .. " has a nonexistent item assigned, item # " .. assigned_item_id .. "; removing it")
+            end
+            printed = true
+            squad_position.equipment.assigned_items:erase(i)
         end
     end
 
@@ -202,10 +239,10 @@ local function process(unit, args, need_newline)
     -- unless --multi is specified, in which we don't care
     local covered = {} -- map of body part id to true/nil
     if not args.multi then
-        for id, item in pairs(present_ids) do
+        for item_id, item in pairs(present_ids) do
             -- weapons and shields don't "cover" the bodypart they're assigned to. (Needed to figure out if we're missing gloves.)
             if item._type ~= df.item_weaponst and item._type ~= df.item_shieldst then
-                covered[worn_parts[id]] = true
+                covered[worn_parts[item_id]] = true
             end
         end
     end
@@ -221,19 +258,23 @@ local function process(unit, args, need_newline)
     end
 
     -- Drop everything (except uniform pieces) from body parts which should be covered but aren't
-    for w_id, item in pairs(worn_items) do
-        if assigned_items[w_id] == nil then -- don't drop uniform pieces (including shields, weapons for hands)
-            if uncovered[worn_parts[w_id]] then
-                need_newline = print_line(unit_name ..
-                    " potentially has object #" ..
-                    w_id .. " '" .. item_description(item) .. "' blocking a missing uniform item.", need_newline)
+    for worn_item_id, item in pairs(worn_items) do
+        if uniform_assigned_items[worn_item_id] == nil then -- don't drop uniform pieces (including shields, weapons for hands)
+            if uncovered[worn_parts[worn_item_id]] then
+                print(unit_name .. " potentially has " .. item_description(item) .. " blocking a missing uniform item.")
+                printed = true
                 if args.drop then
-                    to_drop[w_id] = item
+                    to_drop[worn_item_id] = item
                 end
             end
         end
     end
 
+    -- add a spacing line if there was any output
+    if printed then
+        print()
+    end
+
     return to_drop
 end
 
@@ -242,15 +283,15 @@ local function do_drop(item_list)
         return
     end
 
-    for id, item in pairs(item_list) do
-        local pos = get_item_pos(item)
+    for _, item in pairs(item_list) do
+        local pos = get_visible_item_pos(item)
         if not pos then
-            dfhack.printerr("Could not find drop location for item #" .. id .. "  " .. item_description(item))
+            dfhack.printerr("Could not find drop location for " .. item_description(item))
         else
             if dfhack.items.moveToGround(item, pos) then
-                print("Dropped item #" .. id .. " '" .. item_description(item) .. "'")
+                print("Dropped " .. item_description(item))
             else
-                dfhack.printerr("Could not drop object #" .. id .. "  " .. item_description(item))
+                dfhack.printerr("Could not drop " .. item_description(item))
             end
         end
     end
@@ -265,10 +306,8 @@ local function main(args)
     end
 
     if args.all then
-        local need_newline = false
         for _, unit in ipairs(dfhack.units.getCitizens(true)) do
-            do_drop(process(unit, args, need_newline))
-            need_newline = true
+            do_drop(process(unit, args))
         end
     else
         local unit = dfhack.gui.getSelectedUnit()
