Reduce memory and CPU costs due to SegmentedList usage

src/Compilers/Core/CodeAnalysisTest/Collections/List/SegmentedList.Generic.Tests.cs

@@ -73,5 +73,88 @@ public void CopyTo_ArgumentValidity(int count)
             Assert.Throws<ArgumentException>(null, () => list.CopyTo(0, new T[0], 0, count + 1));
             Assert.Throws<ArgumentException>(null, () => list.CopyTo(count, new T[0], 0, 1));
         }
+
+        [Theory]
+        [MemberData(nameof(ValidCollectionSizes))]
+        public void Capacity_ArgumentValidity(int count)
+        {
+            var list = new SegmentedList<T>(count);
+            for (var i = 0; i < count; i++)
+                list.Add(CreateT(i));
+
+            Assert.Throws<ArgumentOutOfRangeException>(() => list.Capacity = count - 1);
+        }
+
+        [Theory]
+        [InlineData(0, 0, 1)]
+        [InlineData(0, 0, 10)]
+        [InlineData(4, 4, 6)]
+        [InlineData(4, 4, 10)]
+        [InlineData(4, 4, 100_000)]
+        public void Capacity_MatchesSizeRequested(int initialCapacity, int initialSize, int requestedCapacity)
+        {
+            var list = new SegmentedList<T>(initialCapacity);
+
+            for (var i = 0; i < initialSize; i++)
+                list.Add(CreateT(i));
+
+            list.Capacity = requestedCapacity;
+
+            Assert.Equal(requestedCapacity, list.Capacity);
+        }
+
+        [Theory]
+        [InlineData(0, 0, 1, 4)]
+        [InlineData(0, 0, 10, 10)]
+        [InlineData(4, 4, 6, 8)]
+        [InlineData(4, 4, 10, 10)]
+        public void EnsureCapacity_ResizesAppropriately(int initialCapacity, int initialSize, int requestedCapacity, int expectedCapacity)
+        {
+            var list = new SegmentedList<T>(initialCapacity);
+
+            for (var i = 0; i < initialSize; i++)
+                list.Add(CreateT(i));
+
+            list.EnsureCapacity(requestedCapacity);
+
+            Assert.Equal(expectedCapacity, list.Capacity);
+        }
+
+        [Theory]
+        [InlineData(1)]
+        [InlineData(2)]
+        [InlineData(4)]
+        public void EnsureCapacity_GrowsBySegment(int segmentCount)
+        {
+            var elementCount = SegmentedArray<T>.TestAccessor.SegmentSize * segmentCount;
+            var list = new SegmentedList<T>(elementCount);
+
+            for (var i = 0; i < elementCount; i++)
+                list.Add(CreateT(i));
+
+            Assert.Equal(elementCount, list.Capacity);
+
+            list.EnsureCapacity(elementCount + 1);
+            Assert.Equal(elementCount + SegmentedArray<T>.TestAccessor.SegmentSize, list.Capacity);
+        }
+
+        [Theory]
+        [InlineData(1)]
+        [InlineData(2)]
+        [InlineData(4)]
+        public void EnsureCapacity_MatchesSizeWithLargeCapacityRequest(int segmentCount)
+        {
+            var elementCount = SegmentedArray<T>.TestAccessor.SegmentSize * segmentCount;
+            var list = new SegmentedList<T>(elementCount);
+
+            for (var i = 0; i < elementCount; i++)
+                list.Add(CreateT(i));
+
+            Assert.Equal(elementCount, list.Capacity);
+
+            var requestedCapacity = 2 * elementCount + 10;
+            list.EnsureCapacity(requestedCapacity);
+            Assert.Equal(requestedCapacity, list.Capacity);
+        }
     }
 }

src/Compilers/Core/CodeAnalysisTest/Collections/SegmentedArrayTests.cs

@@ -29,6 +29,23 @@ public static IEnumerable<object[]> TestLengths
             }
         }
 
+        public static IEnumerable<object[]> TestLengthsAndSegmentCounts
+        {
+            get
+            {
+                for (var segmentsToAdd = 1; segmentsToAdd < 4; segmentsToAdd++)
+                {
+                    yield return new object[] { 1, segmentsToAdd };
+                    yield return new object[] { 10, segmentsToAdd };
+                    yield return new object[] { 100, segmentsToAdd };
+                    yield return new object[] { SegmentedArray<object>.TestAccessor.SegmentSize / 2, segmentsToAdd };
+                    yield return new object[] { SegmentedArray<object>.TestAccessor.SegmentSize, segmentsToAdd };
+                    yield return new object[] { SegmentedArray<object>.TestAccessor.SegmentSize * 2, segmentsToAdd };
+                    yield return new object[] { 100000, segmentsToAdd };
+                }
+            }
+        }
+
         private static void ResetToSequence(SegmentedArray<IntPtr> array)
         {
             for (var i = 0; i < array.Length; i++)
@@ -241,5 +258,74 @@ static void initialize(int[] array, SegmentedArray<int> segmented)
                 }
             }
         }
+
+        [Theory]
+        [MemberData(nameof(TestLengths))]
+        public void Resize_ArgumentValidity(int length)
+        {
+            var o = new object();
+
+            var segmented = new SegmentedArray<object>(length);
+            for (var i = 0; i < length; i++)
+                segmented[i] = o;
+
+            Assert.Throws<ArgumentOutOfRangeException>(() => segmented.Resize(length - 1));
+            Assert.Throws<ArgumentOutOfRangeException>(() => segmented.Resize(length));
+        }
+
+        [Theory]
+        [MemberData(nameof(TestLengthsAndSegmentCounts))]
+        public void Resize_ReusesSegments(
+            int length,
+            int addSegmentCount)
+        {
+            var elementCountToAdd = SegmentedArray<object>.TestAccessor.SegmentSize * addSegmentCount;
+            var o = new object();
+
+            var oldSegmented = new SegmentedArray<object>(length);
+            for (var i = 0; i < length; i++)
+                oldSegmented[i] = o;
+
+            var oldSegments = SegmentedArray<object>.TestAccessor.GetSegments(oldSegmented);
+            var oldSegmentCount = oldSegments.Length;
+
+            var resizedSegmented = oldSegmented.Resize(length + elementCountToAdd);
+            var resizedSegments = SegmentedArray<object>.TestAccessor.GetSegments(resizedSegmented);
+            var resizedSegmentCount = resizedSegments.Length;
+
+            Assert.Equal(oldSegmentCount + addSegmentCount, resizedSegmentCount);
+
+            for (var i = 0; i < oldSegmentCount - 1; i++)
+                Assert.Same(resizedSegments[i], oldSegments[i]);
+
+            for (var i = oldSegmentCount - 1; i < resizedSegmentCount - 1; i++)
+                Assert.Equal(resizedSegments[i].Length, SegmentedArray<object>.TestAccessor.SegmentSize);
+
+            Assert.NotSame(resizedSegments[resizedSegmentCount - 1], oldSegments[oldSegmentCount - 1]);
+            Assert.Equal(resizedSegments[resizedSegmentCount - 1].Length, oldSegments[oldSegmentCount - 1].Length);
+        }
+
+        [Theory]
+        [CombinatorialData]
+        public void Resize_InOnlySingleSegment(
+            [CombinatorialValues(1, 2, 10, 100)] int length,
+            [CombinatorialValues(1, 2, 10, 100)] int addItemCount)
+        {
+            var o = new object();
+
+            var oldSegmented = new SegmentedArray<object>(length);
+            for (var i = 0; i < length; i++)
+                oldSegmented[i] = o;
+
+            var oldSegments = SegmentedArray<object>.TestAccessor.GetSegments(oldSegmented);
+
+            var resizedSegmented = oldSegmented.Resize(length + addItemCount);
+            var resizedSegments = SegmentedArray<object>.TestAccessor.GetSegments(resizedSegmented);
+
+            Assert.Equal(1, oldSegments.Length);
+            Assert.Equal(1, resizedSegments.Length);
+            Assert.NotSame(resizedSegments[0], oldSegments[0]);
+            Assert.Equal(resizedSegments[0].Length, oldSegments[0].Length + addItemCount);
+        }
     }
 }

src/Dependencies/Collections/SegmentedArray`1.cs

@@ -94,6 +94,52 @@ public SegmentedArray(int length)
             }
         }
 
+        /// <summary>
+        /// Creates a new SegmentedArray of the specified size. All pages from this
+        /// SegmentedArray are copied into the new array. Note as this reuses the
+        /// pages, operations on the returned SegmentedArray and this SegmentedArray
+        /// will affect each other.
+        /// </summary>
+        /// <param name="newLength">The desired length of the returned SegmentedArray.
+        /// Note that this is currently limited to be larger than the current SegmentedArray's
+        /// length.
+        /// </param>
+        /// <returns>The new SegmentedArray</returns>
+        /// <exception cref="ArgumentOutOfRangeException">Thrown when the desired length is not
+        /// greater than the current SegmentedArray's length.</exception>
+        public SegmentedArray<T> Resize(int newLength)
+        {
+            // For now, only allow growing resizes
+            if (newLength <= _length)
+                throw new ArgumentOutOfRangeException(nameof(newLength));
+
+            var newItems = new T[(newLength + SegmentSize - 1) >> SegmentShift][];
+            var lastPageSize = newLength - ((newItems.Length - 1) << SegmentShift);
+
+            // Copy over all old pages
+            for (var i = 0; i < _items.Length; i++)
+                newItems[i] = _items[i];
+
+            // If the previous last page is still the last page, resize it to lastPageSize.
+            // Otherwise, resize it to SegmentSize.
+            if (_items.Length > 0)
+            {
+                Array.Resize(
+                    ref newItems[_items.Length - 1],
+                    _items.Length == newItems.Length ? lastPageSize : SegmentSize);
+            }
+
+            // Create all new pages (except the last one which is done separately)
+            for (var i = _items.Length; i < newItems.Length - 1; i++)
+                newItems[i] = new T[SegmentSize];
+
+            // Create a new last page if necessary
+            if (_items.Length < newItems.Length)
+                newItems[newItems.Length - 1] = new T[lastPageSize];
+
+            return new SegmentedArray<T>(newLength, newItems);
+        }
+
         private SegmentedArray(int length, T[][] items)
         {
             _length = length;
@@ -419,6 +465,8 @@ public void Reset()
         internal static class TestAccessor
         {
             public static int SegmentSize => SegmentedArray<T>.SegmentSize;
+            public static T[][] GetSegments(SegmentedArray<T> array)
+                => array._items;
         }
     }
 }

src/Dependencies/Collections/SegmentedList`1.cs

@@ -135,12 +135,7 @@ public int Capacity
                 {
                     if (value > 0)
                     {
-                        var newItems = new SegmentedArray<T>(value);
-                        if (_size > 0)
-                        {
-                            SegmentedArray.Copy(_items, newItems, _size);
-                        }
-                        _items = newItems;
+                        _items = _items.Resize(value);
                     }
                     else
                     {
@@ -502,7 +497,17 @@ internal void Grow(int capacity)
             // If the computed capacity is still less than specified, set to the original argument.
             // Capacities exceeding Array.MaxLength will be surfaced as OutOfMemoryException by Array.Resize.
             if (newCapacity < capacity)
+            {
                 newCapacity = capacity;
+            }
+            else
+            {
+                var segmentSize = SegmentedArrayHelper.GetSegmentSize<T>();
+
+                // If caller didn't request a large capacity increase, limit the increase to a single page
+                if (newCapacity > segmentSize)
+                    newCapacity = (((capacity - 1) / segmentSize) + 1) * segmentSize;
+            }
 
             Capacity = newCapacity;
         }

src/Tools/IdeCoreBenchmarks/SegmentedListBenchmarks_InsertRange.cs

@@ -8,6 +8,7 @@
 
 namespace IdeCoreBenchmarks
 {
+    [MemoryDiagnoser]
     [DisassemblyDiagnoser]
     public class SegmentedListBenchmarks_InsertRange
     {
@@ -23,7 +24,7 @@ public class SegmentedListBenchmarks_InsertRange
         private Microsoft.CodeAnalysis.Collections.SegmentedList<object?> _segmentedValuesObject = null!;
         private SegmentedArray<object?> _segmentedInsertValuesObject;
 
-        [Params(100000)]
+        [Params(1_000, 10_000, 100_000, 1_000_000)]
         public int Count { get; set; }
 
         [GlobalSetup]
@@ -40,6 +41,17 @@ public void GlobalSetup()
             _segmentedInsertValuesObject = new SegmentedArray<object?>(100);
         }
 
+        [Benchmark(Description = "AddToSegmentedList<object>", Baseline = true)]
+        public void AddList()
+        {
+            var array = new Microsoft.CodeAnalysis.Collections.SegmentedList<object?>();
+            var iterations = Count;
+            for (var i = 0; i < iterations; i++)
+            {
+                array.Add(null);
+            }
+        }
+
         [Benchmark(Description = "List<int>", Baseline = true)]
         public void InsertRangeList()
         {