From 9e8f98f9feddc544df8da618e5c0663c81f97ef1 Mon Sep 17 00:00:00 2001
From: Atif Aziz <code@raboof.com>
Date: Mon, 23 Sep 2024 23:37:49 +0200
Subject: [PATCH] Reuse thread-static arrays for handles/marshalers

---
 src/CSnakes.Runtime/Python/Pack.cs | 92 ++++++------------------------
 1 file changed, 19 insertions(+), 73 deletions(-)

diff --git a/src/CSnakes.Runtime/Python/Pack.cs b/src/CSnakes.Runtime/Python/Pack.cs
index d2a8a1e7..ccbdcf0c 100644
--- a/src/CSnakes.Runtime/Python/Pack.cs
+++ b/src/CSnakes.Runtime/Python/Pack.cs
@@ -59,86 +59,34 @@ private struct ArrayOf8<T>
         private T _;
     }
 
-    [DebuggerDisplay($"{{{nameof(DebuggerDisplay)}(),nq}}")]
-    struct RentalState
+    private class StockArray<T>(int length)
     {
-        private bool returned;
+        private T[]? array;
 
-        /// <remarks>
-        /// This method should only be called from a state manager like <see
-        /// cref="RentedArray{T}"/>. It should not be called directly from user
-        /// code and thus is named "dangerous" to attract attention.
-        /// </remarks>
-        public void DangerousReturn() => returned = true;
-
-        public static implicit operator bool(RentalState b) => !b.returned;
-
-        private string DebuggerDisplay() => this ? "rented" : "returned";
-    }
-
-    [DebuggerDisplay($"{{{nameof(DebuggerDisplay)}(),nq}}")]
-    private ref struct RentedArray<T>
-    {
-        private readonly T[] array;
-        private ref RentalState rented;
-        private readonly ArrayPool<T> pool;
-        private readonly Span<T> span;
-
-        public RentedArray(ArrayPool<T> pool, int length, ref RentalState rental)
+        public Span<T> GetArray(int minimumLength)
         {
-            Debug.Assert(rental);
-            this.rented = ref rental;
-            this.pool = pool;
-            this.array = pool.Rent(length);
-            this.span = array.AsSpan(..length);
-        }
-
-        public int Length => Span.Length;
-
-        private Span<T> Span
-        {
-            [MethodImpl(MethodImplOptions.AggressiveInlining)]
-            get
-            {
-                if (!this.rented)
-                {
-                    ThrowObjectDisposedException();
-                }
-
-                return this.span;
-            }
-        }
+            if (minimumLength == 0)
+                return [];
 
-        public void Dispose()
-        {
-            if (!this.rented)
-                return;
+            if (minimumLength > length)
+                return new T[minimumLength];
 
-            this.rented.DangerousReturn();
-            this.pool.Return(this.array);
+            return (this.array ??= new T[length])[..minimumLength];
         }
-
-        public Span<T>.Enumerator GetEnumerator() => Span.GetEnumerator();
-
-        public static implicit operator Span<T>(RentedArray<T> rented) => rented.Span;
-
-        private string DebuggerDisplay() =>
-            this.rented ? $"Length = {Length} (Capacity = {this.array.Length})" : "(returned)";
-
-        [DoesNotReturn]
-        private void ThrowObjectDisposedException() => throw new ObjectDisposedException(nameof(RentedArray<T>));
     }
 
-    private static class ArrayPools
+    static class ThreadStatics
     {
-        private const int MaxLength = 100;
-        private const int MaxPerBucket = 10;
+        [ThreadStatic]
+        private static StockArray<nint>? spilledHandles;
+
+        [ThreadStatic]
+        private static StockArray<PyObjectMarshaller.ManagedToUnmanagedIn>? spilledMarshallers;
 
-        private static readonly ArrayPool<nint> Handles = ArrayPool<nint>.Create(MaxLength, MaxPerBucket);
-        private static readonly ArrayPool<PyObjectMarshaller.ManagedToUnmanagedIn> Marshallers = ArrayPool<PyObjectMarshaller.ManagedToUnmanagedIn>.Create(MaxLength, MaxPerBucket);
+        private const int StockArrayThresholdLength = 100;
 
-        public static RentedArray<nint> RentHandles(int length, ref RentalState rental) => new(Handles, length, ref rental);
-        public static RentedArray<PyObjectMarshaller.ManagedToUnmanagedIn> RentMarshallers(int length, ref RentalState returned) => new(Marshallers, length, ref returned);
+        public static StockArray<nint> SpilledHandles => spilledHandles ??= new StockArray<nint>(StockArrayThresholdLength);
+        public static StockArray<PyObjectMarshaller.ManagedToUnmanagedIn> SpilledMarshallers => spilledMarshallers ??= new StockArray<PyObjectMarshaller.ManagedToUnmanagedIn>(StockArrayThresholdLength);
     }
 
     private static nint CreateListOrTuple<TBuilder>(Span<PyObject> items)
@@ -152,12 +100,10 @@ private static nint CreateListOrTuple<TBuilder>(Span<PyObject> items)
         var spillLength = Math.Max(0, items.Length - stackSpillThreshold);
 
         Span<nint> initialHandles = stackalloc nint[Math.Min(stackSpillThreshold, items.Length)];
-        var spilledHandlesRental = new RentalState();
-        using var spilledHandles = ArrayPools.RentHandles(spillLength, ref spilledHandlesRental);
+        var spilledHandles = ThreadStatics.SpilledHandles.GetArray(spillLength);
 
         var initialMarshallers = new ArrayOf8<PyObjectMarshaller.ManagedToUnmanagedIn>();
-        var spilledMarshallersRental = new RentalState();
-        using var spilledMarshallers = ArrayPools.RentMarshallers(spillLength, ref spilledMarshallersRental);
+        var spilledMarshallers = ThreadStatics.SpilledMarshallers.GetArray(spillLength);
 
         scoped var uninitializedMarshallers = MemoryMarshal.CreateSpan(ref Unsafe.As<ArrayOf8<PyObjectMarshaller.ManagedToUnmanagedIn>, PyObjectMarshaller.ManagedToUnmanagedIn>(ref initialMarshallers), stackSpillThreshold);
         var uninitializedHandles = initialHandles;