1 files changed, 104 insertions, 0 deletions
diff --git a/ARMeilleure/Common/ThreadStaticPool.cs b/ARMeilleure/Common/ThreadStaticPool.cs
new file mode 100644
index 00000000..cf3a7bb4
--- /dev/null
+++ b/ARMeilleure/Common/ThreadStaticPool.cs
@@ -0,0 +1,104 @@
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Threading;
+
+namespace ARMeilleure.Common
+{
+    internal class ThreadStaticPool<T> where T : class, new()
+    {
+        private const int PoolSizeIncrement = 200;
+
+        [ThreadStatic]
+        private static ThreadStaticPool<T> _instance;
+        public static ThreadStaticPool<T> Instance
+        {
+            get
+            {
+                if (_instance == null)
+                {
+                    PreparePool(0); // So that we can still use a pool when blindly initializing one.
+                }
+                return _instance;
+            }
+        }
+
+        private static ConcurrentDictionary<int, Stack<ThreadStaticPool<T>>> _pools = new ConcurrentDictionary<int, Stack<ThreadStaticPool<T>>>();
+
+        private static Stack<ThreadStaticPool<T>> GetPools(int groupId)
+        {
+            return _pools.GetOrAdd(groupId, x => new Stack<ThreadStaticPool<T>>());
+        }
+
+        public static void PreparePool(int groupId)
+        {
+            // Prepare the pool for this thread, ideally using an existing one from the specified group.
+            if (_instance == null)
+            {
+                Stack<ThreadStaticPool<T>> pools = GetPools(groupId);
+                lock (pools)
+                {
+                    _instance = (pools.Count != 0) ? pools.Pop() : new ThreadStaticPool<T>(PoolSizeIncrement * 2);
+                }
+            }
+        }
+
+        public static void ReturnPool(int groupId)
+        {
+            // Reset and return the pool for this thread to the specified group.
+            Stack<ThreadStaticPool<T>> pools = GetPools(groupId);
+            lock (pools)
+            {
+                _instance.Clear();
+                pools.Push(_instance);
+                _instance = null;
+            }
+        }
+
+        private T[] _pool;
+        private int _poolUsed = -1;
+        private int _poolSize;
+
+        public ThreadStaticPool(int initialSize)
+        {
+            _pool = new T[initialSize];
+
+            for (int i = 0; i < initialSize; i++)
+            {
+                _pool[i] = new T();
+            }
+
+            _poolSize = initialSize;
+        }
+
+        public T Allocate()
+        {
+            int index = Interlocked.Increment(ref _poolUsed);
+            if (index >= _poolSize)
+            {
+                IncreaseSize();
+            }
+            return _pool[index];
+        }
+
+        private void IncreaseSize()
+        {
+            _poolSize += PoolSizeIncrement;
+
+            T[] newArray = new T[_poolSize];
+            Array.Copy(_pool, 0, newArray, 0, _pool.Length);
+
+            for (int i = _pool.Length; i < _poolSize; i++)
+            {
+                newArray[i] = new T();
+            }
+
+            Interlocked.Exchange(ref _pool, newArray);
+        }
+
+        public void Clear()
+        {
+            _poolUsed = -1;
+        }
+    }
+}