Add mechanism events when being added/removed to/from body/parts (#2271)

* Add mechanism events when added/removed to/from body/parts

* Change old usages

* Add TODO

Content.Client/GameObjects/Components/Body/Mechanism/MechanismComponent.cs

@@ -11,9 +11,9 @@ namespace Content.Client.GameObjects.Components.Body.Mechanism
     [ComponentReference(typeof(IMechanism))]
     public class MechanismComponent : SharedMechanismComponent
     {
-        protected override void OnPartAdd(IBodyPart? old, IBodyPart current)
+        protected override void OnAddedToPart()
         {
-            base.OnPartAdd(old, current);
+            base.OnAddedToPart();
 
             if (Owner.TryGetComponent(out ISpriteComponent? sprite))
             {
@@ -21,9 +21,9 @@ namespace Content.Client.GameObjects.Components.Body.Mechanism
             }
         }
 
-        protected override void OnPartRemove(IBodyPart old)
+        protected override void OnRemovedFromPart(IBodyPart old)
         {
-            base.OnPartRemove(old);
+            base.OnRemovedFromPart(old);
 
             if (Owner.TryGetComponent(out ISpriteComponent? sprite))
             {

Content.Server/GameObjects/Components/Body/MechanismComponent.cs

@@ -158,9 +158,9 @@ namespace Content.Server.GameObjects.Components.Body
             }
         }
 
-        protected override void OnPartAdd(IBodyPart? old, IBodyPart current)
+        protected override void OnAddedToPart()
         {
-            base.OnPartAdd(old, current);
+            base.OnAddedToPart();
 
             if (Owner.TryGetComponent(out SpriteComponent? sprite))
             {
@@ -168,9 +168,9 @@ namespace Content.Server.GameObjects.Components.Body
             }
         }
 
-        protected override void OnPartRemove(IBodyPart old)
+        protected override void OnRemovedFromPart(IBodyPart old)
         {
-            base.OnPartRemove(old);
+            base.OnRemovedFromPart(old);
 
             if (Owner.TryGetComponent(out SpriteComponent? sprite))
             {

Content.Shared/GameObjects/Components/Body/Behavior/IMechanismBehavior.cs

@@ -15,5 +15,55 @@ namespace Content.Shared.GameObjects.Components.Body.Behavior
         IMechanism? Mechanism { get; }
 
         void Update(float frameTime);
+
+        /// <summary>
+        ///     Called when the containing <see cref="IBodyPart"/> is attached to a
+        ///     <see cref="IBody"/>.
+        ///     For instance, attaching a head with a brain inside to a body.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void AddedToBody();
+
+        /// <summary>
+        ///     Called when the parent <see cref="IMechanism"/> is
+        ///     added into a <see cref="IBodyPart"/> that is not attached to a
+        ///     <see cref="IBody"/>.
+        ///     For instance, adding a brain to a dismembered head.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void AddedToPart();
+
+        /// <summary>
+        ///     Called when the parent <see cref="IBodyPart"/> is removed from a
+        ///     <see cref="IBody"/>.
+        ///     For instance, removing a head with a brain inside from a body.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void RemovedFromBody(IBody old);
+
+        /// <summary>
+        ///     Called when the parent <see cref="IMechanism"/> is
+        ///     removed from a <see cref="IBodyPart"/> that is not attached to a
+        ///     <see cref="IBody"/>.
+        ///     For instance, removing a brain from a dismembered head.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void RemovedFromPart(IBodyPart old);
+
+        /// <summary>
+        ///     Called when the parent <see cref="IMechanism"/> is added to a
+        ///     <see cref="IBodyPart"/> that is attached to a <see cref="IBody"/>.
+        ///     For instance, adding a brain to a head that is attached to a body.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void AddedToPartInBody();
+
+        /// <summary>
+        ///     Called when the parent <see cref="IMechanism"/> is removed from a
+        ///     <see cref="IBodyPart"/> that is attached to a <see cref="IBody"/>.
+        ///     For instance, removing a brain from a head that is attached to a body.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void RemovedFromPartInBody(IBody? oldBody, IBodyPart? oldPart);
     }
 }

Content.Shared/GameObjects/Components/Body/Behavior/MechanismBehaviorComponent.cs

@@ -15,72 +15,46 @@ namespace Content.Shared.GameObjects.Components.Body.Behavior
 
         public abstract void Update(float frameTime);
 
-        /// <summary>
-        ///     Called when the containing <see cref="IBodyPart"/> is attached to a
-        ///     <see cref="IBody"/>.
-        ///     For instance, attaching a head to a body will call this on the brain inside.
-        /// </summary>
         public void AddedToBody()
         {
             OnAddedToBody();
         }
 
-        /// <summary>
-        ///     Called when the parent <see cref="Mechanism"/> is
-        ///     added into a <see cref="IBodyPart"/>.
-        ///     For instance, putting a brain into an empty head.
-        /// </summary>
         public void AddedToPart()
         {
             OnAddedToPart();
         }
 
-        /// <summary>
-        ///     Called when the containing <see cref="IBodyPart"/> is removed from a
-        ///     <see cref="IBody"/>.
-        ///     For instance, cutting off ones head will call this on the brain inside.
-        /// </summary>
         public void RemovedFromBody(IBody old)
         {
             OnRemovedFromBody(old);
         }
 
-        /// <summary>
-        ///     Called when the parent <see cref="Mechanism"/> is
-        ///     removed from a <see cref="IBodyPart"/>.
-        ///     For instance, taking a brain out of ones head.
-        /// </summary>
         public void RemovedFromPart(IBodyPart old)
         {
             OnRemovedFromPart(old);
         }
 
-        /// <summary>
+        public void AddedToPartInBody()
-        ///     Called when the containing <see cref="IBodyPart"/> is attached to a
+        {
-        ///     <see cref="IBody"/>.
+            OnAddedToPart();
-        ///     For instance, attaching a head to a body will call this on the brain inside.
+        }
-        /// </summary>
+
+        public void RemovedFromPartInBody(IBody? oldBody, IBodyPart? oldPart)
+        {
+            OnRemovedFromPartInBody(oldBody, oldPart);
+        }
+
         protected virtual void OnAddedToBody() { }
 
-        /// <summary>
-        ///     Called when the parent <see cref="Mechanism"/> is
-        ///     added into a <see cref="IBodyPart"/>.
-        ///     For instance, putting a brain into an empty head.
-        /// </summary>
         protected virtual void OnAddedToPart() { }
 
-        /// <summary>
-        ///     Called when the containing <see cref="IBodyPart"/> is removed from a
-        ///     <see cref="IBody"/>.
-        ///     For instance, cutting off ones head will call this on the brain inside.
-        /// </summary>
         protected virtual void OnRemovedFromBody(IBody old) { }
 
-        /// <summary>
-        ///     Called when the parent <see cref="Mechanism"/> is
-        ///     removed from a <see cref="IBodyPart"/>.
-        ///     For instance, taking a brain out of ones head.
-        /// </summary>
         protected virtual void OnRemovedFromPart(IBodyPart old) { }
+
+        protected virtual void OnAddedToPartInBody() { }
+
+        protected virtual void OnRemovedFromPartInBody(IBody? oldBody, IBodyPart? oldPart) { }
     }
 }

Content.Shared/GameObjects/Components/Body/Mechanism/IMechanism.cs

@@ -52,20 +52,55 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
         /// </summary>
         BodyPartCompatibility Compatibility { get; set; }
 
+        // TODO BODY Turn these into event listeners so they dont need to be exposed
         /// <summary>
-        ///     Called when the part housing this mechanism is added to a body.
+        ///     Called when the containing <see cref="IBodyPart"/> is attached to a
-        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY PART CODE!
+        ///     <see cref="IBody"/>.
+        ///     For instance, attaching a head with a brain inside to a body.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
         /// </summary>
-        /// <param name="old">The previous body, if any.</param>
+        void AddedToBody();
-        /// <param name="current">The new body.</param>
-        void OnBodyAdd(IBody? old, IBody current);
 
         /// <summary>
-        ///     Called when the part housing this mechanism is removed from
+        ///     Called when the parent <see cref="IMechanism"/> is
-        ///     a body.
+        ///     added into a <see cref="IBodyPart"/> that is not attached to a
-        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY PART CODE!
+        ///     <see cref="IBody"/>.
+        ///     For instance, adding a brain to a dismembered head.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
         /// </summary>
-        /// <param name="old">The old body.</param>
+        void AddedToPart();
-        void OnBodyRemove(IBody old);
+
+        /// <summary>
+        ///     Called when the parent <see cref="IBodyPart"/> is removed from a
+        ///     <see cref="IBody"/>.
+        ///     For instance, removing a head with a brain inside from a body.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void RemovedFromBody(IBody old);
+
+        /// <summary>
+        ///     Called when the parent <see cref="IMechanism"/> is
+        ///     removed from a <see cref="IBodyPart"/> that is not attached to a
+        ///     <see cref="IBody"/>.
+        ///     For instance, removing a brain from a dismembered head.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void RemovedFromPart(IBodyPart old);
+
+        /// <summary>
+        ///     Called when the parent <see cref="IMechanism"/> is added to a
+        ///     <see cref="IBodyPart"/> that is attached to a <see cref="IBody"/>.
+        ///     For instance, adding a brain to a head that is attached to a body.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void AddedToPartInBody();
+
+        /// <summary>
+        ///     Called when the parent <see cref="IMechanism"/> is removed from a
+        ///     <see cref="IBodyPart"/> that is attached to a <see cref="IBody"/>.
+        ///     For instance, removing a brain from a head that is attached to a body.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void RemovedFromPartInBody(IBody? oldBody, IBodyPart? oldPart);
     }
 }

Content.Shared/GameObjects/Components/Body/Mechanism/MechanismExtensions.cs

@@ -16,6 +16,35 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
                    body.Parts.Values.Any(p => p.Mechanisms.Any(m => m.Owner.HasComponent<T>()));
         }
 
+        public static IEnumerable<IMechanismBehavior> GetMechanismBehaviors(this IEntity entity)
+        {
+            if (!entity.TryGetBody(out var body))
+            {
+                yield break;
+            }
+
+            foreach (var part in body.Parts.Values)
+            foreach (var mechanism in part.Mechanisms)
+            foreach (var behavior in mechanism.Owner.GetAllComponents<IMechanismBehavior>())
+            {
+                yield return behavior;
+            }
+        }
+
+        public static bool TryGetMechanismBehaviors(this IEntity entity,
+            [NotNullWhen(true)] out List<IMechanismBehavior>? behaviors)
+        {
+            behaviors = entity.GetMechanismBehaviors().ToList();
+
+            if (behaviors.Count == 0)
+            {
+                behaviors = null;
+                return false;
+            }
+
+            return true;
+        }
+
         public static IEnumerable<T> GetMechanismBehaviors<T>(this IEntity entity) where T : class, IMechanismBehavior
         {
             if (!entity.TryGetBody(out var body))

Content.Shared/GameObjects/Components/Body/Mechanism/SharedMechanismComponent.cs

@@ -4,6 +4,7 @@ using Content.Shared.GameObjects.Components.Body.Part;
 using Robust.Shared.GameObjects;
 using Robust.Shared.Interfaces.GameObjects;
 using Robust.Shared.Serialization;
+using Robust.Shared.Utility;
 
 namespace Content.Shared.GameObjects.Components.Body.Mechanism
 {
@@ -36,13 +37,14 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
                 var old = _part;
                 _part = value;
 
+                if (old != null)
+                {
+                    OnRemovedFromPart(old);
+                }
+
                 if (value != null)
                 {
-                    OnPartAdd(old, value);
+                    OnAddedToPart();
-                }
-                else if (old != null)
-                {
-                    OnPartRemove(old);
                 }
             }
         }
@@ -86,18 +88,87 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
             serializer.DataField(this, m => m.Compatibility, "compatibility", BodyPartCompatibility.Universal);
         }
 
-        public virtual void OnBodyAdd(IBody? old, IBody current) { }
+        public void AddedToBody()
-
-        public virtual void OnBodyRemove(IBody old) { }
-
-        protected virtual void OnPartAdd(IBodyPart? old, IBodyPart current)
         {
-            Owner.Transform.AttachParent(current.Owner);
+            DebugTools.AssertNotNull(Body);
+
+            OnAddedToBody();
+
+            foreach (var behavior in Owner.GetMechanismBehaviors())
+            {
+                behavior.AddedToBody();
+            }
         }
 
-        protected virtual void OnPartRemove(IBodyPart old)
+        public void RemovedFromBody(IBody old)
+        {
+            OnRemovedFromBody(old);
+
+            foreach (var behavior in Owner.GetMechanismBehaviors())
+            {
+                behavior.RemovedFromBody(old);
+            }
+        }
+
+        public void AddedToPart()
+        {
+            DebugTools.AssertNotNull(Part);
+
+            Owner.Transform.AttachParent(Part!.Owner);
+            OnAddedToPart();
+
+            foreach (var behavior in Owner.GetMechanismBehaviors())
+            {
+                behavior.AddedToPart();
+            }
+        }
+
+        public void RemovedFromPart(IBodyPart old)
         {
             Owner.Transform.AttachToGridOrMap();
+            OnRemovedFromPart(old);
+
+            foreach (var behavior in Owner.GetMechanismBehaviors())
+            {
+                behavior.RemovedFromPart(old);
+            }
         }
+
+        public void AddedToPartInBody()
+        {
+            DebugTools.AssertNotNull(Body);
+            DebugTools.AssertNotNull(Part);
+
+            Owner.Transform.AttachParent(Part!.Owner);
+            OnAddedToPartInBody();
+
+            foreach (var behavior in Owner.GetMechanismBehaviors())
+            {
+                behavior.AddedToPartInBody();
+            }
+        }
+
+        public void RemovedFromPartInBody(IBody? oldBody, IBodyPart? oldPart)
+        {
+            Owner.Transform.AttachToGridOrMap();
+            OnRemovedFromPartInBody();
+
+            foreach (var behavior in Owner.GetMechanismBehaviors())
+            {
+                behavior.RemovedFromPartInBody(oldBody, oldPart);
+            }
+        }
+
+        protected virtual void OnAddedToBody() { }
+
+        protected virtual void OnRemovedFromBody(IBody old) { }
+
+        protected virtual void OnAddedToPart() { }
+
+        protected virtual void OnRemovedFromPart(IBodyPart old) { }
+
+        protected virtual void OnAddedToPartInBody() { }
+
+        protected virtual void OnRemovedFromPartInBody() { }
     }
 }

Content.Shared/GameObjects/Components/Body/Part/SharedBodyPartComponent.cs

@@ -43,18 +43,19 @@ namespace Content.Shared.GameObjects.Components.Body.Part
                 var old = _body;
                 _body = value;
 
+                if (old != null)
+                {
+                    foreach (var mechanism in _mechanisms)
+                    {
+                        mechanism.RemovedFromBody(old);
+                    }
+                }
+
                 if (value != null)
                 {
                     foreach (var mechanism in _mechanisms)
                     {
-                        mechanism.OnBodyAdd(old, value);
+                        mechanism.AddedToBody();
-                    }
-                }
-                else if (old != null)
-                {
-                    foreach (var mechanism in _mechanisms)
-                    {
-                        mechanism.OnBodyRemove(old);
                     }
                 }
             }
@@ -113,6 +114,15 @@ namespace Content.Shared.GameObjects.Components.Body.Part
             mechanism.Part = this;
             SizeUsed += mechanism.Size;
 
+            if (Body == null)
+            {
+                mechanism.AddedToPart();
+            }
+            else
+            {
+                mechanism.AddedToPartInBody();
+            }
+
             Dirty();
         }
 
@@ -122,6 +132,15 @@ namespace Content.Shared.GameObjects.Components.Body.Part
             mechanism.Part = null;
             SizeUsed -= mechanism.Size;
 
+            if (Body == null)
+            {
+                mechanism.RemovedFromPart(this);
+            }
+            else
+            {
+                mechanism.RemovedFromPartInBody(Body, this);
+            }
+
             Dirty();
         }
 
@@ -305,7 +324,7 @@ namespace Content.Shared.GameObjects.Components.Body.Part
     [Serializable, NetSerializable]
     public class BodyPartComponentState : ComponentState
     {
-        private List<IMechanism>? _mechanisms;
+        [NonSerialized] private List<IMechanism>? _mechanisms;
 
         public readonly EntityUid[] MechanismIds;