diff --git a/encompass-cs/ComponentMessageManager.cs b/encompass-cs/ComponentMessageManager.cs index 2ecefdd..432c369 100644 --- a/encompass-cs/ComponentMessageManager.cs +++ b/encompass-cs/ComponentMessageManager.cs @@ -69,7 +69,7 @@ namespace Encompass internal void AddExistingComponentMessage(ComponentMessage componentMessage) where TComponent : struct, IComponent { - RegisterExistingOrPendingComponentMessage(componentMessage.componentID, componentMessage.component); + RegisterExistingOrPendingComponentMessage(componentMessage.entity, componentMessage.componentID, componentMessage.component); if (!componentMessageTypeToExistingComponentIDs.ContainsKey(typeof(TComponent))) { @@ -92,7 +92,7 @@ namespace Encompass internal void AddPendingComponentMessage(PendingComponentMessage pendingComponentMessage) where TComponent : struct, IComponent { - RegisterExistingOrPendingComponentMessage(pendingComponentMessage.componentID, pendingComponentMessage.component); + RegisterExistingOrPendingComponentMessage(pendingComponentMessage.entity, pendingComponentMessage.componentID, pendingComponentMessage.component); if (!componentMessageTypeToPendingComponentIDs.ContainsKey(typeof(TComponent))) { @@ -113,14 +113,26 @@ namespace Encompass entityToTypeToPendingComponentIDs[pendingComponentMessage.entity][typeof(TComponent)].Add(pendingComponentMessage.componentID); } - private void RegisterExistingOrPendingComponentMessage(Guid componentID, TComponent component) where TComponent: struct, IComponent + private void RegisterExistingOrPendingComponentMessage(Entity entity, Guid componentID, TComponent component) where TComponent: struct, IComponent { componentIDToComponent[componentID] = component; + if (!componentMessageTypeToComponentIDs.ContainsKey(typeof(TComponent))) { componentMessageTypeToComponentIDs.Add(typeof(TComponent), new HashSet()); } componentMessageTypeToComponentIDs[typeof(TComponent)].Add(componentID); + + if (!entityToTypeToComponentIDs.ContainsKey(entity)) + { + entityToTypeToComponentIDs.Add(entity, new Dictionary>()); + } + if (!entityToTypeToComponentIDs[entity].ContainsKey(typeof(TComponent))) + { + entityToTypeToComponentIDs[entity].Add(typeof(TComponent), new HashSet()); + } + + entityToTypeToComponentIDs[entity][typeof(TComponent)].Add(componentID); } // general component reads by type