Units::teleport: keep occupancy flags while other units remain on the tile

docs/changelog.txt

@@ -68,6 +68,7 @@ Template for new versions:
 - In ``Screen`` module, Fix out-of-bounds color table access when Lua pens use ``COLOR_RESET``
 - rename the ``stockpiles`` tool tag to ``stockpile`` so it no longer collides with the `stockpiles` plugin in `gui/launcher`
 - ``Units::teleport``: update unit occupancy on the full 3x3 footprint of EQUIPMENT units (e.g. wagons) instead of only their center tile
+- ``Units::teleport``: only clear tile unit occupancy flags when no other unit of the same kind remains on the tile
 - `3dveins`: fix failure on embarks that cross midmap tiles
 - `aquifer`: make ``--skip-top`` and top-relative ``--levels`` take effect for the ``drain``, ``convert``, and ``add`` actions instead of being silently ignored
 - `autodump`: ``destroy`` no longer leaves the contents of destroyed containers in limbo or crashes when destroying unit-held items

library/modules/Units.cpp

@@ -100,6 +100,7 @@ distribution.
 #include <numeric>
 #include <stddef.h>
 #include <string>
+#include <unordered_set>
 #include <vector>
 
 using std::max;
@@ -773,32 +774,58 @@ df::coord Units::getPosition(df::unit *unit) {
 
 bool Units::teleport(df::unit *unit, df::coord target_pos)
 {   // Make sure source and dest map blocks are valid
-    auto old_occ = Maps::getTileOccupancy(unit->pos);
+    // getPosition resolves through a containing cage, so the source tile is
+    // correct even for caged units
+    const df::coord src_pos = getPosition(unit);
+    auto old_occ = Maps::getTileOccupancy(src_pos);
     auto new_occ = Maps::getTileOccupancy(target_pos);
     if (!old_occ || !new_occ)
         return false;
 
     // EQUIPMENT units (e.g. wagons) occupy a 3x3 footprint centered on their
     // position; all other units occupy just their position tile
-    int extent = 0;
-    if (auto craw = df::creature_raw::find(unit->race); craw &&
-        craw->flags.is_set(df::creature_raw_flags::EQUIPMENT_WAGON))
-        extent = 1;
+    auto unit_extent = [](df::unit *u) {
+        if (auto craw = df::creature_raw::find(u->race); craw &&
+            craw->flags.is_set(df::creature_raw_flags::EQUIPMENT_WAGON))
+            return 1;
+        return 0;
+    };
+    const int extent = unit_extent(unit);
 
     auto for_each_occupied_tile = [&](df::coord center, auto &&fn) {
         for (int dy = -extent; dy <= extent; ++dy)
-            for (int dx = -extent; dx <= extent; ++dx)
-                if (auto occ = Maps::getTileOccupancy(center.x+dx, center.y+dy, center.z))
-                    fn(*occ);
+            for (int dx = -extent; dx <= extent; ++dx) {
+                df::coord tile(center.x+dx, center.y+dy, center.z);
+                if (auto occ = Maps::getTileOccupancy(tile))
+                    fn(tile, *occ);
+            }
     };
 
+    // Occupancy flags are per-tile, not per-unit: only clear a flag when no
+    // other unit of the same kind still occupies the tile, matching how the
+    // game itself updates the flags when a unit leaves a tile
+    std::unordered_set<df::coord> grounded_tiles, standing_tiles;
+    for (auto other : world->units.active) {
+        // caged units and in-flight projectiles don't set tile occupancy, so
+        // they must not keep occupancy flags alive on their pos tile
+        if (other == unit || other->flags1.bits.caged ||
+            other->flags1.bits.projectile)
+            continue;
+        auto &tiles = other->flags1.bits.on_ground ? grounded_tiles : standing_tiles;
+        const int other_extent = unit_extent(other);
+        for (int dy = -other_extent; dy <= other_extent; ++dy)
+            for (int dx = -other_extent; dx <= other_extent; ++dx)
+                tiles.emplace(other->pos.x+dx, other->pos.y+dy, other->pos.z);
+    }
+
     // Clear appropriate occupancy flags at old tile
-    for_each_occupied_tile(unit->pos, [&](df::tile_occupancy &occ) {
-        if (unit->flags1.bits.on_ground)
-            // This is potentially wrong, but the game will recompute this as needed
-            occ.bits.unit_grounded = false;
-        else
+    for_each_occupied_tile(src_pos, [&](df::coord tile, df::tile_occupancy &occ) {
+        if (unit->flags1.bits.on_ground) {
+            if (!grounded_tiles.contains(tile))
+                occ.bits.unit_grounded = false;
+        } else if (!standing_tiles.contains(tile)) {
             occ.bits.unit = false;
+        }
     });
 
     // Clear unit projectile info
@@ -818,7 +845,7 @@ bool Units::teleport(df::unit *unit, df::coord target_pos)
         unit->flags1.bits.on_ground = true;
 
     // Set appropriate occupancy flags at new tile
-    for_each_occupied_tile(target_pos, [&](df::tile_occupancy &occ) {
+    for_each_occupied_tile(target_pos, [&](df::coord, df::tile_occupancy &occ) {
         if (unit->flags1.bits.on_ground)
             occ.bits.unit_grounded = true;
         else

test/modules/units_fortress.lua

@@ -0,0 +1,140 @@
+config.mode = 'fortress'
+config.target = 'core'
+
+local function tile_occupancy(pos)
+    return select(2, dfhack.maps.getTileFlags(pos))
+end
+
+local function two_citizens()
+    local a, b
+    for _, unit in ipairs(df.global.world.units.active) do
+        if dfhack.units.isCitizen(unit) then
+            if not a then
+                a = unit
+            else
+                b = unit
+                break
+            end
+        end
+    end
+    return a, b
+end
+
+-- find an allocated, walkable tile with no unit occupancy near pos
+local function free_tile_near(pos)
+    for dx = -4, 4 do for dy = -4, 4 do
+        if dx ~= 0 or dy ~= 0 then
+            local other = xyz2pos(pos.x + dx, pos.y + dy, pos.z)
+            local occ = tile_occupancy(other)
+            if occ and not occ.unit and not occ.unit_grounded
+                    and occ.building == df.tile_building_occ.None
+                    and df.tiletype_shape.attrs[
+                        df.tiletype.attrs[
+                            dfhack.maps.getTileType(other)].shape].walkable then
+                return other
+            end
+        end
+    end end
+    return nil, 'no free tile near the shared tile'
+end
+
+-- recompute the unit occupancy flags of the given tiles from the given
+-- units, so fabricated flag states do not leak into later tests
+local function resync_occupancy(tiles, units)
+    for _, pos in ipairs(tiles) do
+        local occ = tile_occupancy(pos)
+        if occ then
+            occ.unit = false
+            occ.unit_grounded = false
+        end
+    end
+    for _, unit in ipairs(units) do
+        local occ = tile_occupancy(unit.pos)
+        if occ then
+            if unit.flags1.on_ground then
+                occ.unit_grounded = true
+            else
+                occ.unit = true
+            end
+        end
+    end
+end
+
+-- teleporting one of two grounded units off a shared tile must keep the
+-- unit_grounded flag, since the other unit is still there (#5938)
+function test.teleport_keeps_grounded_flag_with_other_grounded_unit()
+    local a, b = two_citizens()
+    expect.ne(nil, b, 'need at least two citizens')
+    if not a or not b then return end
+
+    local shared = copyall(b.pos)
+    local dest, err = free_tile_near(b.pos)
+    expect.ne(nil, dest, err)
+    if not dest then return end
+
+    local orig_a_pos, orig_a_ground = copyall(a.pos), a.flags1.on_ground
+    local orig_b_ground = b.flags1.on_ground
+
+    return dfhack.with_finalize(function()
+        dfhack.units.teleport(a, orig_a_pos)
+        dfhack.units.teleport(b, shared)
+        a.flags1.on_ground = orig_a_ground
+        b.flags1.on_ground = orig_b_ground
+        resync_occupancy({orig_a_pos, shared, dest}, {a, b})
+    end, function()
+        -- teleporting onto a standing unit forces the mover to lie down
+        expect.true_(dfhack.units.teleport(a, shared))
+        expect.true_(a.flags1.on_ground)
+        expect.true_(tile_occupancy(shared).unit_grounded)
+
+        -- fabricate a second grounded unit on the shared tile
+        b.flags1.on_ground = true
+
+        expect.true_(dfhack.units.teleport(a, dest))
+        -- b is still grounded on the shared tile, so the flag must remain
+        expect.true_(tile_occupancy(shared).unit_grounded)
+        expect.true_(tile_occupancy(dest).unit_grounded)
+
+        -- removing the last grounded unit still clears the flag
+        expect.true_(dfhack.units.teleport(b, orig_a_pos))
+        expect.false_(tile_occupancy(shared).unit_grounded)
+    end)
+end
+
+-- the same invariant applies to the standing 'unit' flag
+function test.teleport_keeps_unit_flag_with_other_standing_unit()
+    local a, b = two_citizens()
+    expect.ne(nil, b, 'need at least two citizens')
+    if not a or not b then return end
+
+    local shared = copyall(b.pos)
+    local dest, err = free_tile_near(b.pos)
+    expect.ne(nil, dest, err)
+    if not dest then return end
+
+    local orig_a_pos, orig_a_ground = copyall(a.pos), a.flags1.on_ground
+    local orig_b_ground = b.flags1.on_ground
+
+    return dfhack.with_finalize(function()
+        dfhack.units.teleport(a, orig_a_pos)
+        dfhack.units.teleport(b, shared)
+        a.flags1.on_ground = orig_a_ground
+        b.flags1.on_ground = orig_b_ground
+        resync_occupancy({orig_a_pos, shared, dest}, {a, b})
+    end, function()
+        expect.true_(dfhack.units.teleport(a, shared))
+        expect.true_(a.flags1.on_ground)
+
+        -- make a stand again so both units are standing on the shared tile
+        a.flags1.on_ground = false
+        tile_occupancy(shared).unit_grounded = false
+
+        expect.true_(dfhack.units.teleport(a, dest))
+        expect.true_(tile_occupancy(shared).unit)
+        expect.true_(tile_occupancy(dest).unit)
+
+        -- removing the last standing unit still clears the flag
+        expect.true_(dfhack.units.teleport(b, orig_a_pos))
+        expect.false_(tile_occupancy(shared).unit)
+    end)
+end