add more record methods to EntityQueryParallelExtensions

lunditoph-siago · Oct 20, 2023 · 36d9cfd · 36d9cfd
1 parent 0f6e1f5
commit 36d9cfd
Showing 1 changed file with 101 additions and 24 deletions.
diff --git a/Sia/Entities/Extensions/EntityQueryParallelExtensions.cs b/Sia/Entities/Extensions/EntityQueryParallelExtensions.cs
@@ -6,6 +6,9 @@ namespace Sia;
 
 public static class EntityQueryParallelExtensions
 {
+    public delegate void RecordFunc<TResult>(in EntityRef entity, ref TResult result);
+    public delegate void RecordFunc<TResult, TData>(in TData data, in EntityRef entity, ref TResult result);
+
     private readonly struct ForEachParallelAction
     {
         private readonly ArraySegment<EntityRef> _array;
@@ -46,40 +49,46 @@ public void Invoke(System.Tuple<int, int> range)
         }
     }
 
-    private unsafe struct ForEachParallelData
+    private readonly unsafe struct RecordData
     {
-        public int* Index;
-        public ArraySegment<EntityRef> Array;
+        public readonly int* Index;
+        public readonly ArraySegment<EntityRef> Array;
 
-        public ForEachParallelData(int* index, ArraySegment<EntityRef> array)
+        public RecordData(int* index, ArraySegment<EntityRef> array)
         {
             Index = index;
             Array = array;
         }
     }
 
-    [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private unsafe static void DoRecord(IEntityQuery query, ArraySegment<EntityRef> array)
+    private readonly unsafe struct RecordData<T>
     {
-        int index = 0;
-        var indexPtr = (int*)Unsafe.AsPointer(ref index);
-
-        query.ForEach(new(indexPtr, array), static (in ForEachParallelData data, in EntityRef entity) => {
-            ref int index = ref *data.Index;
-            data.Array[index] = entity;
-            ++index;
-        });
+        public readonly int* Index;
+        public readonly ArraySegment<T> Array;
+        public readonly RecordFunc<T> RecordFunc;
+
+        public RecordData(int* index, ArraySegment<T> array, RecordFunc<T> recordFunc)
+        {
+            Index = index;
+            Array = array;
+            RecordFunc = recordFunc;
+        }
     }
 
-    [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public static unsafe SpanOwner<EntityRef> Record(this IEntityQuery query)
+    private readonly unsafe struct RecordData<TResult, TData>
     {
-        var count = query.Count;
-        if (count == 0) { return default; }
-
-        var spanOwner = SpanOwner<EntityRef>.Allocate(count);
-        DoRecord(query, spanOwner.DangerousGetArray());
-        return spanOwner;
+        public readonly int* Index;
+        public readonly TData Data;
+        public readonly ArraySegment<TResult> Array;
+        public readonly RecordFunc<TResult, TData> RecordFunc;
+
+        public RecordData(int* index, in TData data, ArraySegment<TResult> array, RecordFunc<TResult, TData> recordFunc)
+        {
+            Index = index;
+            Data = data;
+            Array = array;
+            RecordFunc = recordFunc;
+        }
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -90,7 +99,7 @@ public static unsafe void ForEachParallel(this IEntityQuery query, EntityHandler
 
         var spanOwner = SpanOwner<EntityRef>.Allocate(count);
         var array = spanOwner.DangerousGetArray();
-        DoRecord(query, array);
+        RecordEntities(query, array);
 
         var action = new ForEachParallelAction(array, handler);
         Partitioner.Create(0, count)
@@ -106,7 +115,7 @@ public static unsafe void ForEachParallel<TData>(this IEntityQuery query, in TDa
 
         var spanOwner = SpanOwner<EntityRef>.Allocate(count);
         var array = spanOwner.DangerousGetArray();
-        DoRecord(query, array);
+        RecordEntities(query, array);
 
         var action = new ForEachParallelAction<TData>(array, data, handler);
         Partitioner.Create(0, count)
@@ -123,4 +132,72 @@ public static unsafe void ForEachParallel(this IEntityQuery query, SimpleEntityH
     public static unsafe void ForEachParallel<TData>(this IEntityQuery query, in TData data, SimpleEntityHandler<TData> handler)
         => ForEachParallel(query, (handler, data),
             static (in (SimpleEntityHandler<TData>, TData) data, in EntityRef entity) => data.Item1(data.Item2, entity));
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static unsafe SpanOwner<EntityRef> Record(this IEntityQuery query)
+    {
+        var count = query.Count;
+        if (count == 0) { return default; }
+
+        var spanOwner = SpanOwner<EntityRef>.Allocate(count);
+        RecordEntities(query, spanOwner.DangerousGetArray());
+        return spanOwner;
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static unsafe SpanOwner<TResult> Record<TResult>(
+        this IEntityQuery query, RecordFunc<TResult> recordFunc)
+    {
+        var count = query.Count;
+        if (count == 0) { return default; }
+
+        var spanOwner = SpanOwner<TResult>.Allocate(count);
+        var array = spanOwner.DangerousGetArray();
+
+        int index = 0;
+        var indexPtr = (int*)Unsafe.AsPointer(ref index);
+
+        query.ForEach(new(indexPtr, array, recordFunc), static (in RecordData<TResult> data, in EntityRef entity) => {
+            ref int index = ref *data.Index;
+            data.RecordFunc(entity, ref data.Array.AsSpan()[index]);
+            ++index;
+        });
+
+        return spanOwner;
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static unsafe SpanOwner<TResult> Record<TResult, TData>(
+        this IEntityQuery query, in TData data, RecordFunc<TResult, TData> recordFunc)
+    {
+        var count = query.Count;
+        if (count == 0) { return default; }
+
+        var spanOwner = SpanOwner<TResult>.Allocate(count);
+        var array = spanOwner.DangerousGetArray();
+
+        int index = 0;
+        var indexPtr = (int*)Unsafe.AsPointer(ref index);
+
+        query.ForEach(new(indexPtr, data, array, recordFunc), static (in RecordData<TResult, TData> data, in EntityRef entity) => {
+            ref int index = ref *data.Index;
+            data.RecordFunc(data.Data, entity, ref data.Array.AsSpan()[index]);
+            ++index;
+        });
+
+        return spanOwner;
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private unsafe static void RecordEntities(IEntityQuery query, ArraySegment<EntityRef> array)
+    {
+        int index = 0;
+        var indexPtr = (int*)Unsafe.AsPointer(ref index);
+
+        query.ForEach(new(indexPtr, array), static (in RecordData data, in EntityRef entity) => {
+            ref int index = ref *data.Index;
+            data.Array[index] = entity;
+            ++index;
+        });
+    }
 }