From ebe6093af1d0ee008d6c908c67f9bb9f2019eef2 Mon Sep 17 00:00:00 2001
From: Leon Friedrich <60421075+ElectroJr@users.noreply.github.com>
Date: Fri, 29 Jul 2022 14:13:43 +1200
Subject: [PATCH] Pulling state handling changes (#9952)

---
 .../Physics/Controllers/MoverController.cs    |  9 ++--
 .../Pulling/Components/PullableComponent.cs   |  5 +--
 .../SharedPullingStateManagementSystem.cs     | 45 +++++++++++--------
 .../Pulling/Systems/SharedPullingSystem.cs    |  2 +-
 4 files changed, 35 insertions(+), 26 deletions(-)

diff --git a/Content.Client/Physics/Controllers/MoverController.cs b/Content.Client/Physics/Controllers/MoverController.cs
index aa899bea4e..d9697ddea0 100644
--- a/Content.Client/Physics/Controllers/MoverController.cs
+++ b/Content.Client/Physics/Controllers/MoverController.cs
@@ -71,10 +71,13 @@ namespace Content.Client.Physics.Controllers
 
             if (TryComp(player, out JointComponent? jointComponent))
             {
-                foreach (var joint in jointComponent.GetJoints)
+                foreach (var joint in jointComponent.GetJoints.Values)
                 {
-                    joint.BodyA.Predict = true;
-                    joint.BodyB.Predict = true;
+                    if (TryComp(joint.BodyAUid, out PhysicsComponent? physics))
+                        physics.Predict = true;
+
+                    if (TryComp(joint.BodyBUid, out physics))
+                        physics.Predict = true;
                 }
             }
 
diff --git a/Content.Shared/Pulling/Components/PullableComponent.cs b/Content.Shared/Pulling/Components/PullableComponent.cs
index a822e67910..7b1c77bef9 100644
--- a/Content.Shared/Pulling/Components/PullableComponent.cs
+++ b/Content.Shared/Pulling/Components/PullableComponent.cs
@@ -13,14 +13,13 @@ namespace Content.Shared.Pulling.Components
     {
         /// <summary>
         /// The current entity pulling this component.
-        /// SharedPullingStateManagementSystem should be writing this. This means definitely not you.
         /// </summary>
         public EntityUid? Puller { get; set; }
+
         /// <summary>
         /// The pull joint.
-        /// SharedPullingStateManagementSystem should be writing this. This means probably not you.
         /// </summary>
-        public DistanceJoint? PullJoint { get; set; }
+        public string? PullJointId { get; set; }
 
         public bool BeingPulled => Puller != null;
 
diff --git a/Content.Shared/Pulling/Systems/SharedPullingStateManagementSystem.cs b/Content.Shared/Pulling/Systems/SharedPullingStateManagementSystem.cs
index 5d1d19aa35..4cbdd13f54 100644
--- a/Content.Shared/Pulling/Systems/SharedPullingStateManagementSystem.cs
+++ b/Content.Shared/Pulling/Systems/SharedPullingStateManagementSystem.cs
@@ -4,6 +4,7 @@ using Content.Shared.Pulling.Components;
 using JetBrains.Annotations;
 using Robust.Shared.Map;
 using Robust.Shared.Physics;
+using Robust.Shared.Timing;
 
 namespace Content.Shared.Pulling
 {
@@ -16,6 +17,7 @@ namespace Content.Shared.Pulling
     {
         [Dependency] private readonly SharedJointSystem _jointSystem = default!;
         [Dependency] private readonly SharedPhysicsSystem _physics = default!;
+        [Dependency] private readonly IGameTiming _timing = default!;
 
         public override void Initialize()
         {
@@ -42,14 +44,14 @@ namespace Content.Shared.Pulling
             ForceSetMovingTo(pullable, null);
 
             // Joint shutdown
-            if (EntityManager.TryGetComponent<JointComponent?>(puller.Owner, out var jointComp))
+            if (!_timing.ApplyingState && // During state-handling, joint component will handle its own state.
+                pullable.PullJointId != null &&
+                TryComp(puller.Owner, out JointComponent? jointComp))
             {
-                if (jointComp.GetJoints.Contains(pullable.PullJoint!))
-                {
-                    _jointSystem.RemoveJoint(pullable.PullJoint!);
-                }
+                if (jointComp.GetJoints.TryGetValue(pullable.PullJointId, out var j))
+                    _jointSystem.RemoveJoint(j);
             }
-            pullable.PullJoint = null;
+            pullable.PullJointId = null;
 
             // State shutdown
             puller.Pulling = null;
@@ -64,8 +66,8 @@ namespace Content.Shared.Pulling
                 RaiseLocalEvent(pullable.Owner, message, true);
 
             // Networking
-            puller.Dirty();
-            pullable.Dirty();
+            Dirty(puller);
+            Dirty(pullable);
         }
 
         public void ForceRelationship(SharedPullerComponent? puller, SharedPullableComponent? pullable)
@@ -92,26 +94,31 @@ namespace Content.Shared.Pulling
 
             // And now for the actual connection (if any).
 
-            if ((puller != null) && (pullable != null))
+            if (puller != null && pullable != null)
             {
                 var pullerPhysics = EntityManager.GetComponent<PhysicsComponent>(puller.Owner);
                 var pullablePhysics = EntityManager.GetComponent<PhysicsComponent>(pullable.Owner);
+                pullable.PullJointId = $"pull-joint-{pullable.Owner}";
 
                 // State startup
                 puller.Pulling = pullable.Owner;
                 pullable.Puller = puller.Owner;
 
-                // Joint startup
-                var union = _physics.GetHardAABB(pullerPhysics).Union(_physics.GetHardAABB(pullablePhysics));
-                var length = Math.Max(union.Size.X, union.Size.Y) * 0.75f;
+                // joint state handling will manage its own state
+                if (!_timing.ApplyingState)
+                {
+                    // Joint startup
+                    var union = _physics.GetHardAABB(pullerPhysics).Union(_physics.GetHardAABB(pullablePhysics));
+                    var length = Math.Max(union.Size.X, union.Size.Y) * 0.75f;
 
-                pullable.PullJoint = _jointSystem.CreateDistanceJoint(pullablePhysics.Owner, pullerPhysics.Owner, id:$"pull-joint-{pullablePhysics.Owner}");
-                pullable.PullJoint.CollideConnected = false;
-                // This maximum has to be there because if the object is constrained too closely, the clamping goes backwards and asserts.
-                pullable.PullJoint.MaxLength = Math.Max(1.0f, length);
-                pullable.PullJoint.Length = length * 0.75f;
-                pullable.PullJoint.MinLength = 0f;
-                pullable.PullJoint.Stiffness = 1f;
+                    var joint = _jointSystem.CreateDistanceJoint(pullablePhysics.Owner, pullerPhysics.Owner, id: pullable.PullJointId);
+                    joint.CollideConnected = false;
+                    // This maximum has to be there because if the object is constrained too closely, the clamping goes backwards and asserts.
+                    joint.MaxLength = Math.Max(1.0f, length);
+                    joint.Length = length * 0.75f;
+                    joint.MinLength = 0f;
+                    joint.Stiffness = 1f;
+                }
 
                 // Messaging
                 var message = new PullStartedMessage(pullerPhysics, pullablePhysics);
diff --git a/Content.Shared/Pulling/Systems/SharedPullingSystem.cs b/Content.Shared/Pulling/Systems/SharedPullingSystem.cs
index 83e7abe9a2..55b35688c3 100644
--- a/Content.Shared/Pulling/Systems/SharedPullingSystem.cs
+++ b/Content.Shared/Pulling/Systems/SharedPullingSystem.cs
@@ -65,7 +65,7 @@ namespace Content.Shared.Pulling
 
             if (TryComp(uid, out JointComponent? joints))
             {
-                foreach (var jt in joints.GetJoints)
+                foreach (var jt in joints.GetJoints.Values)
                 {
                     if (jt.BodyAUid == component.Puller || jt.BodyBUid == component.Puller)
                         return;