Better gameobject initialization

workers/unity/Packages/io.improbable.gdk.gameobjectcreation/EntityTypeExpectations.cs

@@ -0,0 +1,36 @@
+using System;
+using System.Collections.Generic;
+
+namespace Improbable.Gdk.GameObjectCreation
+{
+    /// <summary>
+    ///     Holds mappings from entity types to the set of components expected before creating their GameObjects.
+    /// </summary>
+    public class EntityTypeExpectations
+    {
+        private Type[] defaultExpectation;
+
+        private readonly Dictionary<string, Type[]> entityExpectations
+            = new Dictionary<string, Type[]>();
+
+        public void RegisterDefault(Type[] defaultComponentTypes)
+        {
+            defaultExpectation = defaultComponentTypes;
+        }
+
+        public void RegisterEntityType(string entityType, Type[] expectedComponentTypes)
+        {
+            entityExpectations.Add(entityType, expectedComponentTypes ?? new Type[] { });
+        }
+
+        internal Type[] GetExpectedTypes(string entityType)
+        {
+            if (!entityExpectations.TryGetValue(entityType, out var types))
+            {
+                return defaultExpectation;
+            }
+
+            return types;
+        }
+    }
+}

workers/unity/Packages/io.improbable.gdk.gameobjectcreation/GameObjectCreatorFromMetadata.cs

@@ -3,7 +3,6 @@
 using System.IO;
 using Improbable.Gdk.Core;
 using Improbable.Gdk.Subscriptions;
-using Unity.Entities;
 using UnityEngine;
 using Object = UnityEngine.Object;
 
@@ -17,39 +16,32 @@ private readonly Dictionary<string, GameObject> cachedPrefabs
         private readonly string workerType;
         private readonly Vector3 workerOrigin;
 
-        private readonly ILogDispatcher logger;
-
         private readonly Dictionary<EntityId, GameObject> entityIdToGameObject = new Dictionary<EntityId, GameObject>();
 
         private readonly Type[] componentsToAdd =
         {
-            typeof(Transform),
-            typeof(Rigidbody),
-            typeof(MeshRenderer)
-        };
-
-        public ComponentType[] MinimumComponentTypes { get; } =
-        {
-            ComponentType.ReadOnly<Metadata.Component>(),
-            ComponentType.ReadOnly<Position.Component>()
+            typeof(Transform), typeof(Rigidbody), typeof(MeshRenderer)
         };
 
         public GameObjectCreatorFromMetadata(string workerType, Vector3 workerOrigin, ILogDispatcher logger)
         {
             this.workerType = workerType;
             this.workerOrigin = workerOrigin;
-            this.logger = logger;
         }
 
-        public void OnEntityCreated(SpatialOSEntity entity, EntityGameObjectLinker linker)
+        public void PopulateEntityTypeExpectations(EntityTypeExpectations entityTypeExpectations)
         {
-            if (!entity.TryGetComponent<Metadata.Component>(out var metadata) ||
-                !entity.TryGetComponent<Position.Component>(out var spatialOSPosition))
+            entityTypeExpectations.RegisterDefault(new[]
             {
-                return;
-            }
+                typeof(Position.Component)
+            });
+        }
+
+        public void OnEntityCreated(string entityType, SpatialOSEntity entity, EntityGameObjectLinker linker)
+        {
+            var spatialOSPosition = entity.GetComponent<Position.Component>();
 
-            var prefabName = metadata.EntityType;
+            var prefabName = entityType;
             var position = spatialOSPosition.Coords.ToUnityVector() + workerOrigin;
 
             if (!cachedPrefabs.TryGetValue(prefabName, out var prefab))

workers/unity/Packages/io.improbable.gdk.gameobjectcreation/GameObjectInitializationSystem.cs

@@ -1,4 +1,3 @@
-using System.Linq;
 using Improbable.Gdk.Core;
 using Improbable.Gdk.Subscriptions;
 using Unity.Entities;
@@ -7,7 +6,7 @@
 namespace Improbable.Gdk.GameObjectCreation
 {
     /// <summary>
-    ///     For each newly added SpatialOS entity, calls IEntityGameObjectCreator to get an associated GameObject
+    ///     For each newly added SpatialOS entity, calls IEntityGagit meObjectCreator to get an associated GameObject
     ///     and links it to the entity via the EntityGameObjectLinker. Also checks for entity removal and calls the
     ///     IEntityGameObjectCreator for cleanup.
     /// </summary>
@@ -26,23 +25,14 @@ internal class GameObjectInitializationSystem : ComponentSystem
         private EntityQuery newEntitiesQuery;
         private EntityQuery removedEntitiesQuery;
 
-        private ComponentType[] minimumComponentSet = new[]
-        {
-            ComponentType.ReadOnly<SpatialEntityId>()
-        };
+        private readonly EntityTypeExpectations entityTypeExpectations = new EntityTypeExpectations();
 
         public GameObjectInitializationSystem(IEntityGameObjectCreator gameObjectCreator, GameObject workerGameObject)
         {
             this.gameObjectCreator = gameObjectCreator;
             this.workerGameObject = workerGameObject;
 
-            var minCreatorComponentSet = gameObjectCreator.MinimumComponentTypes;
-            if (minCreatorComponentSet != null)
-            {
-                minimumComponentSet = minimumComponentSet
-                    .Concat(minCreatorComponentSet)
-                    .ToArray();
-            }
+            gameObjectCreator.PopulateEntityTypeExpectations(entityTypeExpectations);
         }
 
         protected override void OnCreate()
@@ -58,6 +48,11 @@ protected override void OnCreate()
                 Linker.LinkGameObjectToSpatialOSEntity(new EntityId(0), workerGameObject);
             }
 
+            var minimumComponentSet = new[]
+            {
+                ComponentType.ReadOnly<SpatialEntityId>(), ComponentType.ReadOnly<Metadata.Component>()
+            };
+
             newEntitiesQuery = GetEntityQuery(new EntityQueryDesc()
             {
                 All = minimumComponentSet,
@@ -87,15 +82,19 @@ protected override void OnDestroy()
 
         protected override void OnUpdate()
         {
-            Entities.With(newEntitiesQuery).ForEach((Entity entity, ref SpatialEntityId spatialEntityId) =>
+            if (!newEntitiesQuery.IsEmptyIgnoreFilter)
             {
-                gameObjectCreator.OnEntityCreated(new SpatialOSEntity(entity, EntityManager), Linker);
-                PostUpdateCommands.AddComponent(entity, new GameObjectInitSystemStateComponent
-                {
-                    EntityId = spatialEntityId.EntityId
-                });
-            });
+                ProcessNewEntities();
+            }
 
+            if (!removedEntitiesQuery.IsEmptyIgnoreFilter)
+            {
+                ProcessRemovedEntities();
+            }
+        }
+
+        private void ProcessRemovedEntities()
+        {
             Entities.With(removedEntitiesQuery).ForEach((ref GameObjectInitSystemStateComponent state) =>
             {
                 Linker.UnlinkAllGameObjectsFromEntityId(state.EntityId);
@@ -106,5 +105,32 @@ protected override void OnUpdate()
 
             EntityManager.RemoveComponent<GameObjectInitSystemStateComponent>(removedEntitiesQuery);
         }
+
+        private void ProcessNewEntities()
+        {
+            Entities.With(newEntitiesQuery).ForEach(
+                (Entity entity, ref SpatialEntityId spatialEntityId, ref Metadata.Component metadata) =>
+                {
+                    var entityType = metadata.EntityType;
+                    var expectedTypes = entityTypeExpectations.GetExpectedTypes(entityType);
+
+                    if (expectedTypes != null)
+                    {
+                        foreach (var expectedType in expectedTypes)
+                        {
+                            if (!EntityManager.HasComponent(entity, expectedType))
+                            {
+                                return;
+                            }
+                        }
+                    }
+
+                    gameObjectCreator.OnEntityCreated(entityType, new SpatialOSEntity(entity, EntityManager), Linker);
+                    PostUpdateCommands.AddComponent(entity, new GameObjectInitSystemStateComponent
+                    {
+                        EntityId = spatialEntityId.EntityId
+                    });
+                });
+        }
     }
 }

workers/unity/Packages/io.improbable.gdk.gameobjectcreation/IEntityGameObjectCreator.cs

@@ -11,17 +11,17 @@ namespace Improbable.Gdk.GameObjectCreation
     public interface IEntityGameObjectCreator
     {
         /// <summary>
-        ///     The minimum set of components required on an entity to create a GameObject.
+        ///     Called to register the components expected on an entity to create a GameObject for a given entity type.
         /// </summary>
-        ComponentType[] MinimumComponentTypes { get; }
+        void PopulateEntityTypeExpectations(EntityTypeExpectations entityTypeExpectations);
 
         /// <summary>
         ///     Called when a new SpatialOS Entity is checked out by the worker.
         /// </summary>
         /// <returns>
         ///     A GameObject to be linked to the entity, or null if no GameObject should be linked.
         /// </returns>
-        void OnEntityCreated(SpatialOSEntity entity, EntityGameObjectLinker linker);
+        void OnEntityCreated(string entityType, SpatialOSEntity entity, EntityGameObjectLinker linker);
 
         /// <summary>
         ///     Called when a SpatialOS Entity is removed from the worker's view.

workers/unity/Packages/io.improbable.gdk.gameobjectcreation/Tests/Editmode/GameObjectCreationTests.cs

@@ -115,18 +115,20 @@ private class TestGameObjectCreator : IEntityGameObjectCreator
 
             private readonly Dictionary<EntityId, GameObject> entityIdToGameObject = new Dictionary<EntityId, GameObject>();
 
-            public ComponentType[] MinimumComponentTypes { get; } =
-            {
-                ComponentType.ReadOnly<Position.Component>(),
-                ComponentType.ReadOnly<Metadata.Component>()
-            };
-
             public TestGameObjectCreator(string workerType)
             {
                 this.workerType = workerType;
             }
 
-            public void OnEntityCreated(SpatialOSEntity entity, EntityGameObjectLinker linker)
+            public void PopulateEntityTypeExpectations(EntityTypeExpectations entityTypeExpectations)
+            {
+                entityTypeExpectations.RegisterDefault(new[]
+                {
+                    typeof(Metadata.Component), typeof(Position.Component)
+                });
+            }
+
+            public void OnEntityCreated(string entityType, SpatialOSEntity entity, EntityGameObjectLinker linker)
             {
                 var gameObject = new GameObject();
                 gameObject.transform.position = Vector3.one;

workers/unity/Packages/io.improbable.gdk.gameobjectcreation/Tests/Editmode/MockGameObjectCreator.cs

@@ -7,16 +7,18 @@ namespace Improbable.Gdk.GameObjectCreation.EditmodeTests
 {
     public class MockGameObjectCreator : IEntityGameObjectCreator
     {
-        public ComponentType[] MinimumComponentTypes { get; } = { };
+        public void PopulateEntityTypeExpectations(EntityTypeExpectations entityTypeExpectations)
+        {
+        }
 
-        public void OnEntityCreated(SpatialOSEntity entity, EntityGameObjectLinker linker)
+        public void OnEntityCreated(string entityType, SpatialOSEntity entity, EntityGameObjectLinker linker)
         {
-            throw new System.NotImplementedException();
+            throw new NotImplementedException();
         }
 
         public void OnEntityRemoved(EntityId entityId)
         {
-            throw new System.NotImplementedException();
+            throw new NotImplementedException();
         }
     }
 }