From b8d3d706ad23703ac70aca85035215738f88a46f Mon Sep 17 00:00:00 2001
From: thatcosmonaut <evan@moonside.games>
Date: Tue, 20 Aug 2019 15:44:01 -0700
Subject: [PATCH] fix crash when entity is added and removed same frame

---
 encompass-cs/ComponentManager.cs | 13 +++++++++--
 encompass-cs/Engine.cs           |  2 +-
 test/EngineTest.cs               | 40 ++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/encompass-cs/ComponentManager.cs b/encompass-cs/ComponentManager.cs
index 78539e8..7c34369 100644
--- a/encompass-cs/ComponentManager.cs
+++ b/encompass-cs/ComponentManager.cs
@@ -20,6 +20,7 @@ namespace Encompass
         private readonly Dictionary<Type, HashSet<Guid>> typeToComponentIDs = new Dictionary<Type, HashSet<Guid>>();
 
         private readonly List<(Entity, Type, Guid, IComponent)> componentAddData = new List<(Entity, Type, Guid, IComponent)>();
+        private readonly HashSet<Guid> componentIDsMarkedForAdd = new HashSet<Guid>();
         private readonly HashSet<Guid> componentsMarkedForRemoval = new HashSet<Guid>();
         private readonly Dictionary<Guid, IComponent> pendingUpdates = new Dictionary<Guid, IComponent>();
 
@@ -43,7 +44,7 @@ namespace Encompass
         {
             var id = NextID();
             componentAddData.Add((entity, typeof(TComponent), id, component));
-            // add these here so entity and component lookups dont break on pending components
+            componentIDsMarkedForAdd.Add(id);
             return id;
         }
 
@@ -85,6 +86,7 @@ namespace Encompass
             }
 
             componentAddData.Clear();
+            componentIDsMarkedForAdd.Clear();
         }
 
         internal IEnumerable<Guid> GetComponentIDsByEntityID(Guid entityID)
@@ -182,7 +184,14 @@ namespace Encompass
         {
             foreach (var componentID in componentsMarkedForRemoval)
             {
-                Remove(componentID);
+                if (componentIDsMarkedForAdd.Contains(componentID))
+                {
+                    componentIDsMarkedForAdd.Remove(componentID);
+                }
+                else
+                {
+                    Remove(componentID);
+                }
             }
 
             componentsMarkedForRemoval.Clear();
diff --git a/encompass-cs/Engine.cs b/encompass-cs/Engine.cs
index 5ba8489..0ae92a8 100644
--- a/encompass-cs/Engine.cs
+++ b/encompass-cs/Engine.cs
@@ -348,7 +348,7 @@ namespace Encompass
 
         protected TMessage ReadMessage<TMessage>() where TMessage : struct, IMessage
         {
-            return ReadMessages<TMessage>().Single();
+            return ReadMessages<TMessage>().First();
         }
 
         protected bool SomeMessage<TMessage>() where TMessage : struct, IMessage
diff --git a/test/EngineTest.cs b/test/EngineTest.cs
index 381be18..79ad717 100644
--- a/test/EngineTest.cs
+++ b/test/EngineTest.cs
@@ -804,5 +804,45 @@ namespace Tests
             resultMessages.Should().NotBeEmpty();
             resultMessages.First().Should().BeOfType<MockMessage>();
         }
+
+        [Receives(typeof(MockMessage))]
+        [Activates(typeof(MockComponent))]
+        class ActivateComponentEngine : Engine
+        {
+            public override void Update(double dt)
+            {
+                foreach (var message in ReadMessages<MockMessage>())
+                {
+                    var entity = CreateEntity();
+                    AddComponent(entity, new MockComponent {});
+                }
+            }
+        }
+
+        [ReadsPending(typeof(MockComponent))]
+        class RemoveComponentEngine : Engine
+        {
+            public override void Update(double dt)
+            {
+                foreach (var (componentID, component) in ReadComponents<MockComponent>())
+                {
+                    RemoveComponent(componentID);
+                }
+            }
+        }
+
+        [Test]
+        public void EngineAddAndRemoveComponentSameFrame()
+        {
+            var worldBuilder = new WorldBuilder();
+            worldBuilder.AddEngine(new ActivateComponentEngine());
+            worldBuilder.AddEngine(new RemoveComponentEngine());
+
+            worldBuilder.SendMessage(new MockMessage {});
+
+            var world = worldBuilder.Build();
+
+            Assert.DoesNotThrow(() => world.Update(0.01));
+        }
     }
 }