Ensure PooledArrayBuilder<T> resets capacity of builder after drain

It is perfectly legal for a PooledArrayBuilder<T> to be drained via one of the DrainToImmutable* methods and then continue adding items to the builder. However, if the inner ImmutableArray<T>.Builder's capacity was set to 0 during the drain, its capacity will not be reset to any specified value when adding new items. This change fixes that.

src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared.Test/PooledObjects/PooledArrayBuilderTests.cs

@@ -86,6 +86,61 @@ public class PooledArrayBuilderTests
         }
     }
 
+    [Fact]
+    public void UseDrainAndReuse()
+    {
+        var builderPool = TestArrayBuilderPool<int>.Create();
+        using var builder = new PooledArrayBuilder<int>(capacity: 10, builderPool);
+
+        for (var i = 0; i < 10; i++)
+        {
+            builder.Add(i);
+        }
+
+        // Verify that the builder contains 10 items within its inner array builder.
+        builder.Validate(static t =>
+        {
+            Assert.Equal(0, t.InlineItemCount);
+            Assert.Equal(10, t.Capacity);
+            Assert.NotNull(t.InnerArrayBuilder);
+
+            Assert.Equal(10, t.InnerArrayBuilder.Count);
+            Assert.Equal(10, t.InnerArrayBuilder.Capacity);
+        });
+
+        var result = builder.DrainToImmutable();
+
+        // After draining, the builder should contain 0 items in its inner array builder
+        // and the capacity should have been set to 0.
+        builder.Validate(static t =>
+        {
+            Assert.Equal(0, t.InlineItemCount);
+            Assert.Equal(10, t.Capacity);
+            Assert.NotNull(t.InnerArrayBuilder);
+
+            Assert.Empty(t.InnerArrayBuilder);
+            Assert.Equal(0, t.InnerArrayBuilder.Capacity);
+        });
+
+        // Add another 10 items to the builder. At the end, the inner array builder
+        // should hold 10 items with a capacity of 10.
+        for (var i = 0; i < 10; i++)
+        {
+            builder.Add(i);
+        }
+
+        // Verify that the builder contains 10 items within its inner array builder.
+        builder.Validate(static t =>
+        {
+            Assert.Equal(0, t.InlineItemCount);
+            Assert.Equal(10, t.Capacity);
+            Assert.NotNull(t.InnerArrayBuilder);
+
+            Assert.Equal(10, t.InnerArrayBuilder.Count);
+            Assert.Equal(10, t.InnerArrayBuilder.Capacity);
+        });
+    }
+
     private static Func<int, bool> IsEven => x => x % 2 == 0;
     private static Func<int, bool> IsOdd => x => x % 2 != 0;
 

src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared/PooledObjects/PooledArrayBuilder`1.cs

@@ -83,13 +83,50 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
         // Return _builder to the pool if necessary. Note that we don't need to clear the inline elements here
         // because this type is intended to be allocated on the stack and the GC can reclaim objects from the
         // stack after the last use of a reference to them.
-        if (_builder is { } builder)
+        if (_builder is { } innerBuilder)
         {
-            _builderPool?.Return(builder);
+            _builderPool?.Return(innerBuilder);
             _builder = null;
         }
     }
 
+    /// <summary>
+    ///  Retrieves the inner <see cref="_builder"/>.
+    /// </summary>
+    /// <returns>
+    ///  Returns <see langword="true"/> if <see cref="_builder"/> is available; otherwise <see langword="false"/>
+    /// </returns>
+    /// <remarks>
+    ///  This should only be used by methods that will not add to the inner <see cref="_builder"/>.
+    /// </remarks>
+    private readonly bool TryGetBuilder([NotNullWhen(true)] out ImmutableArray<T>.Builder? builder)
+    {
+        builder = _builder;
+        return builder is not null;
+    }
+
+    /// <summary>
+    ///  Retrieves the inner <see cref="_builder"/> and resets its capacity if necessary.
+    /// </summary>
+    /// <returns>
+    ///  Returns <see langword="true"/> if <see cref="_builder"/> is available; otherwise <see langword="false"/>
+    /// </returns>
+    /// <remarks>
+    ///  This should only be used by methods that will add to the inner <see cref="_builder"/>.
+    /// </remarks>
+    private readonly bool TryGetBuilderAndEnsureCapacity([NotNullWhen(true)] out ImmutableArray<T>.Builder? builder)
+    {
+        if (TryGetBuilder(out builder))
+        {
+            if (builder.Capacity == 0 && _capacity is int capacity)
+            {
+                builder.Capacity = capacity;
+            }
+        }
+
+        return builder is not null;
+    }
+
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     private void ClearInlineElement(int index)
     {
@@ -112,7 +149,7 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         readonly get
         {
-            if (_builder is { } builder)
+            if (TryGetBuilder(out var builder))
             {
                 return builder[index];
             }
@@ -128,7 +165,7 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         set
         {
-            if (_builder is { } builder)
+            if (TryGetBuilder(out var builder))
             {
                 builder[index] = value;
                 return;
@@ -198,7 +235,7 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
 
     public void Add(T item)
     {
-        if (_builder is { } builder)
+        if (TryGetBuilderAndEnsureCapacity(out var builder))
         {
             builder.Add(item);
         }
@@ -233,7 +270,7 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
             return;
         }
 
-        if (_builder is { } builder)
+        if (TryGetBuilderAndEnsureCapacity(out var builder))
         {
             builder.AddRange(items);
         }
@@ -254,7 +291,7 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
 
     public void AddRange(IEnumerable<T> items)
     {
-        if (_builder is { } builder)
+        if (TryGetBuilderAndEnsureCapacity(out var builder))
         {
             builder.AddRange(items);
             return;
@@ -296,7 +333,7 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
 
     public void Clear()
     {
-        if (_builder is { } builder)
+        if (TryGetBuilder(out var builder))
         {
             // Keep using a real builder to avoid churn in the object pool.
             builder.Clear();
@@ -314,7 +351,7 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
 
     public void RemoveAt(int index)
     {
-        if (_builder is { } builder)
+        if (TryGetBuilderAndEnsureCapacity(out var builder))
         {
             builder.RemoveAt(index);
             return;
@@ -384,7 +421,7 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
     /// <returns>An immutable array.</returns>
     public ImmutableArray<T> DrainToImmutable()
     {
-        if (_builder is { } builder)
+        if (TryGetBuilder(out var builder))
         {
             return builder.DrainToImmutable();
         }
@@ -400,7 +437,7 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
 
     public readonly ImmutableArray<T> ToImmutable()
     {
-        if (_builder is { } builder)
+        if (TryGetBuilder(out var builder))
         {
             return builder.ToImmutable();
         }
@@ -424,7 +461,7 @@ internal partial struct PooledArrayBuilder<T> : IDisposable
 
     public readonly T[] ToArray()
     {
-        if (_builder is { } builder)
+        if (TryGetBuilder(out var builder))
         {
             return builder.ToArray();
         }