From 162dbe52a2e35485b660a5c2d7b0dc771b0a3dd3 Mon Sep 17 00:00:00 2001
From: Evan Hemsley <2342303+ehemsley@users.noreply.github.com>
Date: Sun, 11 Aug 2019 11:26:31 -0700
Subject: [PATCH] fix broken entity ID lookup on pending components

---
 encompass-cs/ComponentManager.cs        |  5 ++-
 encompass-cs/ComponentMessageManager.cs |  9 +----
 test/EngineTest.cs                      | 44 ++++++++++++++++++++++++-
 3 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/encompass-cs/ComponentManager.cs b/encompass-cs/ComponentManager.cs
index 4c77e95..22446fe 100644
--- a/encompass-cs/ComponentManager.cs
+++ b/encompass-cs/ComponentManager.cs
@@ -43,6 +43,8 @@ namespace Encompass
         {
             var id = NextID();
             componentAddData.Add((entity, typeof(TComponent), id, component));
+            IDToComponent[id] = component; // add these here so entity lookup doesnt break on pending components
+            componentIDToEntityID[id] = entity.ID;
             return id;
         }
 
@@ -55,7 +57,6 @@ namespace Encompass
 
         internal void AddComponent(Entity entity, Type type, Guid componentID, IComponent component)
         {
-            IDToComponent[componentID] = component;
             componentIDToType[componentID] = type;
 
             if (!typeToComponentIDs.ContainsKey(type))
@@ -74,8 +75,6 @@ namespace Encompass
             {
                 throw new MultipleComponentOfSameTypeException("Entity {0} cannot have multiple components of type {1}", entity.ID, type.Name);
             }
-
-            componentIDToEntityID[componentID] = entity.ID;
         }
 
         internal void AddMarkedComponents()
diff --git a/encompass-cs/ComponentMessageManager.cs b/encompass-cs/ComponentMessageManager.cs
index 0a110df..6685d07 100644
--- a/encompass-cs/ComponentMessageManager.cs
+++ b/encompass-cs/ComponentMessageManager.cs
@@ -124,14 +124,7 @@ namespace Encompass
             }
             componentMessageTypeToComponentIDs[typeof(TComponent)].Add(componentID);
 
-            if (!entityToTypeToComponentID[entity].ContainsKey(typeof(TComponent)))
-            {
-                entityToTypeToComponentID[entity].Add(typeof(TComponent), componentID);
-            }
-            else
-            {
-                throw new MultipleComponentOfSameTypeException("Entity {0} cannot have multiple components of type {1}", entity.ID, typeof(TComponent).Name);
-            }
+            entityToTypeToComponentID[entity][typeof(TComponent)] = componentID;
         }
 
         // general component reads by type
diff --git a/test/EngineTest.cs b/test/EngineTest.cs
index d40c3de..5fc58d6 100644
--- a/test/EngineTest.cs
+++ b/test/EngineTest.cs
@@ -546,11 +546,53 @@ namespace Tests
             worldBuilder.AddComponent(entity, component);
 
             var world = worldBuilder.Build();
-            world.Update(0.01f);
+            world.Update(0.01);
 
             Assert.That(entity, Is.EqualTo(entityFromComponentIDResult));
         }
 
+        [Activates(typeof(MockComponent))]
+        [Reads(typeof(MockComponent))]
+        class AddAndRemoveMockComponentEngine : Engine
+        {
+            public override void Update(double dt)
+            {
+                foreach (var (mockComponentID, mockComponent) in ReadComponents<MockComponent>())
+                {
+                    var entity = GetEntityByComponentID(mockComponentID);
+                    RemoveComponent(mockComponentID);
+                    AddComponent(entity, new MockComponent());
+                }
+            }
+        }
+
+        [ReadsPending(typeof(MockComponent))]
+        class GetEntityFromPendingComponentIDEngine : Engine
+        {
+            public override void Update(double dt)
+            {
+                foreach (var (mockComponentID, mockComponent) in ReadComponents<MockComponent>())
+                {
+                    entityFromComponentIDResult = GetEntityByComponentID(mockComponentID);
+                }
+            }
+        }
+
+        [Test]
+        public void GetEntityFromPendingComponentID()
+        {
+            var worldBuilder = new WorldBuilder();
+            worldBuilder.AddEngine(new AddAndRemoveMockComponentEngine());
+            worldBuilder.AddEngine(new GetEntityFromPendingComponentIDEngine());
+
+            var entity = worldBuilder.CreateEntity();
+            worldBuilder.AddComponent(entity, new MockComponent());
+
+            var world = worldBuilder.Build();
+
+            Assert.DoesNotThrow(() => world.Update(0.01));
+        }
+
         static MockComponent mockComponentByIDResult;
 
         [Reads(typeof(MockComponent))]