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

* Add ImmutableArray<T>.Builder.DrainToImmutable

* Refactor DrainToImmutable for a more optimized implementation

* Augment and refine DrainToImmutable unit tests

diff --git a/src/libraries/System.Collections.Immutable/ref/System.Collections.Immutable.cs b/src/libraries/System.Collections.Immutable/ref/System.Collections.Immutable.cs
index add568d..f934a9a 100644 (file)
--- a/src/libraries/System.Collections.Immutable/ref/System.Collections.Immutable.cs
+++ b/src/libraries/System.Collections.Immutable/ref/System.Collections.Immutable.cs
@@ -364,6 +364,7 @@ namespace System.Collections.Immutable
             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; }

diff --git a/src/libraries/System.Collections.Immutable/src/System/Collections/Immutable/ImmutableArray_1.Builder.cs b/src/libraries/System.Collections.Immutable/src/System/Collections/Immutable/ImmutableArray_1.Builder.cs
index 205b416..879fb65 100644 (file)
--- a/src/libraries/System.Collections.Immutable/src/System/Collections/Immutable/ImmutableArray_1.Builder.cs
+++ b/src/libraries/System.Collections.Immutable/src/System/Collections/Immutable/ImmutableArray_1.Builder.cs
@@ -214,6 +214,30 @@ namespace System.Collections.Immutable
             }
 
             /// <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>
             public void Clear()

diff --git a/src/libraries/System.Collections.Immutable/tests/ImmutableArrayBuilderTest.cs b/src/libraries/System.Collections.Immutable/tests/ImmutableArrayBuilderTest.cs
index aaddd60..0bb44cb 100644 (file)
--- a/src/libraries/System.Collections.Immutable/tests/ImmutableArrayBuilderTest.cs
+++ b/src/libraries/System.Collections.Immutable/tests/ImmutableArrayBuilderTest.cs
@@ -881,6 +881,137 @@ namespace System.Collections.Immutable.Tests
         }
 
         [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()
         {
             var builder = ImmutableArray.CreateBuilder<int>(initialCapacity: 10);