Optimise broadphase re-insertions

Robust.Benchmarks/Physics/BroadphaseReinsertBenchmark.cs

@@ -0,0 +1,79 @@
+using System.Numerics;
+using BenchmarkDotNet.Attributes;
+using Robust.Shared.Analyzers;
+using Robust.Shared.GameObjects;
+using Robust.Shared.Map;
+using Robust.Shared.Maths;
+using Robust.Shared.Physics;
+using Robust.Shared.Physics.Collision.Shapes;
+using Robust.Shared.Physics.Components;
+using Robust.Shared.Physics.Dynamics;
+using Robust.Shared.Physics.Systems;
+using Robust.UnitTesting.Server;
+
+namespace Robust.Benchmarks.Physics;
+
+[Virtual, MemoryDiagnoser]
+public class BroadphaseReinsertBenchmark
+{
+    private IEntityManager _entManager = default!;
+    private EntityLookupSystem _lookup = default!;
+    private EntityUid _root;
+    private TransformComponent _rootXform = default!;
+
+    [Params(10, 100, 500)]
+    public int Children;
+
+    [GlobalSetup]
+    public void GlobalSetup()
+    {
+        var sim = RobustServerSimulation.NewSimulation().InitializeInstance();
+        _entManager = sim.Resolve<IEntityManager>();
+        _lookup = _entManager.System<EntityLookupSystem>();
+        var xforms = _entManager.System<SharedTransformSystem>();
+        var physics = _entManager.System<SharedPhysicsSystem>();
+        var fixtures = _entManager.System<FixtureSystem>();
+        var mapSys = _entManager.System<SharedMapSystem>();
+
+        var (map, mapId) = sim.CreateMap();
+        var grid = mapSys.CreateGridEntity(mapId);
+        var gridUid = grid.Owner;
+        mapSys.SetTile(grid, Vector2i.Zero, new Tile(1));
+        xforms.SetCoordinates(gridUid, new EntityCoordinates(map, new Vector2(10f, 10f)));
+        xforms.SetLocalRotation(gridUid, Angle.FromDegrees(35));
+
+        _root = _entManager.SpawnEntity(null, new EntityCoordinates(gridUid, new Vector2(0.5f, 0.5f)));
+        _rootXform = _entManager.GetComponent<TransformComponent>(_root);
+
+        var parent = _root;
+        for (var i = 0; i < Children; i++)
+        {
+            var child = _entManager.SpawnEntity(null, new EntityCoordinates(parent, new Vector2(0.001f, 0f)));
+
+            if (i % 4 == 0)
+            {
+                var xform = _entManager.GetComponent<TransformComponent>(child);
+                var body = _entManager.AddComponent<PhysicsComponent>(child);
+                var shape = new PolygonShape();
+                shape.SetAsBox(0.01f, 0.01f);
+                fixtures.CreateFixture(
+                    child,
+                    "fix1",
+                    new Fixture(shape, 0, 0, false),
+                    body: body,
+                    xform: xform);
+                physics.SetCanCollide(child, true, body: body);
+                _lookup.FindAndAddToEntityTree(child, false, xform);
+            }
+
+            parent = child;
+        }
+    }
+
+    [Benchmark]
+    public void RemoveAndReinsertRecursive()
+    {
+        _lookup.RemoveFromEntityTree(_root, _rootXform);
+        _lookup.FindAndAddToEntityTree(_root, true, _rootXform);
+    }
+}

Robust.Shared.IntegrationTests/Physics/Broadphase_Test.cs

@@ -230,6 +230,85 @@ public void BroadphaseRecursiveUpdate()
         Assert.That(lookup.FindBroadphase(child2), Is.EqualTo(mapBroadphase));
     }
 
+    [Test]
+    public void BroadphaseRecursiveReinsert()
+    {
+        var sim = RobustServerSimulation.NewSimulation().InitializeInstance();
+        var entManager = sim.Resolve<IEntityManager>();
+        var system = entManager.EntitySysManager;
+        var lookup = system.GetEntitySystem<EntityLookupSystem>();
+        var xforms = system.GetEntitySystem<SharedTransformSystem>();
+        var physics = system.GetEntitySystem<SharedPhysicsSystem>();
+        var fixtures = system.GetEntitySystem<FixtureSystem>();
+        var mapSys = entManager.System<SharedMapSystem>();
+
+        var (map, mapId) = sim.CreateMap();
+        var grid = mapSys.CreateGridEntity(mapId);
+        var gridUid = grid.Owner;
+        mapSys.SetTile(grid, Vector2i.Zero, new Tile(1));
+        xforms.SetCoordinates(gridUid, new EntityCoordinates(map, new Vector2(10f, 10f)));
+        xforms.SetLocalRotation(gridUid, Angle.FromDegrees(90));
+
+        var gridBroadphase = entManager.GetComponent<BroadphaseComponent>(gridUid);
+        var broadphaseData = new BroadphaseData(gridUid, false, false);
+
+        var parent = entManager.SpawnEntity(null, new EntityCoordinates(gridUid, new Vector2(0.5f, 0.5f)));
+        var child = entManager.SpawnEntity(null, new EntityCoordinates(parent, new Vector2(0.2f, 0f)));
+        var collidableChild = entManager.SpawnEntity(null, new EntityCoordinates(child, new Vector2(0.1f, 0f)));
+        var parentXform = entManager.GetComponent<TransformComponent>(parent);
+        var childXform = entManager.GetComponent<TransformComponent>(child);
+        var collidableXform = entManager.GetComponent<TransformComponent>(collidableChild);
+        var collidableBody = entManager.AddComponent<PhysicsComponent>(collidableChild);
+
+        var shape = new PolygonShape();
+        shape.SetAsBox(0.25f, 0.25f);
+        fixtures.CreateFixture(
+            collidableChild,
+            "fix1",
+            new Fixture(shape, 0, 0, false),
+            body: collidableBody,
+            xform: collidableXform);
+        physics.SetCanCollide(collidableChild, true, body: collidableBody);
+        lookup.FindAndAddToEntityTree(collidableChild, false, collidableXform);
+
+        Assert.That(parentXform.Broadphase, Is.EqualTo(broadphaseData));
+        Assert.That(childXform.Broadphase, Is.EqualTo(broadphaseData));
+        Assert.That(collidableXform.Broadphase, Is.EqualTo(new BroadphaseData(gridUid, true, true)));
+        Assert.That(gridBroadphase.SundriesTree, Does.Contain(parent));
+        Assert.That(gridBroadphase.SundriesTree, Does.Contain(child));
+        Assert.That(gridBroadphase.StaticTree.Count, Is.EqualTo(1));
+
+        lookup.RemoveFromEntityTree(parent, parentXform);
+        Assert.That(parentXform.Broadphase, Is.Null);
+        Assert.That(childXform.Broadphase, Is.Null);
+        Assert.That(collidableXform.Broadphase, Is.Null);
+        Assert.That(gridBroadphase.SundriesTree, Does.Not.Contain(parent));
+        Assert.That(gridBroadphase.SundriesTree, Does.Not.Contain(child));
+        Assert.That(gridBroadphase.StaticTree.Count, Is.EqualTo(0));
+
+        lookup.FindAndAddToEntityTree(parent, true, parentXform);
+
+        Assert.That(parentXform.Broadphase, Is.EqualTo(broadphaseData));
+        Assert.That(childXform.Broadphase, Is.EqualTo(broadphaseData));
+        Assert.That(collidableXform.Broadphase, Is.EqualTo(new BroadphaseData(gridUid, true, true)));
+        Assert.That(gridBroadphase.SundriesTree, Does.Contain(parent));
+        Assert.That(gridBroadphase.SundriesTree, Does.Contain(child));
+        Assert.That(gridBroadphase.StaticTree.Count, Is.EqualTo(1));
+
+        var parentWorld = xforms.GetWorldPosition(parent);
+        var childWorld = xforms.GetWorldPosition(child);
+        var collidableWorld = xforms.GetWorldPosition(collidableChild);
+
+        var found = lookup.GetEntitiesIntersecting(mapId, Box2.CenteredAround(parentWorld, Vector2.One));
+        Assert.That(found, Does.Contain(parent));
+
+        found = lookup.GetEntitiesIntersecting(mapId, Box2.CenteredAround(childWorld, Vector2.One));
+        Assert.That(found, Does.Contain(child));
+
+        found = lookup.GetEntitiesIntersecting(mapId, Box2.CenteredAround(collidableWorld, Vector2.One));
+        Assert.That(found, Does.Contain(collidableChild));
+    }
+
     /// <summary>
     /// Check that broadphases properly recursively update when entities move between maps and grids. The broadphase
     /// updating handles grids separately from other entities, this is intended to be an exhaustive check that the

Robust.Shared/GameObjects/Systems/EntityLookupSystem.cs

@@ -433,13 +433,27 @@ private void AddOrUpdatePhysicsTree(
         xform.Broadphase ??= new(broadUid, body.CanCollide, body.BodyType == BodyType.Static);
         var tree = body.BodyType == BodyType.Static ? broadphase.StaticTree : broadphase.DynamicTree;
 
-        // TOOD optimize this. This function iterates UP through parents, while we are currently iterating down.
         var (worldPos, worldRot) = _transform.GetWorldPositionRotation(xform);
         var mapTransform = new Transform(worldPos, worldRot);
 
         // TODO BROADPHASE PARENTING this just assumes local = world
         var broadphaseTransform = new Transform(Vector2.Transform(mapTransform.Position, broadphaseXform.InvLocalMatrix), mapTransform.Quaternion2D.Angle - broadphaseXform.LocalRotation);
 
+        AddOrUpdatePhysicsTree(uid, broadUid, xform, body, manager, tree, broadphaseTransform);
+    }
+
+    private void AddOrUpdatePhysicsTree(
+        EntityUid uid,
+        EntityUid broadUid,
+        TransformComponent xform,
+        PhysicsComponent body,
+        FixturesComponent manager,
+        IBroadPhase tree,
+        Transform broadphaseTransform)
+    {
+        DebugTools.Assert(body.Owner == uid);
+        xform.Broadphase ??= new(broadUid, body.CanCollide, body.BodyType == BodyType.Static);
+
         foreach (var (id, fixture) in manager.Fixtures)
         {
             AddOrMoveProxies((uid, body, xform), id, fixture, tree, broadphaseTransform);
@@ -595,7 +609,13 @@ public void FindAndAddToEntityTree(EntityUid uid, bool recursive = true, Transfo
             return;
 
         if (TryFindBroadphase(xform, out var broadphase))
-            AddOrUpdateEntityTree(broadphase.Owner, broadphase, uid, xform, recursive);
+            AddOrUpdateEntityTreeDown(
+                broadphase.Owner,
+                broadphase,
+                _xformQuery.GetComponent(broadphase.Owner),
+                uid,
+                xform,
+                recursive);
     }
 
     /// <summary>
@@ -636,6 +656,39 @@ private void AddOrUpdateEntityTree(
         EntityUid uid,
         TransformComponent xform,
         bool recursive = true)
+    {
+        AddOrUpdateEntityTreeDown(broadUid, broadphase, broadphaseXform, uid, xform, recursive);
+    }
+
+    private void AddOrUpdateEntityTreeDown(
+        EntityUid broadUid,
+        BroadphaseComponent broadphase,
+        TransformComponent broadphaseXform,
+        EntityUid uid,
+        TransformComponent xform,
+        bool recursive = true)
+    {
+        var (worldPos, worldRot) = _transform.GetWorldPositionRotation(xform);
+        AddOrUpdateEntityTreeDown(
+            broadUid,
+            broadphase,
+            broadphaseXform,
+            uid,
+            xform,
+            worldPos,
+            worldRot,
+            recursive);
+    }
+
+    private void AddOrUpdateEntityTreeDown(
+        EntityUid broadUid,
+        BroadphaseComponent broadphase,
+        TransformComponent broadphaseXform,
+        EntityUid uid,
+        TransformComponent xform,
+        Vector2 worldPos,
+        Angle worldRot,
+        bool recursive = true)
     {
         if (xform.Broadphase != null && !xform.Broadphase.Value.IsValid())
         {
@@ -644,34 +697,42 @@ private void AddOrUpdateEntityTree(
             return;
         }
 
+        var relativePosition = Vector2.Transform(worldPos, broadphaseXform.InvLocalMatrix);
+        var relativeRotation = worldRot - broadphaseXform.LocalRotation;
+
         if (!_physicsQuery.TryGetComponent(uid, out var body) || !body.CanCollide)
         {
-            // TODO optimize this. This function iterates UP through parents, while we are currently iterating down.
-            var (coordinates, rotation) = _transform.GetMoverCoordinateRotation(uid, xform);
-
-            // TODO BROADPHASE PARENTING this just assumes local = world
-            var relativeRotation = rotation - broadphaseXform.LocalRotation;
-
-            var aabb = GetAABBNoContainer(uid, coordinates.Position, relativeRotation);
+            var aabb = GetAABBNoContainer(uid, relativePosition, relativeRotation);
             AddOrUpdateSundriesTree(broadUid, broadphase, uid, xform, body?.BodyType == BodyType.Static, aabb);
         }
         else
         {
-            AddOrUpdatePhysicsTree(uid, broadUid, broadphase, broadphaseXform, xform, body, _fixturesQuery.GetComponent(uid));
+            var tree = body.BodyType == BodyType.Static ? broadphase.StaticTree : broadphase.DynamicTree;
+            var broadphaseTransform = new Transform(relativePosition, relativeRotation);
+            AddOrUpdatePhysicsTree(uid, broadUid, xform, body, _fixturesQuery.GetComponent(uid), tree, broadphaseTransform);
         }
 
         if (xform.ChildCount == 0 || !recursive)
             return;
 
-        // TODO can this be removed?
-        // AFAIK the separate container check is redundant now that we check for an invalid broadphase at the beginning of this function.
         if (!_containerQuery.HasComponent(uid))
         {
             foreach (var child in xform._children)
             {
                 var childXform = _xformQuery.GetComponent(child);
-                AddOrUpdateEntityTree(broadUid, broadphase, broadphaseXform, child, childXform, recursive);
+                var childWorldPos = worldRot.RotateVec(childXform.LocalPosition) + worldPos;
+                var childWorldRot = worldRot + childXform.LocalRotation;
+                AddOrUpdateEntityTreeDown(
+                    broadUid,
+                    broadphase,
+                    broadphaseXform,
+                    child,
+                    childXform,
+                    childWorldPos,
+                    childWorldRot,
+                    recursive);
             }
+
             return;
         }
 
@@ -681,7 +742,17 @@ private void AddOrUpdateEntityTree(
                 continue;
 
             var childXform = _xformQuery.GetComponent(child);
-            AddOrUpdateEntityTree(broadUid, broadphase, broadphaseXform, child, childXform, recursive);
+            var childWorldPos = worldRot.RotateVec(childXform.LocalPosition) + worldPos;
+            var childWorldRot = worldRot + childXform.LocalRotation;
+            AddOrUpdateEntityTreeDown(
+                broadUid,
+                broadphase,
+                broadphaseXform,
+                child,
+                childXform,
+                childWorldPos,
+                childWorldRot,
+                recursive);
         }
     }