EntityStore.CopyEntity() - update indexes when copy indexed components

src/ECS/Archetype/Archetype.cs

@@ -308,20 +308,20 @@ internal static void MoveLastComponentsTo(Archetype arch, int newIndex, bool upd
     }
 
     /// <summary>
-    /// <see cref="TreeNode"/> components are not copied.<br/>
+    /// <see cref="TreeNode"/> components (child entities) are not copied.<br/>
     /// Otherwise, two different entities would have the same child entities.
     /// </summary>
-    internal static void CloneComponents(Archetype sourceArch, Archetype targetArch, in CopyContext context)
+    internal static void CopyComponents(Archetype sourceArch, Archetype targetArch, in CopyContext context, long updateIndexTypes)
     {
         var sourceIndex = context.source.compIndex;
         var targetIndex = context.target.compIndex;
         if (sourceArch == targetArch) {
             foreach (var sourceHeap in sourceArch.structHeaps) {
-                sourceHeap.CloneComponent(sourceIndex, sourceHeap, targetIndex, context);
+                sourceHeap.CopyComponent(sourceIndex, sourceHeap, targetIndex, context, updateIndexTypes);
             }
         } else {
             foreach (var sourceHeap in sourceArch.structHeaps) {
-                sourceHeap.CloneComponent(sourceIndex, targetArch.heapMap[sourceHeap.structIndex], targetIndex, context);
+                sourceHeap.CopyComponent(sourceIndex, targetArch.heapMap[sourceHeap.structIndex], targetIndex, context, updateIndexTypes);
             }
         }
     }

src/ECS/Archetype/StructHeap.cs

@@ -36,7 +36,7 @@ internal abstract class StructHeap : IComponentStash
     internal  abstract  void        ResizeComponents        (int capacity, int count);
     internal  abstract  void        MoveComponent           (int from, int to);
     internal  abstract  void        CopyComponentTo         (int sourcePos, StructHeap targetHeap, int targetPos);
-    internal  abstract  void        CloneComponent          (int sourcePos, StructHeap targetHeap, int targetPos, in CopyContext context);
+    internal  abstract  void        CopyComponent           (int sourcePos, StructHeap targetHeap, int targetPos, in CopyContext context, long updateIndexTypes);
     internal  abstract  void        SetComponentDefault     (int compIndex);
     internal  abstract  void        SetComponentsDefault    (int compIndexStart, int count);
     internal  abstract  object      GetComponentDebug       (int compIndex);

src/ECS/Archetype/StructHeap.generic.cs

@@ -73,24 +73,33 @@ internal override void CopyComponentTo(int sourcePos, StructHeap target, int tar
         targetHeap.components[targetPos] = components[sourcePos];
     }
 
-    internal override void CloneComponent(int sourcePos, StructHeap targetHeap, int targetPos, in CopyContext context)
+    internal override void CopyComponent(int sourcePos, StructHeap targetHeap, int targetPos, in CopyContext context, long updateIndexTypes)
     {
         if (typeof(T) == typeof(TreeNode)) {
             return;
         }
-        var copyValue = CopyValueUtils<T>.CopyValue;
-        ref T source = ref components[sourcePos];
-        ref T target = ref ((StructHeap<T>)targetHeap).components[targetPos];
+        var copyValue       = CopyValueUtils<T>.CopyValue;
+        ref T source        = ref components[sourcePos];
+        var typedTargetHeap = (StructHeap<T>)targetHeap;
+        ref T target        = ref typedTargetHeap.components[targetPos];
+        if (StructInfo<T>.HasIndex) {
+            AddOrUpdateIndex(source, target, context.target, typedTargetHeap, updateIndexTypes);
+        }
         if (copyValue == null) {
             target = source;
         } else {
             copyValue(source, ref target, context);
         }
-        if (!StructInfo<T>.HasIndex) {
-            return;
+    }
+
+    private static void AddOrUpdateIndex(in T source, in T target, in Entity targetEntity, StructHeap<T> targetHeap, long updateIndexTypes)
+    {
+        if (((1 << StructInfo<T>.Index) & updateIndexTypes) == 0) {
+            StoreIndex.AddIndex(targetEntity.store, targetEntity.Id, source);
+        } else {
+            targetHeap.componentStash = target;
+            StoreIndex.UpdateIndex(targetEntity.store, targetEntity.Id, source, targetHeap);
         }
-        var targetEntity = context.target;
-        StoreIndex.AddIndex(targetEntity.store, targetEntity.Id, source);
     }
 
     internal  override  void SetComponentDefault (int compIndex) {

src/ECS/Entity/Extensions/EntityExtensions.cs

@@ -15,7 +15,7 @@ namespace Friflo.Engine.ECS;
 public static partial class EntityExtensions
 {
 #region add components
-    private static readonly long IndexTypesMask = EntityStoreBase.Static.EntitySchema.indexTypes.bitSet.l0;
+    internal static readonly long IndexTypesMask = EntityStoreBase.Static.EntitySchema.indexTypes.bitSet.l0;
 
     private static void UpdateIndexedComponents(Entity entity, Archetype archetype, long indexTypesMask)
     {

src/ECS/Entity/Store/Entities.cs

@@ -76,6 +76,15 @@ internal int CreateEntityInternal(Archetype archetype, int id, out short revisio
         return CreateEntityNode(archetype, id, out revision);
     }
 
+    private static void RemoveIndexedComponents(Entity entity, long removedIndexTypes) {
+        var removeTypes = new ComponentTypes();
+        removeTypes.bitSet.l0 = removedIndexTypes;
+        var heapMap = entity.archetype.heapMap;
+        foreach (var removeType in removeTypes) {
+            heapMap[removeType.StructIndex].RemoveIndex(entity);
+        }
+    }
+
     /// <summary>
     /// Copy all components, tags and scripts of the <paramref name="source"/> entity to the <paramref name="target"/> entity.<br/>
     /// The <paramref name="source"/> and <paramref name="target"/> entities can be in the same or different stores.
@@ -89,14 +98,27 @@ public static void CopyEntity(Entity source, Entity target)
         var sourceArch      = source.GetArchetype() ?? throw EntityArgumentNullException(source, nameof(source));
         var curTargetArch   = target.GetArchetype() ?? throw EntityArgumentNullException(target, nameof(target));
         var targetStore     = target.store;
-        var targetArch      = targetStore.GetArchetype(sourceArch.componentTypes, sourceArch.Tags);
+        if (source.store == targetStore) {
+            if (targetStore.internBase.activeQueryLoops > 0) {
+                throw StructuralChangeWithinQueryLoop();
+            }
+        }
+        var targetArch = targetStore.GetArchetype(sourceArch.componentTypes, sourceArch.Tags);
         if (targetArch != curTargetArch) {
+            var removedIndexTypes = (curTargetArch.componentTypes.bitSet.l0 & ~targetArch.componentTypes.bitSet.l0) & EntityExtensions.IndexTypesMask;
+            // --- remove indexes of removed indexed components
+            if (removedIndexTypes != 0) {
+                RemoveIndexedComponents(target, removedIndexTypes);
+            }
+            // --- move entity targetArch 
             ref var node    = ref targetStore.nodes[target.Id];
             node.compIndex  = Archetype.MoveEntityTo(curTargetArch, target.Id, node.compIndex, targetArch);
             node.archetype  = targetArch;
         }
-        var context = new CopyContext(source, target);
-        Archetype.CloneComponents(sourceArch, targetArch, context);
+        // bit == 1: update component index.    bit == 0: add component index
+        var updateIndexTypes    = curTargetArch.componentTypes.bitSet.l0 & targetArch.componentTypes.bitSet.l0;
+        var context             = new CopyContext(source, target);
+        Archetype.CopyComponents(sourceArch, targetArch, context, updateIndexTypes);
 
         targetStore.CloneScrips(source, target);
     }
@@ -119,7 +141,7 @@ public Entity CloneEntity(Entity entity)
         // if (true) {
 
         var context = new CopyContext(entity, clone);
-        Archetype.CloneComponents(archetype, archetype, context);
+        Archetype.CopyComponents(archetype, archetype, context, 0);
 
         CloneScrips(entity, clone);
 

src/Tests/ECS/Arch/Test_StructuralChangeException.cs

@@ -51,6 +51,11 @@ private static void TestExceptions(EntityStore store, Entity entity)
             entity.RemoveComponent<Position>();
         });
 
+        var target = store.CreateEntity();
+        Assert.Throws<StructuralChangeException>(() => {
+            EntityStore.CopyEntity(entity, target);
+        });
+
         var buffer = store.GetCommandBuffer();
         Assert.Throws<StructuralChangeException>(() => {
             buffer.Playback();

src/Tests/ECS/Entity/Test_CopyEntity.cs

@@ -1,6 +1,7 @@
 using System;
 using Friflo.Engine.ECS;
 using NUnit.Framework;
+using Tests.ECS.Index;
 using static NUnit.Framework.Assert;
 
 // ReSharper disable EqualExpressionComparison
@@ -11,6 +12,31 @@ namespace Tests.ECS {
 
 public static class Test_CopyEntity
 {
+    [Test]
+    public static void Test_CopyEntity_CloneStore_subset()
+    {
+        var store       = new EntityStore();
+        var targetStore = new EntityStore();
+
+        store.CreateEntity(new Position(1,1,1));                        // 1
+        store.CreateEntity(new Position(2,2,2), Tags.Get<TestTag>());   // 2
+        store.CreateEntity(new Position(3,3,3));                        // 3
+        store.CreateEntity(new Position(4,4,4), Tags.Get<TestTag>());   // 4
+        store.CreateEntity(new Position(5,5,5));                        // 5
+
+        // Query will copy only entities [2, 4]
+        var query = store.Query().AllTags(Tags.Get<TestTag>());
+        foreach (var entity in query.Entities) {
+            // preserve same entity ids in target store
+            if (!targetStore.TryGetEntityById(entity.Id, out Entity targetEntity)) {
+                targetEntity = targetStore.CreateEntity(entity.Id);
+            }
+            EntityStore.CopyEntity(entity, targetEntity);
+        }
+        AreEqual(new Position(2,2,2), targetStore.GetEntityById(2).GetComponent<Position>());
+        AreEqual(new Position(4,4,4), targetStore.GetEntityById(4).GetComponent<Position>());
+    }
+
     [Test]
     public static void Test_CopyEntity_different_stores()
     {
@@ -68,6 +94,39 @@ public static void Test_CopyEntity_same_stores()
         IsTrue(target6.Tags.Has<TestTag2>());
     }
 
+    [Test]
+    public static void Test_CopyEntity_IndexedComponent()
+    {
+        var store  = new EntityStore();
+        var index  = store.ComponentIndex<IndexedInt,int>();
+
+        var entity1 = store.CreateEntity(new IndexedInt { value = 11 });
+        var entity2 = store.CreateEntity(2);
+        AreEqual("{ 1 }",       index[11].Debug());
+
+        EntityStore.CopyEntity(entity1, entity2);
+        AreEqual("{ 1, 2 }",    index[11].Debug());
+        AreEqual(2, store.Count);
+        AreEqual(new IndexedInt { value = 11 }, entity2.GetComponent<IndexedInt>());
+
+        entity1.AddComponent(new IndexedInt { value = 42 });
+        AreEqual("{ 2 }",       index[11].Debug());
+        AreEqual("{ 1 }",       index[42].Debug());
+
+        EntityStore.CopyEntity(entity1, entity2);
+        AreEqual("{ }",         index[11].Debug());
+        AreEqual("{ 1, 2 }",    index[42].Debug());
+        AreEqual(new IndexedInt { value = 42 }, entity2.GetComponent<IndexedInt>());
+
+        entity1.RemoveComponent<IndexedInt>();
+        AreEqual("{ }",         index[11].Debug());
+        AreEqual("{ 2 }",       index[42].Debug());
+
+        EntityStore.CopyEntity(entity1, entity2);
+        AreEqual("{ }",         index[11].Debug());
+        AreEqual("{ }",         index[42].Debug());
+    }
+
     [Test]
     public static void Test_CopyEntity_exceptions()
     {