Shuffle body container subs slightly (#21084)

Content.Shared/Body/Systems/SharedBodySystem.Body.cs

@@ -25,63 +25,11 @@ public partial class SharedBodySystem
 
     private void InitializeBody()
     {
-        // Body here to handle root body parts.
-        SubscribeLocalEvent<BodyComponent, EntInsertedIntoContainerMessage>(OnBodyInserted);
-        SubscribeLocalEvent<BodyComponent, EntRemovedFromContainerMessage>(OnBodyRemoved);
-
         SubscribeLocalEvent<BodyComponent, ComponentInit>(OnBodyInit);
         SubscribeLocalEvent<BodyComponent, MapInitEvent>(OnBodyMapInit);
         SubscribeLocalEvent<BodyComponent, CanDragEvent>(OnBodyCanDrag);
     }
 
-    private void OnBodyInserted(EntityUid uid, BodyComponent component, EntInsertedIntoContainerMessage args)
-    {
-        // Root body part?
-        var slotId = args.Container.ID;
-
-        if (slotId != BodyRootContainerId)
-            return;
-
-        var entity = args.Entity;
-
-        if (TryComp(entity, out BodyPartComponent? childPart))
-        {
-            AddPart(uid, entity, slotId, childPart);
-            RecursiveBodyUpdate(entity, uid, childPart);
-        }
-
-        if (TryComp(entity, out OrganComponent? organ))
-        {
-            AddOrgan(entity, uid, uid, organ);
-        }
-    }
-
-    private void OnBodyRemoved(EntityUid uid, BodyComponent component, EntRemovedFromContainerMessage args)
-    {
-        // TODO: lifestage shenanigans
-        if (TerminatingOrDeleted(uid))
-            return;
-
-        // Root body part?
-        var slotId = args.Container.ID;
-
-        if (slotId != BodyRootContainerId)
-            return;
-
-        var entity = args.Entity;
-
-        if (TryComp(entity, out BodyPartComponent? childPart))
-        {
-            RemovePart(uid, entity, slotId, childPart);
-            RecursiveBodyUpdate(entity, null, childPart);
-        }
-
-        if (TryComp(entity, out OrganComponent? organ))
-        {
-            RemoveOrgan(entity, uid, organ);
-        }
-    }
-
     private void OnBodyInit(EntityUid bodyId, BodyComponent body, ComponentInit args)
     {
         // Setup the initial container.
@@ -229,6 +177,8 @@ public partial class SharedBodySystem
             yield break;
         }
 
+        yield return (body.RootContainer.ContainedEntity.Value, rootPart);
+
         foreach (var child in GetBodyPartChildren(body.RootContainer.ContainedEntity.Value, rootPart))
         {
             yield return child;

Content.Shared/Body/Systems/SharedBodySystem.Organs.cs

@@ -9,26 +9,50 @@ namespace Content.Shared.Body.Systems;
 
 public partial class SharedBodySystem
 {
-    private void AddOrgan(EntityUid uid, EntityUid bodyUid, EntityUid parentPartUid, OrganComponent component)
+    private void InitializeOrgans()
     {
-        component.Body = bodyUid;
+        SubscribeLocalEvent<OrganComponent, EntGotInsertedIntoContainerMessage>(OnOrganInserted);
-        RaiseLocalEvent(uid, new AddedToPartEvent(parentPartUid));
+        SubscribeLocalEvent<OrganComponent, EntGotRemovedFromContainerMessage>(OnOrganRemoved);
-
-        if (component.Body != null)
-            RaiseLocalEvent(uid, new AddedToPartInBodyEvent(component.Body.Value, parentPartUid));
-
-        Dirty(uid, component);
     }
 
-    private void RemoveOrgan(EntityUid uid, EntityUid parentPartUid, OrganComponent component)
+    private void OnOrganInserted(EntityUid uid, OrganComponent component, EntGotInsertedIntoContainerMessage args)
     {
-        RaiseLocalEvent(uid, new RemovedFromPartEvent(parentPartUid));
+        // No recursive updates for these as you can't insert organs into organsTM.
+        var parentUid = args.Container.Owner;
 
-        if (component.Body != null)
+        if (HasComp<BodyComponent>(parentUid))
         {
-            RaiseLocalEvent(uid, new RemovedFromPartInBodyEvent(component.Body.Value, parentPartUid));
+            component.Body = parentUid;
+            Dirty(uid, component);
         }
 
+        // Organ inserted into body part.
+        if (TryComp(parentUid, out BodyPartComponent? partComp))
+        {
+            RaiseLocalEvent(uid, new AddedToPartEvent(parentUid));
+
+            if (partComp.Body != null)
+            {
+                component.Body = partComp.Body;
+                RaiseLocalEvent(uid, new AddedToPartInBodyEvent(partComp.Body.Value, parentUid));
+                Dirty(uid, component);
+            }
+        }
+    }
+
+    private void OnOrganRemoved(EntityUid uid, OrganComponent component, EntGotRemovedFromContainerMessage args)
+    {
+        // Lifestage shenanigans.
+        if (component.Body != null && TerminatingOrDeleted(component.Body.Value))
+            return;
+
+        if (HasComp<BodyPartComponent>(args.Container.Owner))
+            RaiseLocalEvent(uid, new RemovedFromPartEvent(args.Container.Owner));
+
+        if (component.Body == null)
+            return;
+
+        RaiseLocalEvent(uid, new RemovedFromPartInBodyEvent(component.Body.Value, args.Container.Owner));
         component.Body = null;
         Dirty(uid, component);
     }

Content.Shared/Body/Systems/SharedBodySystem.Parts.cs

@@ -19,89 +19,99 @@ public partial class SharedBodySystem
         // TODO: This doesn't handle comp removal on child ents.
 
         // If you modify this also see the Body partial for root parts.
-        SubscribeLocalEvent<BodyPartComponent, EntInsertedIntoContainerMessage>(OnBodyPartInserted);
+        SubscribeLocalEvent<BodyPartComponent, EntGotInsertedIntoContainerMessage>(OnBodyPartInserted);
-        SubscribeLocalEvent<BodyPartComponent, EntRemovedFromContainerMessage>(OnBodyPartRemoved);
+        SubscribeLocalEvent<BodyPartComponent, EntGotRemovedFromContainerMessage>(OnBodyPartRemoved);
     }
 
-    private void OnBodyPartInserted(EntityUid uid, BodyPartComponent component, EntInsertedIntoContainerMessage args)
+    private void OnBodyPartInserted(EntityUid uid, BodyPartComponent component, EntGotInsertedIntoContainerMessage args)
     {
-        // Body part inserted into another body part.
+        var parentUid = args.Container.Owner;
-        var entity = args.Entity;
-        var slotId = args.Container.ID;
 
+        if (HasComp<BodyComponent>(parentUid))
+        {
+            AddPart(parentUid, uid, GetPartSlotContainerId(args.Container.ID), component);
+            RecursiveBodyUpdate(uid, parentUid, component);
+        }
+
+        // Body part inserted into another body part.
+        if (TryComp(parentUid, out BodyPartComponent? partComp))
+        {
+            if (partComp.Body == null)
+                return;
+
+            AddPart(partComp.Body.Value, uid, GetPartSlotContainerId(args.Container.ID), component);
+            RecursiveBodyUpdate(uid, partComp.Body, component);
+        }
+    }
+
+    private void OnBodyPartRemoved(EntityUid uid, BodyPartComponent component, EntGotRemovedFromContainerMessage args)
+    {
         if (component.Body != null)
         {
-            if (TryComp(entity, out BodyPartComponent? childPart))
+            // Lifestage shenanigans.
-            {
+            if (TerminatingOrDeleted(component.Body.Value))
-                AddPart(component.Body.Value, entity, slotId, childPart);
+                return;
-                RecursiveBodyUpdate(entity, component.Body.Value, childPart);
-            }
 
-            if (TryComp(entity, out OrganComponent? organ))
+            RemovePart(component.Body.Value, uid, GetPartSlotContainerId(args.Container.ID), component);
-            {
+            RecursiveBodyUpdate(uid, null, component);
-                AddOrgan(entity, component.Body.Value, uid, organ);
-            }
-        }
-    }
-
-    private void OnBodyPartRemoved(EntityUid uid, BodyPartComponent component, EntRemovedFromContainerMessage args)
-    {
-        // TODO: lifestage shenanigans
-        if (TerminatingOrDeleted(uid))
-            return;
-
-        // Body part removed from another body part.
-        var entity = args.Entity;
-        var slotId = args.Container.ID;
-
-        if (TryComp(entity, out BodyPartComponent? childPart) && childPart.Body != null)
-        {
-            RemovePart(childPart.Body.Value, entity, slotId, childPart);
-            RecursiveBodyUpdate(entity, null, childPart);
-        }
-
-        if (TryComp(entity, out OrganComponent? organ))
-        {
-            RemoveOrgan(entity, uid, organ);
         }
     }
 
+    /// <summary>
+    /// Updates body parts recursively for a particular entity, including itself.
+    /// </summary>
     private void RecursiveBodyUpdate(EntityUid uid, EntityUid? bodyUid, BodyPartComponent component)
     {
+        // Recursively iterate all child parts and organs and ensure their body refs are updated.
         foreach (var children in GetBodyPartChildren(uid, component))
         {
-            if (children.Component.Body != bodyUid)
+            if (children.Component.Body == bodyUid)
+                continue;
+
+            // Raise the body part change event if relevant.
+            if (Containers.TryGetContainingContainer(children.Id, out var childContainer))
             {
-                children.Component.Body = bodyUid;
+                var slotId = GetPartSlotContainerId(childContainer.ID);
-                Dirty(children.Id, children.Component);
 
-                foreach (var slotId in children.Component.Organs.Keys)
+                if (bodyUid != null)
                 {
-                    var organContainerId = GetOrganContainerId(slotId);
+                    AddPart(bodyUid.Value, children.Id, slotId, children.Component);
+                }
+                else if (component.Body != null)
+                {
+                    RemovePart(component.Body.Value, children.Id, slotId, children.Component);
+                }
+            }
 
-                    if (!Containers.TryGetContainer(children.Id, organContainerId, out var container))
+            children.Component.Body = bodyUid;
-                        continue;
+            Dirty(children.Id, children.Component);
 
-                    foreach (var organ in container.ContainedEntities)
+            foreach (var slotId in children.Component.Organs.Keys)
+            {
+                var organContainerId = GetOrganContainerId(slotId);
+
+                if (!Containers.TryGetContainer(children.Id, organContainerId, out var container))
+                    continue;
+
+                foreach (var organ in container.ContainedEntities)
+                {
+                    if (TryComp(organ, out OrganComponent? organComp))
                     {
-                        if (TryComp(organ, out OrganComponent? organComp))
+                        var oldBody = organComp.Body;
+                        organComp.Body = bodyUid;
+
+                        if (bodyUid != null)
                         {
-                            var oldBody = organComp.Body;
+                            var ev = new AddedToPartInBodyEvent(bodyUid.Value, children.Id);
-                            organComp.Body = bodyUid;
+                            RaiseLocalEvent(organ, ev);
-
-                            if (bodyUid != null)
-                            {
-                                var ev = new AddedToPartInBodyEvent(bodyUid.Value, children.Id);
-                                RaiseLocalEvent(organ, ev);
-                            }
-                            else if (oldBody != null)
-                            {
-                                var ev = new RemovedFromPartInBodyEvent(oldBody.Value, children.Id);
-                                RaiseLocalEvent(organ, ev);
-                            }
-
-                            Dirty(organ, organComp);
                         }
+                        else if (oldBody != null)
+                        {
+                            var ev = new RemovedFromPartInBodyEvent(oldBody.Value, children.Id);
+                            RaiseLocalEvent(organ, ev);
+                        }
+
+                        Dirty(organ, organComp);
                     }
                 }
             }
@@ -116,9 +126,8 @@ public partial class SharedBodySystem
         BodyComponent? bodyComp = null)
     {
         DebugTools.AssertOwner(partUid, component);
-        Dirty(partUid, component);
         component.Body = bodyUid;
-
+        Dirty(partUid, component);
         var ev = new BodyPartAddedEvent(slotId, component);
         RaiseLocalEvent(bodyUid, ref ev);
 
@@ -134,8 +143,8 @@ public partial class SharedBodySystem
     {
         DebugTools.AssertOwner(partUid, component);
         Resolve(bodyUid, ref bodyComp, false);
-        Dirty(partUid, component);
         component.Body = null;
+        Dirty(partUid, component);
 
         var ev = new BodyPartRemovedEvent(slotId, component);
         RaiseLocalEvent(bodyUid, ref ev);
@@ -547,15 +556,13 @@ public partial class SharedBodySystem
     }
 
     /// <summary>
-    /// Returns all body part components for this entity including itself.
+    /// Returns all body part components for this entity not including itself.
     /// </summary>
     public IEnumerable<(EntityUid Id, BodyPartComponent Component)> GetBodyPartChildren(EntityUid partId, BodyPartComponent? part = null)
     {
         if (!Resolve(partId, ref part, false))
             yield break;
 
-        yield return (partId, part);
-
         foreach (var slotId in part.Children.Keys)
         {
             var containerSlotId = GetPartSlotContainerId(slotId);
@@ -567,6 +574,8 @@ public partial class SharedBodySystem
                     if (!TryComp(containedEnt, out BodyPartComponent? childPart))
                         continue;
 
+                    yield return (containedEnt, childPart);
+
                     foreach (var value in GetBodyPartChildren(containedEnt, childPart))
                     {
                         yield return value;
@@ -629,6 +638,9 @@ public partial class SharedBodySystem
             return false;
         }
 
+        if (childId == parentId)
+            return true;
+
         foreach (var (foundId, _) in GetBodyPartChildren(parentId, parent))
         {
             if (foundId == childId)

Content.Shared/Body/Systems/SharedBodySystem.cs

@@ -39,6 +39,7 @@ public abstract partial class SharedBodySystem : EntitySystem
         base.Initialize();
 
         InitializeBody();
+        InitializeOrgans();
         InitializeParts();
     }