Move entity init support into base

And other refactorings and fixes

TechbloxModdingAPI/Block.cs

@@ -28,7 +28,7 @@ namespace TechbloxModdingAPI
 	/// A single (perhaps scaled) block. Properties may return default values if the block is removed and then setting them is ignored.
 	/// For specific block type operations, use the specialised block classes in the TechbloxModdingAPI.Blocks namespace.
 	/// </summary>
-	public class Block : EcsObjectBase, IHasPhysics, IEquatable<Block>, IEquatable<EGID>
+	public class Block : EcsObjectBase<BlockEntityDescriptor>, IHasPhysics, IEquatable<Block>, IEquatable<EGID>
 	{
 		protected static readonly PlacementEngine PlacementEngine = new();
 		protected static readonly RemovalEngine RemovalEngine = new();
@@ -54,9 +54,7 @@ namespace TechbloxModdingAPI
 			if (PlacementEngine.IsInGame && GameClient.IsBuildMode)
 			{
 				var initializer = PlacementEngine.PlaceBlock(block, position, player, autoWire);
-				var egid = initializer.EGID;
-				var bl = New(egid);
-				bl.InitData = initializer;
+				var bl = New(initializer);
 				Placed += bl.OnPlacedInit;
 				return bl;
 			}
@@ -123,20 +121,31 @@ namespace TechbloxModdingAPI
 		/// <param name="signaling">Whether the block is definitely a signaling block</param>
 		/// <returns></returns>
 		internal static Block New(EGID egid, bool signaling = false)
+		{
+			return New(egid, default, signaling);
+		}
+
+		private static Block New(EcsInitData initData, bool signaling = false)
+		{
+			return New(initData.EGID, initData, signaling);
+		}
+
+		private static Block New(EGID egid, EcsInitData initData, bool signaling)
 		{
 			if (egid == default) return null;
+			Func<EGID, Block> constructor;
+			Type type = null;
 			if (GroupToConstructor.TryGetValue(egid.groupID, out var value))
-			{
-				var (constructor, type) = value;
-				return GetInstance(egid, constructor, type);
-			}
+				(constructor, type) = value;
+			else
+				constructor = signaling ? e => new SignalingBlock(e) : e => new Block(e);
 
-			return signaling
-				? GetInstance(egid, e => new SignalingBlock(e))
-				: GetInstance(egid, e => new Block(e));
+			return initData != default
+				? GetInstanceNew(initData, constructor, type)
+				: GetInstanceExisting(egid, constructor, type);
 		}
 
-		public Block(EGID id) : base(id, typeof(BlockEntityDescriptor))
+		public Block(EGID id) : base(id)
 		{
 			Type expectedType;
 			if (GroupToConstructor.ContainsKey(id.groupID) &&
@@ -375,7 +384,7 @@ namespace TechbloxModdingAPI
 				var bgec = GetComponent<BlockGroupEntityComponent>();
 				return blockGroup = bgec.currentBlockGroup == -1
 					? null
-					: GetInstance(new EGID((uint)bgec.currentBlockGroup, BlockGroupExclusiveGroups.BlockGroupEntityGroup),
+					: GetInstanceExisting(new EGID((uint)bgec.currentBlockGroup, BlockGroupExclusiveGroups.BlockGroupEntityGroup),
 						egid => new BlockGroup((int)egid.entityID, this));
 			}
 			set
@@ -466,7 +475,6 @@ namespace TechbloxModdingAPI
 		{ //Member method instead of lambda to avoid constantly creating delegates
 			if (e.ID != Id) return;
 			Placed -= OnPlacedInit; //And we can reference it
-			InitData = default; //Remove initializer as it's no longer valid - if the block gets removed it shouldn't be used again
 			if (copiedFrom != default)
 				BlockCloneEngine.CopyBlockStats(copiedFrom, Id);
 		}

TechbloxModdingAPI/Blocks/Wire.cs

@@ -218,7 +218,7 @@ namespace TechbloxModdingAPI.Blocks
         /// <returns>A copy of the wire object.</returns>
         public Wire OutputToInputCopy()
         {
-            return GetInstance(wireEGID, egid => new Wire(egid));
+            return GetInstanceExisting(wireEGID, egid => new Wire(egid));
         }
         
         /// <summary>

TechbloxModdingAPI/Common/EcsObjectBase.cs

@@ -10,8 +10,25 @@ using TechbloxModdingAPI.Utility;
 
 namespace TechbloxModdingAPI.Common;
 
-public abstract class EcsObjectBase
+public abstract class EcsObjectBase<TDescriptor> : EcsObjectBase where TDescriptor : IEntityDescriptor, new()
 {
+    protected EcsObjectBase(EGID id) : base(id, typeof(TDescriptor))
+    {
+    }
+
+    protected EcsObjectBase(EntityReference reference) : base(reference, typeof(TDescriptor))
+    {
+    }
+
+    protected bool RemoveEntity()
+    {
+        if (!Exists) return false;
+        _engine.Functions.RemoveEntity<TDescriptor>(Id);
+        return true;
+    }
+}
+
+public abstract class EcsObjectBase {
     public EGID Id => _engine.GetEgid(Reference);
     /// <summary>
     /// A reference to a specific entity that persists through group swaps and such.
@@ -23,9 +40,12 @@ public abstract class EcsObjectBase
     /// Whether the entity reference is still valid. Returns false if this object no longer exists.
     /// </summary>
     public bool Exists => Id != default; // TODO: Might need extra code to support IDs during init
+
+    public readonly Type EntityDescriptorType;
+    public readonly Type[] AllowedEntityComponents;
         
     private static readonly Dictionary<Type, WeakDictionary<EntityReference, EcsObjectBase>> _instances = new();
-    private static readonly EcsObjectBaseEngine _engine = new();
+    internal static readonly EcsObjectBaseEngine _engine = new();
 
     private static WeakDictionary<EntityReference, EcsObjectBase> GetInstances(Type type)
     {
@@ -34,13 +54,14 @@ public abstract class EcsObjectBase
 
     /// <summary>
     /// Returns a cached instance if there's an actively used instance of the object already.
-    /// Objects still get garbage collected and then they will be removed from the cache.
+    /// Objects still get garbage collected and then they will be removed from the cache.<br />
+    /// <b>Only use for existing entities!</b> Use the other overload for newly created entities.
     /// </summary>
     /// <param name="egid">The EGID of the entity</param>
     /// <param name="constructor">The constructor to construct the object</param>
     /// <typeparam name="T">The object type</typeparam>
     /// <returns></returns>
-    internal static T GetInstance<T>(EGID egid, Func<EGID, T> constructor, Type type = null) where T : EcsObjectBase
+    internal static T GetInstanceExisting<T>(EGID egid, Func<EGID, T> constructor, Type type = null) where T : EcsObjectBase
     {
         var instances = GetInstances(type ?? typeof(T));
         if (instances == null || !instances.TryGetValue(_engine.GetEntityReference(egid), out var instance))
@@ -48,6 +69,33 @@ public abstract class EcsObjectBase
         return (T)instance;
     }
 
+    /// <summary>
+    /// Returns a cached instance if there's an actively used instance of the object already.
+    /// Objects still get garbage collected and then they will be removed from the cache.<br />
+    /// <b>Only use for newly created entities!</b> Use the other overload for existing entities.
+    /// </summary>
+    /// <param name="egid">The EGID of the entity</param>
+    /// <param name="constructor">The constructor to construct the object</param>
+    /// <typeparam name="T">The object type</typeparam>
+    /// <returns></returns>
+    internal static T GetInstanceNew<T>(EcsInitData initData, Func<EGID, T> constructor, Type type = null) where T : EcsObjectBase
+    {
+        var instances = GetInstances(type ?? typeof(T));
+        if (instances == null || !instances.TryGetValue(initData.Reference, out var instance))
+        {
+            var ret = constructor(initData.EGID);
+            ret.InitData = initData;
+            return ret; // It will be added by the constructor
+        }
+
+        return (T)instance;
+    }
+
+    protected static V CreateEntity<U, V>(EGID egid, Func<EGID, V> constructor, Type type = null) where U : IEntityDescriptor, new() where V : EcsObjectBase<U>
+    {
+        return GetInstanceNew(_engine.Factory.BuildEntity<U>(egid), constructor, type);
+    }
+
     protected EcsObjectBase(EGID id, Type entityDescriptorType) : this(_engine.GetEntityReference(id), entityDescriptorType)
     {
     }
@@ -62,9 +110,11 @@ public abstract class EcsObjectBase
         if (!dict.ContainsKey(reference)) // Multiple instances may be created
             dict.Add(reference, this);
         Reference = reference;
+        EntityDescriptorType = entityDescriptorType;
+        AllowedEntityComponents = EcsUtils.GetValidEntityComponents(entityDescriptorType);
+        // Remove init data once the entity gets submitted so that it won't be used again once the entity is removed
+        if (InitData != default) _engine.TrackNewEntity(this, obj => obj.InitData = default);
     }
-    
-    
 
     protected internal OptionalRef<T> GetComponentOptional<T>() where T : unmanaged, IEntityComponent
     {
@@ -91,28 +141,22 @@ public abstract class EcsObjectBase
         _engine.SetComponent(this, type, name, value);
     }
 
-    protected bool RemoveEntity()
-    {
-        // TODO: _entityFunctions.Remove...()
-    }
-
     #region ECS initializer stuff
 
-    protected internal EcsInitData InitData;
+    internal EcsInitData InitData { get; private set; }
 
     /// <summary>
     /// Holds information needed to construct a component initializer.
     /// Necessary because the initializer is a ref struct which cannot be assigned to a field.
     /// </summary>
-    protected internal struct EcsInitData
+    protected internal readonly record struct EcsInitData(FasterDictionary<RefWrapperType, ITypeSafeDictionary> group, EntityReference Reference, EGID EGID)
     {
-        private FasterDictionary<RefWrapperType, ITypeSafeDictionary> group;
-        private EntityReference reference;
-
-        public static implicit operator EcsInitData(EntityInitializer initializer) => new()
-            { group = GetInitGroup(initializer), reference = initializer.reference };
+        public static implicit operator EcsInitData(EntityInitializer initializer) => new(GetInitGroup(initializer), initializer.reference, initializer.EGID);
 
-        public EntityInitializer Initializer(EGID id) => new(id, group, reference);
+        private readonly FasterDictionary<RefWrapperType, ITypeSafeDictionary> group = group;
+        public readonly EntityReference Reference = Reference;
+        public readonly EGID EGID = EGID;
+        public EntityInitializer Initializer(EGID id = default) => new(id == default ? EGID : id, group, Reference);
         public bool Valid => group != null;
     }
 

TechbloxModdingAPI/Common/EcsObjectBaseEngine.cs

@@ -1,7 +1,11 @@
 using System;
+using System.Collections.Generic;
 using HarmonyLib;
+using RobocraftX.Schedulers;
 using Svelto.ECS;
 using Svelto.ECS.Hybrid;
+using Svelto.Tasks;
+using Svelto.Tasks.Lean;
 using TechbloxModdingAPI.Blocks.Engines;
 using TechbloxModdingAPI.Common.Engines;
 using TechbloxModdingAPI.Utility;
@@ -9,7 +13,7 @@ using TechbloxModdingAPI.Utility.ECS;
 
 namespace TechbloxModdingAPI.Common;
 
-public class EcsObjectBaseEngine : IApiEngine
+public class EcsObjectBaseEngine : IFactoryEngine, IFunEngine
 {
     public void Ready()
     {
@@ -68,4 +72,24 @@ public class EcsObjectBaseEngine : IApiEngine
         AccessTools.Field(str.GetType(), name).SetValue(str, value);
         prop.SetValue(opt, str);
     }
+
+    private readonly Dictionary<EcsObjectBase, Action<EcsObjectBase>> _waitingForSubmission = new();
+
+    public void TrackNewEntity(EcsObjectBase obj, Action<EcsObjectBase> done)
+    {
+        if (_waitingForSubmission.ContainsKey(obj))
+            throw new InvalidOperationException("Something has gone horribly wrong here");
+        _waitingForSubmission.Add(obj, done);
+        WaitUntilEntitySubmission().RunOn(ClientLean.UIScheduler); // TODO: Pick the right scheduler
+    }
+
+    private IEnumerator<TaskContract> WaitUntilEntitySubmission()
+    {
+        // TODO: Get the scheduler instance based on the engine (inject in engine manager)
+        yield return new WaitForSubmissionEnumerator(FullGameFields._mainGameEnginesRoot.scheduler).Continue();
+        foreach (var (obj, done) in _waitingForSubmission) done(obj);
+    }
+
+    public IEntityFactory Factory { get; set; }
+    public IEntityFunctions Functions { get; set; }
 }

TechbloxModdingAPI/Common/Engines/EngineManager.cs

@@ -17,6 +17,8 @@ public class EngineManager
     /// <param name="types">The types to register to</param>
     public static void AddEngine(IApiEngine engine, params ApiEngineType[] types)
     {
+        if (types.Length == 0)
+            Logging.LogWarning($"Engine {engine.GetType().FullName} added without any types! This doesn't do anything.");
         foreach (var type in types)
         {
             if (!_engines.ContainsKey(type))

TechbloxModdingAPI/Common/Traits/HasPhysics.cs

@@ -11,7 +11,8 @@ public interface IHasPhysics
 
 public static class HasPhysicsExtensions
 {
-    internal static void UpdatePhysicsUECSComponent<T, O>(this O obj, T componentData) where O : EcsObjectBase, IHasPhysics where T : struct, IComponentData
+    internal static void UpdatePhysicsUECSComponent<T, O>(this O obj, T componentData)
+        where O : EcsObjectBase, IHasPhysics where T : struct, IComponentData
     {
         var phyStruct = obj.GetComponentOptional<DOTSPhysicsEntityStruct>();
         if (phyStruct) //It exists

TechbloxModdingAPI/Common/Utils/EcsUtils.cs

@@ -0,0 +1,23 @@
+using System;
+using System.Linq;
+using System.Linq.Expressions;
+using HarmonyLib;
+using Svelto.ECS;
+
+namespace TechbloxModdingAPI.Common.Utils
+{
+    public static class EcsUtils
+    {
+        public static Type[] GetValidEntityComponents(Type entityDescriptorType)
+        {
+            // TODO: Cache
+            var templateType = typeof(EntityDescriptorTemplate<>).MakeGenericType(entityDescriptorType);
+            var templateDescriptor = AccessTools.Property(templateType, "descriptor");
+            var getDescriptorExpr = Expression.MakeMemberAccess(null, templateDescriptor ?? throw new InvalidOperationException());
+            var getTemplateDescriptorExpr = Expression.Lambda<Func<IEntityDescriptor>>(getDescriptorExpr);
+            var getTemplateDescriptor = getTemplateDescriptorExpr.Compile();
+            var builders = getTemplateDescriptor().componentsToBuild;
+            return builders.Select(builder => builder.GetEntityComponentType()).ToArray();
+        }
+    }
+}

TechbloxModdingAPI/Main.cs

@@ -8,6 +8,7 @@ using Svelto.Context;
 
 using TechbloxModdingAPI.App;
 using TechbloxModdingAPI.Blocks;
+using TechbloxModdingAPI.Common;
 using TechbloxModdingAPI.Tasks;
 using TechbloxModdingAPI.Utility;
 
@@ -61,6 +62,7 @@ namespace TechbloxModdingAPI
                 // init input
                 Input.FakeInput.Init();
                 // init object-oriented classes
+                EcsObjectBase.Init();
                 Player.Init();
                 Block.Init();
                 BlockGroup.Init();

TechbloxModdingAPI/Player.cs

@@ -84,7 +84,7 @@ namespace TechbloxModdingAPI
 
         internal static Player GetInstance(uint id)
         {
-	        return EcsObjectBase.GetInstance(new EGID(id, CharacterExclusiveGroups.OnFootGroup),
+	        return EcsObjectBase.GetInstanceExisting(new EGID(id, CharacterExclusiveGroups.OnFootGroup),
 		        e => new Player(e.entityID));
         }
 
@@ -469,7 +469,7 @@ namespace TechbloxModdingAPI
         {
 	        var egid = playerEngine.GetThingLookedAt(Id, maxDistance);
 	        return egid != default && egid.groupID == CommonExclusiveGroups.SIMULATION_BODIES_GROUP
-		        ? EcsObjectBase.GetInstance(egid, e => new SimBody(e))
+		        ? EcsObjectBase.GetInstanceExisting(egid, e => new SimBody(e))
 		        : null;
         }
 
@@ -482,7 +482,7 @@ namespace TechbloxModdingAPI
         {
 	        var egid = playerEngine.GetThingLookedAt(Id, maxDistance);
 	        return egid != default && egid.groupID == WiresGUIExclusiveGroups.WireGroup
-		        ? EcsObjectBase.GetInstance(new EGID(egid.entityID, BuildModeWiresGroups.WiresGroup.Group),
+		        ? EcsObjectBase.GetInstanceExisting(new EGID(egid.entityID, BuildModeWiresGroups.WiresGroup.Group),
 			        e => new Wire(e))
 		        : null;
         }

TechbloxModdingAPI/SimBody.cs

@@ -20,7 +20,7 @@ namespace TechbloxModdingAPI
         /// </summary>
         public Cluster Cluster => cluster ??= clusterId == uint.MaxValue // Return cluster or if it's null then set it
             ? Block.BlockEngine.GetCluster(Id.entityID) // If we don't have a clusterId set then get it from the game
-            : GetInstance(new EGID(clusterId, ClustersExclusiveGroups.SIMULATION_CLUSTERS_GROUP),
+            : GetInstanceExisting(new EGID(clusterId, ClustersExclusiveGroups.SIMULATION_CLUSTERS_GROUP),
                 egid => new Cluster(egid)); // Otherwise get the cluster from the ID
 
         private Cluster cluster;

TechbloxModdingAPI/Utility/ECS/ManagedApiExtensions.cs

@@ -1,3 +1,4 @@
+using System.Linq;
 using Svelto.ECS;
 using Svelto.ECS.Hybrid;
 using TechbloxModdingAPI.Common;
@@ -54,7 +55,10 @@ namespace TechbloxModdingAPI.Utility.ECS
             EGID id = group == ExclusiveGroupStruct.Invalid ? obj.Id : new EGID(obj.Id.entityID, group);
             var opt = QueryEntityOptional<T>(entitiesDB, id);
             if (opt) return ref opt.Get();
-            if (obj.InitData.Valid) return ref obj.InitData.Initializer(id).GetOrAdd<T>();
+            // If initializing the entity, check if the component is allowed by the descriptor, otherwise it could cause
+            // issues in the game with Add() calls running unexpectedly
+            if (obj.InitData.Valid && obj.AllowedEntityComponents.Contains(typeof(T)))
+                return ref obj.InitData.Initializer(id).GetOrAdd<T>();
             return ref opt.Get(); //Default value
         }
 

TechbloxModdingAPI/Utility/ECS/NativeApiExtensions.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using Svelto.ECS;
 using Svelto.Tasks;
 using Svelto.Tasks.Lean;
@@ -58,12 +59,10 @@ namespace TechbloxModdingAPI.Utility.ECS
             EGID id = group == ExclusiveGroupStruct.Invalid ? obj.Id : new EGID(obj.Id.entityID, group);
             var opt = QueryEntityOptional<T>(entitiesDB, id);
             if (opt) return ref opt.Get();
-            if (obj.InitData.Valid) return ref obj.InitData.Initializer(id).GetOrAdd<T>();
-            /*if (!obj.InitData.Valid) return ref opt.Get(); //Default value
-            var init = obj.InitData.Initializer(id);
-            // Do not create the component if missing, as that can trigger Add() listeners that, in some cases, may be
-            // invalid if (ab)using the classes in an unusual way - TODO: Check entity descriptor or something
-            if (init.Has<T>()) return ref init.Get<T>();*/
+            // If initializing the entity, check if the component is allowed by the descriptor, otherwise it could cause
+            // issues in the game with Add() calls running unexpectedly
+            if (obj.InitData.Valid && obj.AllowedEntityComponents.Contains(typeof(T)))
+                return ref obj.InitData.Initializer(id).GetOrAdd<T>();
             return ref opt.Get(); //Default value
         }
 

TechbloxModdingAPI/Utility/OptionalRef.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Linq;
 using Svelto.DataStructures;
 using Svelto.ECS;
 using TechbloxModdingAPI.Common;
@@ -14,6 +15,7 @@ namespace TechbloxModdingAPI.Utility
         private MB<T> managedArray;
         private readonly EntityInitializer initializer;
         //The possible fields are: (index && (array || managedArray)) || initializer
+        private readonly EcsObjectBase obj;
 
         public OptionalRef(NB<T> array, uint index, EGID entityId = default)
         {
@@ -23,6 +25,7 @@ namespace TechbloxModdingAPI.Utility
             this.entityId = entityId;
             initializer = default;
             managedArray = default;
+            obj = default;
         }
 
         public OptionalRef(MB<T> array, uint index, EGID entityId = default)
@@ -33,6 +36,7 @@ namespace TechbloxModdingAPI.Utility
             this.entityId = entityId;
             initializer = default;
             this.array = default;
+            obj = default;
         }
 
         /// <summary>
@@ -56,6 +60,7 @@ namespace TechbloxModdingAPI.Utility
             array = default;
             index = default;
             managedArray = default;
+            this.obj = obj;
         }
 
         /// <summary>
@@ -66,7 +71,10 @@ namespace TechbloxModdingAPI.Utility
         {
             CompRefCache.Default = default; //The default value can be changed by mods
             if (state == State.Empty) return ref CompRefCache.Default;
-            if ((state & State.Initializer) != State.Empty) return ref initializer.GetOrAdd<T>();
+            // If initializing the entity, check if the component is allowed by the descriptor, otherwise it could cause
+            // issues in the game with Add() calls running unexpectedly
+            if ((state & State.Initializer) != State.Empty && obj.AllowedEntityComponents.Contains(typeof(T)))
+                return ref initializer.GetOrAdd<T>();
             if ((state & State.Native) != State.Empty) return ref array[index];
             return ref managedArray[index];
         }