Reuse segments during SegmentedDictionary growth (#75943)

* Reuse segments during SegmentedDictionary growth

Mimics changes done in #75661 to enable segment reuse in SegmentedDictionary.

src/Dependencies/Collections/SegmentedDictionary`2.cs

@@ -646,10 +646,10 @@ private void Resize(int newSize)
             Debug.Assert(_entries.Length > 0, "_entries should be non-empty");
             Debug.Assert(newSize >= _entries.Length);
 
-            var entries = new SegmentedArray<Entry>(newSize);
-
             var count = _count;
-            SegmentedArray.Copy(_entries, entries, count);
+
+            // Rather than creating a copy of _entries, instead reuse as much of it's data as possible.
+            var entries = CreateNewSegmentedArrayReusingOldSegments(_entries, newSize);
 
             // Assign member variables after both arrays allocated to guard against corruption from OOM if second fails
             _buckets = new SegmentedArray<int>(newSize);
@@ -667,6 +667,28 @@ private void Resize(int newSize)
             _entries = entries;
         }
 
+        private static SegmentedArray<Entry> CreateNewSegmentedArrayReusingOldSegments(SegmentedArray<Entry> oldArray, int newSize)
+        {
+            var segments = SegmentedCollectionsMarshal.AsSegments(oldArray);
+
+            var oldSegmentCount = segments.Length;
+            var newSegmentCount = (newSize + SegmentedArrayHelper.GetSegmentSize<Entry>() - 1) >> SegmentedArrayHelper.GetSegmentShift<Entry>();
+
+            // Grow the array of segments, if necessary
+            Array.Resize(ref segments, newSegmentCount);
+
+            // Resize all segments to full segment size from the last old segment to the next to last
+            // new segment.
+            for (var i = oldSegmentCount - 1; i < newSegmentCount - 1; i++)
+                Array.Resize(ref segments[i], SegmentedArrayHelper.GetSegmentSize<Entry>());
+
+            // Resize the last segment
+            var lastSegmentSize = newSize - ((newSegmentCount - 1) << SegmentedArrayHelper.GetSegmentShift<Entry>());
+            Array.Resize(ref segments[newSegmentCount - 1], lastSegmentSize);
+
+            return SegmentedCollectionsMarshal.AsSegmentedArray(newSize, segments);
+        }
+
         public bool Remove(TKey key)
         {
             // The overload Remove(TKey key, out TValue value) is a copy of this method with one additional

src/Dependencies/Collections/SegmentedHashSet`1.cs

@@ -899,10 +899,10 @@ private void Resize(int newSize)
             Debug.Assert(_entries.Length > 0, "_entries should be non-empty");
             Debug.Assert(newSize >= _entries.Length);
 
-            var entries = new SegmentedArray<Entry>(newSize);
-
             var count = _count;
-            SegmentedArray.Copy(_entries, entries, count);
+
+            // Rather than creating a copy of _entries, instead reuse as much of it's data as possible.
+            var entries = CreateNewSegmentedArrayReusingOldSegments(_entries, newSize);
 
             // Assign member variables after both arrays allocated to guard against corruption from OOM if second fails
             _buckets = new SegmentedArray<int>(newSize);
@@ -921,6 +921,28 @@ private void Resize(int newSize)
             _entries = entries;
         }
 
+        private static SegmentedArray<Entry> CreateNewSegmentedArrayReusingOldSegments(SegmentedArray<Entry> oldArray, int newSize)
+        {
+            var segments = SegmentedCollectionsMarshal.AsSegments(oldArray);
+
+            var oldSegmentCount = segments.Length;
+            var newSegmentCount = (newSize + SegmentedArrayHelper.GetSegmentSize<Entry>() - 1) >> SegmentedArrayHelper.GetSegmentShift<Entry>();
+
+            // Grow the array of segments, if necessary
+            Array.Resize(ref segments, newSegmentCount);
+
+            // Resize all segments to full segment size from the last old segment to the next to last
+            // new segment.
+            for (var i = oldSegmentCount - 1; i < newSegmentCount - 1; i++)
+                Array.Resize(ref segments[i], SegmentedArrayHelper.GetSegmentSize<Entry>());
+
+            // Resize the last segment
+            var lastSegmentSize = newSize - ((newSegmentCount - 1) << SegmentedArrayHelper.GetSegmentShift<Entry>());
+            Array.Resize(ref segments[newSegmentCount - 1], lastSegmentSize);
+
+            return SegmentedCollectionsMarshal.AsSegmentedArray(newSize, segments);
+        }
+
         /// <summary>
         /// Sets the capacity of a <see cref="SegmentedHashSet{T}"/> object to the actual number of elements it contains,
         /// rounded up to a nearby, implementation-specific value.

src/Dependencies/Collections/SegmentedList`1.cs

@@ -148,26 +148,31 @@ public int Capacity
                 else
                 {
                     // Rather than creating a copy of _items, instead reuse as much of it's data as possible.
-                    var segments = SegmentedCollectionsMarshal.AsSegments(_items);
+                    _items = CreateNewSegmentedArrayReusingOldSegments(_items, value);
+                }
+            }
+        }
 
-                    var oldSegmentCount = segments.Length;
-                    var newSegmentCount = (value + SegmentedArrayHelper.GetSegmentSize<T>() - 1) >> SegmentedArrayHelper.GetSegmentShift<T>();
+        private static SegmentedArray<T> CreateNewSegmentedArrayReusingOldSegments(SegmentedArray<T> oldArray, int newSize)
+        {
+            var segments = SegmentedCollectionsMarshal.AsSegments(oldArray);
 
-                    // Grow the array of segments, if necessary
-                    Array.Resize(ref segments, newSegmentCount);
+            var oldSegmentCount = segments.Length;
+            var newSegmentCount = (newSize + SegmentedArrayHelper.GetSegmentSize<T>() - 1) >> SegmentedArrayHelper.GetSegmentShift<T>();
 
-                    // Resize all segments to full segment size from the last old segment to the next to last
-                    // new segment.
-                    for (var i = oldSegmentCount - 1; i < newSegmentCount - 1; i++)
-                        Array.Resize(ref segments[i], SegmentedArrayHelper.GetSegmentSize<T>());
+            // Grow the array of segments, if necessary
+            Array.Resize(ref segments, newSegmentCount);
 
-                    // Resize the last segment
-                    var lastSegmentSize = value - ((newSegmentCount - 1) << SegmentedArrayHelper.GetSegmentShift<T>());
-                    Array.Resize(ref segments[newSegmentCount - 1], lastSegmentSize);
+            // Resize all segments to full segment size from the last old segment to the next to last
+            // new segment.
+            for (var i = oldSegmentCount - 1; i < newSegmentCount - 1; i++)
+                Array.Resize(ref segments[i], SegmentedArrayHelper.GetSegmentSize<T>());
 
-                    _items = SegmentedCollectionsMarshal.AsSegmentedArray(value, segments);
-                }
-            }
+            // Resize the last segment
+            var lastSegmentSize = newSize - ((newSegmentCount - 1) << SegmentedArrayHelper.GetSegmentShift<T>());
+            Array.Resize(ref segments[newSegmentCount - 1], lastSegmentSize);
+
+            return SegmentedCollectionsMarshal.AsSegmentedArray(newSize, segments);
         }
 
         // Read-only property describing how many elements are in the SegmentedList.

src/Tools/IdeCoreBenchmarks/SegmentedDictionaryBenchmarks_Add.cs

@@ -0,0 +1,94 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using BenchmarkDotNet.Attributes;
+using Microsoft.CodeAnalysis.Collections;
+
+namespace IdeCoreBenchmarks
+{
+    [MemoryDiagnoser]
+    public class SegmentedDictionaryBenchmarks_Add
+    {
+        [Params(1_000, 10_000, 100_000, 1_000_000)]
+        public int Count { get; set; }
+
+        private int[]? _intItems;
+        private object[]? _objectItems;
+        private LargeStruct[]? _largeItems;
+        private EnormousStruct[]? _enormousItems;
+
+        [IterationSetup]
+        public void IterationSetup()
+        {
+            _intItems = new int[Count];
+            _objectItems = new object[Count];
+            _largeItems = new LargeStruct[Count];
+            _enormousItems = new EnormousStruct[Count];
+
+            for (var i = 0; i < Count; i++)
+            {
+                _intItems[i] = i;
+                _objectItems[i] = new object();
+                _largeItems[i] = new LargeStruct() { s1 = new MediumStruct() { i1 = i } };
+                _enormousItems[i] = new EnormousStruct() { s1 = _largeItems[i] };
+            }
+        }
+
+        [Benchmark]
+        public void AddIntToList()
+            => AddToList(_intItems!);
+
+        [Benchmark]
+        public void AddObjectToList()
+            => AddToList(_objectItems!);
+
+        [Benchmark]
+        public void AddLargeStructToList()
+            => AddToList(_largeItems!);
+
+        [Benchmark]
+        public void AddEnormousStructToList()
+            => AddToList(_enormousItems!);
+
+        private void AddToList<T>(T[] items) where T : notnull
+        {
+            var dict = new SegmentedDictionary<T, T>();
+            var iterations = Count;
+
+            for (var i = 0; i < iterations; i++)
+                dict.Add(items[i], items[i]);
+        }
+
+        private struct MediumStruct
+        {
+            public int i1 { get; set; }
+            public int i2 { get; set; }
+            public int i3 { get; set; }
+            public int i4 { get; set; }
+            public int i5 { get; set; }
+        }
+
+        private struct LargeStruct
+        {
+            public MediumStruct s1 { get; set; }
+            public MediumStruct s2 { get; set; }
+            public MediumStruct s3 { get; set; }
+            public MediumStruct s4 { get; set; }
+        }
+
+        private struct EnormousStruct
+        {
+            public LargeStruct s1 { get; set; }
+            public LargeStruct s2 { get; set; }
+            public LargeStruct s3 { get; set; }
+            public LargeStruct s4 { get; set; }
+            public LargeStruct s5 { get; set; }
+            public LargeStruct s6 { get; set; }
+            public LargeStruct s7 { get; set; }
+            public LargeStruct s8 { get; set; }
+            public LargeStruct s9 { get; set; }
+            public LargeStruct s10 { get; set; }
+        }
+    }
+}