From 8d736fde910ef7f2cb8c0bfe7a48ac51cba1bbb3 Mon Sep 17 00:00:00 2001
From: Phlam Sicusa <sicusa@gilatod.art>
Date: Fri, 20 Oct 2023 12:48:22 +0800
Subject: [PATCH] refactor EntityQueryParallelExtensions

---
 .../EntityQueryParallelExtensions.cs          | 37 ++++++++++++-------
 1 file changed, 24 insertions(+), 13 deletions(-)

diff --git a/Sia/Entities/Extensions/EntityQueryParallelExtensions.cs b/Sia/Entities/Extensions/EntityQueryParallelExtensions.cs
index d018eff..124b5d3 100644
--- a/Sia/Entities/Extensions/EntityQueryParallelExtensions.cs
+++ b/Sia/Entities/Extensions/EntityQueryParallelExtensions.cs
@@ -8,10 +8,10 @@ public static class EntityQueryParallelExtensions
 {
     private readonly struct ForEachParallelAction
     {
-        private readonly EntityRef[] _array;
+        private readonly ArraySegment<EntityRef> _array;
         private readonly EntityHandler _handler;
 
-        public ForEachParallelAction(EntityRef[] array, EntityHandler handler)
+        public ForEachParallelAction(ArraySegment<EntityRef> array, EntityHandler handler)
         {
             _array = array;
             _handler = handler;
@@ -27,11 +27,11 @@ public void Invoke(System.Tuple<int, int> range)
 
     private readonly struct ForEachParallelAction<TData>
     {
-        private readonly EntityRef[] _array;
+        private readonly ArraySegment<EntityRef> _array;
         private readonly TData _data;
         private readonly EntityHandler<TData> _handler;
 
-        public ForEachParallelAction(EntityRef[] array, in TData data, EntityHandler<TData> handler)
+        public ForEachParallelAction(ArraySegment<EntityRef> array, in TData data, EntityHandler<TData> handler)
         {
             _array = array;
             _data = data;
@@ -49,9 +49,9 @@ public void Invoke(System.Tuple<int, int> range)
     private unsafe struct ForEachParallelData
     {
         public int* Index;
-        public EntityRef[] Array;
+        public ArraySegment<EntityRef> Array;
         
-        public ForEachParallelData(int* index, EntityRef[] array)
+        public ForEachParallelData(int* index, ArraySegment<EntityRef> array)
         {
             Index = index;
             Array = array;
@@ -59,7 +59,7 @@ public ForEachParallelData(int* index, EntityRef[] array)
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private unsafe static void RecordEntities(IEntityQuery query, EntityRef[] array)
+    private unsafe static void DoRecord(IEntityQuery query, ArraySegment<EntityRef> array)
     {
         int index = 0;
         var indexPtr = (int*)Unsafe.AsPointer(ref index);
@@ -71,15 +71,26 @@ private unsafe static void RecordEntities(IEntityQuery query, EntityRef[] array)
         });
     }
 
+    [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);
+        DoRecord(query, spanOwner.DangerousGetArray());
+        return spanOwner;
+    }
+
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public static unsafe void ForEachParallel(this IEntityQuery query, EntityHandler handler)
     {
         var count = query.Count;
         if (count == 0) { return; }
 
-        using var spanOwner = SpanOwner<EntityRef>.Allocate(count);
-        var array = spanOwner.DangerousGetArray().Array!;
-        RecordEntities(query, array);
+        var spanOwner = SpanOwner<EntityRef>.Allocate(count);
+        var array = spanOwner.DangerousGetArray();
+        DoRecord(query, array);
 
         var action = new ForEachParallelAction(array, handler);
         Partitioner.Create(0, count)
@@ -93,9 +104,9 @@ public static unsafe void ForEachParallel<TData>(this IEntityQuery query, in TDa
         var count = query.Count;
         if (count == 0) { return; }
 
-        using var spanOwner = SpanOwner<EntityRef>.Allocate(count);
-        var array = spanOwner.DangerousGetArray().Array!;
-        RecordEntities(query, array);
+        var spanOwner = SpanOwner<EntityRef>.Allocate(count);
+        var array = spanOwner.DangerousGetArray();
+        DoRecord(query, array);
 
         var action = new ForEachParallelAction<TData>(array, data, handler);
         Partitioner.Create(0, count)