Add ImmutableArray<T>.Builder.DrainToImmutable() (dotnet#79385)

* Add ImmutableArray<T>.Builder.DrainToImmutable

* Refactor DrainToImmutable for a more optimized implementation

* Augment and refine DrainToImmutable unit tests

src/libraries/System.Collections.Immutable/ref/System.Collections.Immutable.cs

@@ -364,6 +364,7 @@ public void CopyTo(T[] array, int index) { }
             public void CopyTo(T[] destination) { throw null; }
             public void CopyTo(int sourceIndex, T[] destination, int destinationIndex, int length) { throw null; }
             public void CopyTo(System.Span<T> destination) { }
+            public System.Collections.Immutable.ImmutableArray<T> DrainToImmutable() { throw null; }
             public System.Collections.Generic.IEnumerator<T> GetEnumerator() { throw null; }
             public int IndexOf(T item) { throw null; }
             public int IndexOf(T item, int startIndex) { throw null; }

src/libraries/System.Collections.Immutable/src/System/Collections/Immutable/ImmutableArray_1.Builder.cs

@@ -213,6 +213,30 @@ public ImmutableArray<T> MoveToImmutable()
                 return new ImmutableArray<T>(temp);
             }
 
+            /// <summary>
+            /// Returns the current contents as an <see cref="ImmutableArray{T}"/> and sets the collection to a zero length array.
+            /// </summary>
+            /// <remarks>
+            /// If <see cref="Capacity"/> equals <see cref="Count"/>, the internal array will be extracted
+            /// as an <see cref="ImmutableArray{T}"/> without copying the contents. Otherwise, the contents
+            /// will be copied into a new array. The collection will then be set to a zero length array.
+            /// </remarks>
+            /// <returns>An immutable array.</returns>
+            public ImmutableArray<T> DrainToImmutable()
+            {
+                T[] result = _elements;
+
+                if (result.Length != _count)
+                {
+                    result = ToArray();
+                }
+
+                _elements = ImmutableArray<T>.Empty.array!;
+                _count = 0;
+
+                return new ImmutableArray<T>(result);
+            }
+
             /// <summary>
             /// Removes all items from the <see cref="ICollection{T}"/>.
             /// </summary>

src/libraries/System.Collections.Immutable/tests/ImmutableArrayBuilderTest.cs

@@ -880,6 +880,137 @@ public void MoveToImmutableAddRemoveAddToCapacity()
             Assert.Equal(new[] { 0, 1, 2 }, builder.MoveToImmutable());
         }
 
+        [Fact]
+        public void DrainToImmutableEmptyZeroCapacity()
+        {
+            var builder = ImmutableArray.CreateBuilder<int>(0);
+            var array = builder.DrainToImmutable();
+            Assert.Equal(0, array.Length);
+            Assert.Equal(0, builder.Capacity);
+            Assert.Equal(0, builder.Count);
+        }
+
+        [Fact]
+        public void DrainToImmutableEmptyNonZeroCapacity()
+        {
+            var builder = ImmutableArray.CreateBuilder<int>(10);
+            var array = builder.DrainToImmutable();
+            Assert.Equal(0, array.Length);
+            Assert.Equal(0, builder.Capacity);
+            Assert.Equal(0, builder.Count);
+        }
+
+        [Fact]
+        public void DrainToImmutableAtCapacity()
+        {
+            var builder = ImmutableArray.CreateBuilder<string>(2);
+            builder.Count = 2;
+            builder[1] = "b";
+            builder[0] = "a";
+            Assert.Equal(2, builder.Count);
+            Assert.Equal(2, builder.Capacity);
+
+            var array = builder.DrainToImmutable();
+            Assert.Equal(new[] { "a", "b" }, array);
+            Assert.Equal(0, builder.Count);
+            Assert.Equal(0, builder.Capacity);
+        }
+
+        [Fact]
+        public void DrainToImmutablePartialFill()
+        {
+            var builder = ImmutableArray.CreateBuilder<int>(4);
+            builder.AddRange(42, 13);
+            Assert.Equal(4, builder.Capacity);
+            Assert.Equal(2, builder.Count);
+
+            var array = builder.DrainToImmutable();
+            Assert.Equal(new[] { 42, 13 }, array);
+            Assert.Equal(0, builder.Capacity);
+            Assert.Equal(0, builder.Count);
+        }
+
+        [Fact]
+        public void DrainToImmutablePartialFillWithCountUpdate()
+        {
+            var builder = ImmutableArray.CreateBuilder<int>(4);
+            builder.AddRange(42, 13);
+            builder.Count = builder.Capacity;
+            Assert.Equal(4, builder.Capacity);
+            Assert.Equal(4, builder.Count);
+
+            var array = builder.DrainToImmutable();
+            Assert.Equal(new[] { 42, 13, 0, 0 }, array);
+            Assert.Equal(0, builder.Capacity);
+            Assert.Equal(0, builder.Count);
+        }
+
+        [Fact]
+        public void DrainToImmutableRepeat()
+        {
+            var builder = ImmutableArray.CreateBuilder<string>(2);
+            builder.AddRange("a", "b");
+
+            var array1 = builder.DrainToImmutable();
+            var array2 = builder.DrainToImmutable();
+            Assert.Equal(new[] { "a", "b" }, array1);
+            Assert.Equal(0, array2.Length);
+            Assert.Equal(0, builder.Capacity);
+            Assert.Equal(0, builder.Count);
+        }
+
+        [Fact]
+        public void DrainToImmutableThenUse()
+        {
+            var builder = ImmutableArray.CreateBuilder<string>(2);
+            builder.Count = 2;
+            Assert.Equal(2, builder.DrainToImmutable().Length);
+            Assert.Equal(0, builder.Capacity);
+            builder.AddRange("a", "b");
+            Assert.Equal(2, builder.Count);
+            Assert.True(builder.Capacity >= 2);
+
+            var array = builder.DrainToImmutable();
+            Assert.Equal(new[] { "a", "b" }, array);
+        }
+
+        [Fact]
+        public void DrainToImmutableAfterClear()
+        {
+            var builder = ImmutableArray.CreateBuilder<string>(2);
+            builder.AddRange("a", "b");
+            builder.Clear();
+
+            var array = builder.DrainToImmutable();
+            Assert.Equal(0, array.Length);
+            Assert.Equal(0, builder.Capacity);
+            Assert.Equal(0, builder.Count);
+        }
+
+        [Fact]
+        public void DrainToImmutableAtCapacityClearsCollection()
+        {
+            var builder = ImmutableArray.CreateBuilder<int>(2);
+            builder.AddRange(1, 2);
+            builder.DrainToImmutable();
+            builder.Count = 4;
+
+            var array = builder.DrainToImmutable();
+            Assert.Equal(new[] { 0, 0, 0, 0 }, array);
+        }
+
+        [Fact]
+        public void DrainToImmutablePartialFillClearsCollection()
+        {
+            var builder = ImmutableArray.CreateBuilder<int>(4);
+            builder.AddRange(1, 2);
+            builder.DrainToImmutable();
+            builder.Count = 6;
+
+            var array = builder.DrainToImmutable();
+            Assert.Equal(new[] { 0, 0, 0, 0, 0, 0 }, array);
+        }
+
         [Fact]
         public void CapacitySetToZero()
         {