perf: add PooledHashSet and PooledStringBuilder

Src/CSharpier/DocPrinter/DocPrinter.cs

@@ -7,7 +7,8 @@ internal class DocPrinter
     protected readonly Stack<PrintCommand> RemainingCommands = new();
     protected readonly Dictionary<string, PrintMode> GroupModeMap = new();
     protected int CurrentWidth;
-    protected readonly StringBuilder Output = new();
+    protected readonly PooledStringBuilder PooledOutput;
+    protected readonly StringBuilder Output;
     protected bool ShouldRemeasure;
     protected bool NewLineNextStringValue;
     protected bool SkipNextNewLine;
@@ -28,6 +29,8 @@ protected DocPrinter(Doc doc, PrinterOptions printerOptions, string endOfLine)
         this.RemainingCommands.Push(
             new PrintCommand(Indenter.GenerateRoot(), PrintMode.Break, doc)
         );
+        this.PooledOutput = PooledStringBuilder.GetInstance();
+        this.Output = this.PooledOutput.Builder;
     }
 
     public static string Print(Doc document, PrinterOptions printerOptions, string endOfLine)
@@ -52,6 +55,8 @@ public string Print()
             result = result.TrimStart('\n', '\r');
         }
 
+        this.PooledOutput.Free();
+
         return result;
     }
 

Src/CSharpier/DocPrinter/PropagateBreaks.cs

@@ -8,7 +8,7 @@ private class MarkerDoc : Doc { }
 
     public static void RunOn(Doc document)
     {
-        var alreadyVisitedSet = new HashSet<Group>();
+        var alreadyVisitedSet = PooledHashSet<Group>.GetInstance();
         var groupStack = new Stack<Group>();
         var forceFlat = 0;
         var canSkipBreak = false;
@@ -132,5 +132,7 @@ void OnExit(Doc doc)
                 docsStack.Push(hasContents.Contents);
             }
         }
+
+        alreadyVisitedSet.Free();
     }
 }

Src/CSharpier/Utilities/ObjectHashSet.cs

@@ -0,0 +1,47 @@
+using System.Diagnostics;
+
+namespace CSharpier.Utilities;
+
+// From https://github.com/dotnet/roslyn/blob/38f239fb81b72bfd313cd18aeff0b0ed40f34c5c/src/Dependencies/PooledObjects/PooledHashSet.cs#L12
+internal sealed class PooledHashSet<T> : HashSet<T>
+{
+    private readonly ObjectPool<PooledHashSet<T>> _pool;
+
+    private PooledHashSet(ObjectPool<PooledHashSet<T>> pool, IEqualityComparer<T> equalityComparer)
+        : base(equalityComparer)
+    {
+        _pool = pool;
+    }
+
+    public void Free()
+    {
+        if (this.Count <= 100_000)
+        {
+            this.Clear();
+            _pool?.Free(this);
+        }
+    }
+
+    // global pool
+    private static readonly ObjectPool<PooledHashSet<T>> s_poolInstance = CreatePool(
+        EqualityComparer<T>.Default
+    );
+
+    // if someone needs to create a pool;
+    public static ObjectPool<PooledHashSet<T>> CreatePool(IEqualityComparer<T> equalityComparer)
+    {
+        ObjectPool<PooledHashSet<T>>? pool = null;
+        pool = new ObjectPool<PooledHashSet<T>>(
+            () => new PooledHashSet<T>(pool!, equalityComparer),
+            16
+        );
+        return pool;
+    }
+
+    public static PooledHashSet<T> GetInstance()
+    {
+        var instance = s_poolInstance.Allocate();
+        Debug.Assert(instance.Count == 0);
+        return instance;
+    }
+}

Src/CSharpier/Utilities/ObjectPool.cs

@@ -0,0 +1,179 @@
+using System.Diagnostics;
+
+namespace CSharpier.Utilities;
+
+// From https://github.com/dotnet/roslyn/blob/38f239fb81b72bfd313cd18aeff0b0ed40f34c5c/src/Dependencies/PooledObjects/ObjectPool%601.cs#L42
+
+/// <summary>
+/// Generic implementation of object pooling pattern with predefined pool size limit. The main
+/// purpose is that limited number of frequently used objects can be kept in the pool for
+/// further recycling.
+///
+/// Notes:
+/// 1) it is not the goal to keep all returned objects. Pool is not meant for storage. If there
+///    is no space in the pool, extra returned objects will be dropped.
+///
+/// 2) it is implied that if object was obtained from a pool, the caller will return it back in
+///    a relatively short time. Keeping checked out objects for long durations is ok, but
+///    reduces usefulness of pooling. Just new up your own.
+///
+/// Not returning objects to the pool in not detrimental to the pool's work, but is a bad practice.
+/// Rationale:
+///    If there is no intent for reusing the object, do not use pool - just use "new".
+/// </summary>
+internal class ObjectPool<T>
+    where T : class
+{
+    [DebuggerDisplay("{Value,nq}")]
+    private struct Element
+    {
+        internal T? Value;
+    }
+
+    /// <remarks>
+    /// Not using System.Func{T} because this file is linked into the (debugger) Formatter,
+    /// which does not have that type (since it compiles against .NET 2.0).
+    /// </remarks>
+    internal delegate T Factory();
+
+    // Storage for the pool objects. The first item is stored in a dedicated field because we
+    // expect to be able to satisfy most requests from it.
+    private T? _firstItem;
+    private readonly Element[] _items;
+
+    // factory is stored for the lifetime of the pool. We will call this only when pool needs to
+    // expand. compared to "new T()", Func gives more flexibility to implementers and faster
+    // than "new T()".
+    private readonly Factory _factory;
+
+    public readonly bool TrimOnFree;
+
+    internal ObjectPool(Factory factory, bool trimOnFree = true)
+        : this(factory, Environment.ProcessorCount * 2, trimOnFree) { }
+
+    internal ObjectPool(Factory factory, int size, bool trimOnFree = true)
+    {
+        Debug.Assert(size >= 1);
+        _factory = factory;
+        _items = new Element[size - 1];
+        TrimOnFree = trimOnFree;
+    }
+
+    internal ObjectPool(Func<ObjectPool<T>, T> factory, int size)
+    {
+        Debug.Assert(size >= 1);
+        _factory = () => factory(this);
+        _items = new Element[size - 1];
+    }
+
+    private T CreateInstance()
+    {
+        var inst = _factory();
+        return inst;
+    }
+
+    /// <summary>
+    /// Produces an instance.
+    /// </summary>
+    /// <remarks>
+    /// Search strategy is a simple linear probing which is chosen for it cache-friendliness.
+    /// Note that Free will try to store recycled objects close to the start thus statistically
+    /// reducing how far we will typically search.
+    /// </remarks>
+    internal T Allocate()
+    {
+        // PERF: Examine the first element. If that fails, AllocateSlow will look at the remaining elements.
+        // Note that the initial read is optimistically not synchronized. That is intentional.
+        // We will interlock only when we have a candidate. in a worst case we may miss some
+        // recently returned objects. Not a big deal.
+        var inst = _firstItem;
+        if (inst == null || inst != Interlocked.CompareExchange(ref _firstItem, null, inst))
+        {
+            inst = AllocateSlow();
+        }
+
+        return inst;
+    }
+
+    private T AllocateSlow()
+    {
+        var items = _items;
+
+        for (var i = 0; i < items.Length; i++)
+        {
+            // Note that the initial read is optimistically not synchronized. That is intentional.
+            // We will interlock only when we have a candidate. in a worst case we may miss some
+            // recently returned objects. Not a big deal.
+            var inst = items[i].Value;
+            if (inst != null)
+            {
+                if (inst == Interlocked.CompareExchange(ref items[i].Value, null, inst))
+                {
+                    return inst;
+                }
+            }
+        }
+
+        return CreateInstance();
+    }
+
+    /// <summary>
+    /// Returns objects to the pool.
+    /// </summary>
+    /// <remarks>
+    /// Search strategy is a simple linear probing which is chosen for it cache-friendliness.
+    /// Note that Free will try to store recycled objects close to the start thus statistically
+    /// reducing how far we will typically search in Allocate.
+    /// </remarks>
+    internal void Free(T obj)
+    {
+        Validate(obj);
+
+        if (_firstItem == null)
+        {
+            // Intentionally not using interlocked here.
+            // In a worst case scenario two objects may be stored into same slot.
+            // It is very unlikely to happen and will only mean that one of the objects will get collected.
+            _firstItem = obj;
+        }
+        else
+        {
+            FreeSlow(obj);
+        }
+    }
+
+    private void FreeSlow(T obj)
+    {
+        var items = _items;
+        for (var i = 0; i < items.Length; i++)
+        {
+            if (items[i].Value == null)
+            {
+                // Intentionally not using interlocked here.
+                // In a worst case scenario two objects may be stored into same slot.
+                // It is very unlikely to happen and will only mean that one of the objects will get collected.
+                items[i].Value = obj;
+                break;
+            }
+        }
+    }
+
+    private void Validate(object obj)
+    {
+        Debug.Assert(obj != null, "freeing null?");
+
+        Debug.Assert(_firstItem != obj, "freeing twice?");
+
+        var items = _items;
+        for (var i = 0; i < items.Length; i++)
+        {
+            var value = items[i].Value;
+            if (value == null)
+            {
+                return;
+            }
+
+            Debug.Assert(value != obj, "freeing twice?");
+        }
+    }
+}

Src/CSharpier/Utilities/PooledStringBuilder.cs

@@ -0,0 +1,86 @@
+using System.Diagnostics;
+using System.Text;
+
+namespace CSharpier.Utilities;
+
+//  From https://github.com/dotnet/roslyn/blob/38f239fb81b72bfd313cd18aeff0b0ed40f34c5c/src/Dependencies/PooledObjects/PooledStringBuilder.cs#L18
+
+/// <summary>
+/// The usage is:
+///        var inst = PooledStringBuilder.GetInstance();
+///        var sb = inst.builder;
+///        ... Do Stuff...
+///        ... sb.ToString() ...
+///        inst.Free();
+/// </summary>
+internal sealed class PooledStringBuilder
+{
+    public readonly StringBuilder Builder = new();
+    private readonly ObjectPool<PooledStringBuilder> _pool;
+
+    private PooledStringBuilder(ObjectPool<PooledStringBuilder> pool)
+    {
+        Debug.Assert(pool != null);
+        _pool = pool!;
+    }
+
+    public int Length
+    {
+        get { return this.Builder.Length; }
+    }
+
+    public void Free()
+    {
+        var builder = this.Builder;
+
+        // do not store builders that are too large.
+        if (builder.Capacity <= 2_000_000)
+        {
+            builder.Clear();
+            _pool.Free(this);
+        }
+    }
+
+    public string ToStringAndFree()
+    {
+        var result = this.Builder.ToString();
+        this.Free();
+
+        return result;
+    }
+
+    public string ToStringAndFree(int startIndex, int length)
+    {
+        var result = this.Builder.ToString(startIndex, length);
+        this.Free();
+
+        return result;
+    }
+
+    // global pool
+    private static readonly ObjectPool<PooledStringBuilder> s_poolInstance = CreatePool();
+
+    // if someone needs to create a private pool;
+    /// <summary>
+    /// If someone need to create a private pool
+    /// </summary>
+    /// <param name="size">The size of the pool.</param>
+    public static ObjectPool<PooledStringBuilder> CreatePool(int size = 16)
+    {
+        ObjectPool<PooledStringBuilder>? pool = null;
+        pool = new ObjectPool<PooledStringBuilder>(() => new PooledStringBuilder(pool!), size);
+        return pool;
+    }
+
+    public static PooledStringBuilder GetInstance()
+    {
+        var builder = s_poolInstance.Allocate();
+        Debug.Assert(builder.Builder.Length == 0);
+        return builder;
+    }
+
+    public static implicit operator StringBuilder(PooledStringBuilder obj)
+    {
+        return obj.Builder;
+    }
+}