feat: pool

Chickensoft.Collections.Tests/src/PoolTest.cs

@@ -0,0 +1,75 @@
+namespace Chickensoft.Collections.Tests;
+
+using System;
+using Shouldly;
+using Xunit;
+
+public class PoolTest {
+  public abstract class Shape : IPooled {
+    public void Reset() { }
+  }
+
+  public class Cube : Shape { }
+  public class Sphere : Shape { }
+
+  [Fact]
+  public void InitializesAndRegisters() {
+    var pool = new Pool<Shape>();
+
+    pool.Register<Cube>();
+    pool.Register<Sphere>();
+
+    pool.Get<Cube>().ShouldBeOfType<Cube>();
+    pool.Get<Sphere>().ShouldBeOfType<Sphere>();
+  }
+
+  [Fact]
+  public void ThrowsIfRegisteringDuplicateType() {
+    var pool = new Pool<Shape>();
+
+    pool.Register<Cube>();
+
+    Should.Throw<InvalidOperationException>(() => pool.Register<Cube>(1));
+  }
+
+  [Fact]
+  public void ThrowsIfGettingUnregisteredType() {
+    var pool = new Pool<Shape>();
+
+    Should.Throw<InvalidOperationException>(pool.Get<Cube>);
+  }
+
+  [Fact]
+  public void CreatesNewInstancesIfPoolIsEmpty() {
+    var pool = new Pool<Shape>();
+
+    pool.Register<Cube>();
+
+    var cube1 = pool.Get<Cube>();
+    var cube2 = pool.Get<Cube>();
+
+    cube1.ShouldNotBeSameAs(cube2);
+  }
+
+  [Fact]
+  public void ReturnThrowsIfTypeNotRegistered() {
+    var pool = new Pool<Shape>();
+
+    Should.Throw<InvalidOperationException>(() => pool.Return(new Cube()));
+  }
+
+  [Fact]
+  public void ReturnResetsAndEnqueues() {
+    var pool = new Pool<Shape>();
+
+    pool.Register<Cube>();
+
+    var cube = pool.Get<Cube>();
+
+    pool.Return(cube);
+
+    var cube2 = pool.Get<Cube>();
+
+    cube.ShouldBeSameAs(cube2);
+  }
+}

Chickensoft.Collections/src/pool/IPooled.cs

@@ -0,0 +1,12 @@
+namespace Chickensoft.Collections;
+
+/// <summary>
+/// Represents a type that can be stored in a pool. Pools maintain
+/// pre-instantiated instances that can be borrowed and returned to avoid
+/// excess memory allocations. Pooled objects are "reset" whenever they are
+/// returned to the pool so that they can be reused.
+/// </summary>
+public interface IPooled {
+  /// <summary>Resets the object to its default state.</summary>
+  void Reset();
+}

Chickensoft.Collections/src/pool/Pool.cs

@@ -0,0 +1,91 @@
+namespace Chickensoft.Collections.Tests;
+
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+
+/// <summary>
+/// Represents a pool of objects that can be borrowed and returned. Objects
+/// are pre-instantiated and stored in the pool. When an object is borrowed,
+/// it is removed from the pool. When an object is returned, it is reset and
+/// placed back into the pool.
+/// </summary>
+public class Pool<TBaseType> where TBaseType : IPooled {
+  private readonly ConcurrentDictionary<Type, ConcurrentQueue<TBaseType>>
+    _pool = [];
+  private readonly Dictionary<Type, Func<TBaseType>> _factories = [];
+
+  /// <summary>Registers a type with the pool.</summary>
+  /// <param name="capacity">The number of items to instantiate for each type
+  /// registered in the pool.</param>
+  /// <typeparam name="TDerivedType">The type to register.</typeparam>
+  public void Register<TDerivedType>(int capacity = 1)
+    where TDerivedType : TBaseType, new() => Register(() => new TDerivedType());
+
+  /// <summary>Registers a type with the pool.</summary>
+  /// <param name="factory">A factory function that creates an instance of the
+  /// type to register.</param>
+  /// <param name="capacity">The number of items to instantiate for each type
+  /// registered in the pool.</param>
+  /// <typeparam name="TDerivedType">The type to register.</typeparam>
+  /// <exception cref="InvalidOperationException" />
+  public void Register<TDerivedType>(
+    Func<TDerivedType> factory, int capacity = 1
+  )
+  where TDerivedType : TBaseType {
+    var type = typeof(TDerivedType);
+    var queue = new ConcurrentQueue<TBaseType>();
+
+    for (var i = 0; i < capacity; i++) {
+      var item = factory();
+      queue.Enqueue(item);
+    }
+
+    lock (_pool) {
+      if (!_pool.TryAdd(type, queue)) {
+        throw new InvalidOperationException(
+          $"Type `{type}` is already registered."
+        );
+      }
+
+      _factories.TryAdd(type, () => factory());
+    }
+  }
+
+  /// <summary>Borrows an object from the pool.</summary>
+  /// <typeparam name="TDerivedType">The type of object to borrow.</typeparam>
+  /// <returns>An object of the specified type.</returns>
+  public TDerivedType Get<TDerivedType>() where TDerivedType : TBaseType, new()
+    => (TDerivedType)Get(typeof(TDerivedType));
+
+  /// <summary>Borrows an object from the pool.</summary>
+  /// <param name="type">The type of object to borrow.</param>
+  /// <returns>An object of the specified type.</returns>
+  public TBaseType Get(Type type) {
+    if (!_pool.TryGetValue(type, out var queue)) {
+      throw new InvalidOperationException($"Type `{type}` is not registered.");
+    }
+
+    if (queue.TryDequeue(out var item)) {
+      return item;
+    }
+
+    // Out of values. Just make one.
+    return _factories[type]();
+  }
+
+  /// <summary>Returns an object to the pool.</summary>
+  /// <param name="item">The object to return.</param>
+  /// <exception cref="InvalidOperationException" />
+  public void Return(TBaseType item) {
+    var type = item.GetType();
+
+    if (!_pool.TryGetValue(type, out var queue)) {
+      throw new InvalidOperationException($"Type `{type}` is not registered.");
+    }
+
+    item.Reset();
+
+    queue.Enqueue(item);
+  }
+}

README.md

@@ -196,6 +196,38 @@ Once you have implemented the `IBoxlessValueHandler`, you can create a boxless q
     queue.Dequeue();
 ```
 
+## Pool
+
+A simple object pool implementation is provided that allows you to pre-allocate objects when memory churn is a concern. Internally, the pool is just a simple wrapper around a .NET concurrent dictionary that maps types to concurrent queues of objects. By leveraging .NET concurrent collections, we can create a type-safe and thread-safe object pool that's easy to use.
+
+Any object you wish to store in a pool must conform to `IPooled` and implement the required `Reset` method. The reset method is called when the object is returned to the pool, allowing you to reset the object's state.
+
+```csharp
+  public abstract class Shape : IPooled {
+    public void Reset() { }
+  }
+
+  public class Cube : Entity { }
+  public class Sphere : Entity { }
+```
+
+A pool can be easily created. Each derived type that you wish to pool can be "registered" with the pool. The pool will create instances of each type registered with it according to the provided capacity.
+
+```csharp
+  var pool = new Pool();
+
+  pool.Register<Cube>(10); // Preallocate 10 cubes.
+  pool.Register<Sphere>(5); // Preallocate 5 spheres.
+
+  // Borrow a cube and a sphere, removing them from the pool:
+  var cube = pool.Get<Cube>();
+  var sphere = pool.Get<Sphere>();
+
+  // Return them to the pool (their Reset() methods will be called):
+  pool.Return(cube);
+  pool.Return(sphere);
+```
+
 ---
 
 🐣 Created with love by Chickensoft 🐤 — <https://chickensoft.games>