Use Span to fill List<T> in more ToList scenarios (#86796)

* Use Span to fill List<T> in more ToList scenarios * Optimize Append/Prepend changes * Update src/libraries/System.Linq/src/System/Linq/Lookup.SpeedOpt.cs Co-authored-by: Stephen Toub <[email protected]> * Update src/libraries/System.Linq/src/System/Linq/Lookup.SpeedOpt.cs Co-authored-by: Stephen Toub <[email protected]> * Address feedback * Remove ToSingleItemList, seems like an overoptimization --------- Co-authored-by: Stephen Toub <[email protected]>
dotnet · Oct 30, 2023 · 67a86a2 · 67a86a2
1 parent 4b5756d
commit 67a86a2
Show file tree

Hide file tree

Showing 7 changed files with 114 additions and 58 deletions.
diff --git a/src/libraries/System.Linq/src/System/Linq/AppendPrepend.SpeedOpt.cs b/src/libraries/System.Linq/src/System/Linq/AppendPrepend.SpeedOpt.cs
@@ -73,6 +73,13 @@ public override TSource[] ToArray()
             public override List<TSource> ToList()
             {
                 int count = GetCount(onlyIfCheap: true);
+
+                if (count == 1)
+                {
+                    // If GetCount returns 1, then _source is empty and only _item should be returned
+                    return new List<TSource>(1) { _item };
+                }
+
                 List<TSource> list = count == -1 ? new List<TSource>() : new List<TSource>(count);
                 if (!_appending)
                 {
@@ -122,17 +129,8 @@ private TSource[] LazyToArray()
 
                 TSource[] array = builder.ToArray();
 
-                int index = 0;
-                for (SingleLinkedNode<TSource>? node = _prepended; node != null; node = node.Linked)
-                {
-                    array[index++] = node.Item;
-                }
-
-                index = array.Length - 1;
-                for (SingleLinkedNode<TSource>? node = _appended; node != null; node = node.Linked)
-                {
-                    array[index--] = node.Item;
-                }
+                _prepended?.Fill(array);
+                _appended?.FillReversed(array);
 
                 return array;
             }
@@ -181,17 +179,11 @@ public override List<TSource> ToList()
                 int count = GetCount(onlyIfCheap: true);
                 List<TSource> list = count == -1 ? new List<TSource>() : new List<TSource>(count);
 
-                for (SingleLinkedNode<TSource>? node = _prepended; node != null; node = node.Linked)
-                {
-                    list.Add(node.Item);
-                }
+                _prepended?.Fill(SetCountAndGetSpan(list, _prependCount));
 
                 list.AddRange(_source);
 
-                if (_appended != null)
-                {
-                    list.AddRange(_appended.ToArray(_appendCount));
-                }
+                _appended?.FillReversed(SetCountAndGetSpan(list, list.Count + _appendCount));
 
                 return list;
             }

diff --git a/src/libraries/System.Linq/src/System/Linq/Lookup.SpeedOpt.cs b/src/libraries/System.Linq/src/System/Linq/Lookup.SpeedOpt.cs
@@ -10,22 +10,16 @@ public partial class Lookup<TKey, TElement> : IIListProvider<IGrouping<TKey, TEl
     {
         IGrouping<TKey, TElement>[] IIListProvider<IGrouping<TKey, TElement>>.ToArray()
         {
-            IGrouping<TKey, TElement>[] array = new IGrouping<TKey, TElement>[_count];
-            int index = 0;
-            Grouping<TKey, TElement>? g = _lastGrouping;
-            if (g != null)
+            IGrouping<TKey, TElement>[] array;
+            if (_count > 0)
             {
-                do
-                {
-                    g = g._next;
-                    Debug.Assert(g != null);
-
-                    array[index] = g;
-                    ++index;
-                }
-                while (g != _lastGrouping);
+                array = new IGrouping<TKey, TElement>[_count];
+                Fill(_lastGrouping, array);
+            }
+            else
+            {
+                array = Array.Empty<IGrouping<TKey, TElement>>();
             }
-
             return array;
         }
 
@@ -53,21 +47,33 @@ internal TResult[] ToArray<TResult>(Func<TKey, IEnumerable<TElement>, TResult> r
 
         List<IGrouping<TKey, TElement>> IIListProvider<IGrouping<TKey, TElement>>.ToList()
         {
-            List<IGrouping<TKey, TElement>> list = new List<IGrouping<TKey, TElement>>(_count);
-            Grouping<TKey, TElement>? g = _lastGrouping;
+            var list = new List<IGrouping<TKey, TElement>>(_count);
+            if (_count > 0)
+            {
+                Fill(_lastGrouping, Enumerable.SetCountAndGetSpan(list, _count));
+            }
+
+            return list;
+        }
+
+        private static void Fill(Grouping<TKey, TElement>? lastGrouping, Span<IGrouping<TKey, TElement>> results)
+        {
+            int index = 0;
+            Grouping<TKey, TElement>? g = lastGrouping;
             if (g != null)
             {
                 do
                 {
                     g = g._next;
                     Debug.Assert(g != null);
 
-                    list.Add(g);
+                    results[index] = g;
+                    ++index;
                 }
-                while (g != _lastGrouping);
+                while (g != lastGrouping);
             }
 
-            return list;
+            Debug.Assert(index == results.Length, "All list elements were not initialized.");
         }
 
         int IIListProvider<IGrouping<TKey, TElement>>.GetCount(bool onlyIfCheap) => _count;

diff --git a/src/libraries/System.Linq/src/System/Linq/Lookup.cs b/src/libraries/System.Linq/src/System/Linq/Lookup.cs
@@ -154,15 +154,20 @@ internal List<TResult> ToList<TResult>(Func<TKey, IEnumerable<TElement>, TResult
             Grouping<TKey, TElement>? g = _lastGrouping;
             if (g != null)
             {
+                Span<TResult> span = Enumerable.SetCountAndGetSpan(list, _count);
+                int index = 0;
                 do
                 {
                     g = g._next;
 
                     Debug.Assert(g != null);
                     g.Trim();
-                    list.Add(resultSelector(g._key, g._elements));
+                    span[index] = resultSelector(g._key, g._elements);
+                    ++index;
                 }
                 while (g != _lastGrouping);
+
+                Debug.Assert(index == _count, "All list elements were not initialized.");
             }
 
             return list;

diff --git a/src/libraries/System.Linq/src/System/Linq/Select.SpeedOpt.cs b/src/libraries/System.Linq/src/System/Linq/Select.SpeedOpt.cs
@@ -611,13 +611,7 @@ private TResult[] PreallocatingToArray(int count)
                 Debug.Assert(count == _source.GetCount(onlyIfCheap: true));
 
                 TResult[] array = new TResult[count];
-                int index = 0;
-                foreach (TSource input in _source)
-                {
-                    array[index] = _selector(input);
-                    ++index;
-                }
-
+                Fill(_source, array, _selector);
                 return array;
             }
 
@@ -640,20 +634,33 @@ public List<TResult> ToList()
                 {
                     case -1:
                         list = new List<TResult>();
+                        foreach (TSource input in _source)
+                        {
+                            list.Add(_selector(input));
+                        }
                         break;
                     case 0:
-                        return new List<TResult>();
+                        list = new List<TResult>();
+                        break;
                     default:
                         list = new List<TResult>(count);
+                        Fill(_source, SetCountAndGetSpan(list, count), _selector);
                         break;
                 }
 
-                foreach (TSource input in _source)
+                return list;
+            }
+
+            private static void Fill(IPartition<TSource> source, Span<TResult> results, Func<TSource, TResult> func)
+            {
+                int index = 0;
+                foreach (TSource item in source)
                 {
-                    list.Add(_selector(input));
+                    results[index] = func(item);
+                    ++index;
                 }
 
-                return list;
+                Debug.Assert(index == results.Length, "All list elements were not initialized.");
             }
 
             public int GetCount(bool onlyIfCheap)

diff --git a/src/libraries/System.Linq/src/System/Linq/SingleLinkedNode.cs b/src/libraries/System.Linq/src/System/Linq/SingleLinkedNode.cs
@@ -92,15 +92,36 @@ public TSource[] ToArray(int count)
             Debug.Assert(count == GetCount());
 
             TSource[] array = new TSource[count];
-            int index = count;
+            FillReversed(array);
+            return array;
+        }
+
+        /// <summary>
+        /// Fills a start of a span with the items of this node's singly-linked list.
+        /// </summary>
+        /// <param name="span">The span to fill. Must be the precise size required.</param>
+        public void Fill(Span<TSource> span)
+        {
+            int index = 0;
             for (SingleLinkedNode<TSource>? node = this; node != null; node = node.Linked)
             {
-                --index;
-                array[index] = node.Item;
+               span[index] = node.Item;
+               index++;
             }
+        }
 
-            Debug.Assert(index == 0);
-            return array;
+        /// <summary>
+        /// Fills the end of a span with the items of this node's singly-linked list in reverse.
+        /// </summary>
+        /// <param name="span">The span to fill.</param>
+        public void FillReversed(Span<TSource> span)
+        {
+            int index = span.Length;
+            for (SingleLinkedNode<TSource>? node = this; node != null; node = node.Linked)
+            {
+                --index;
+                span[index] = node.Item;
+            }
         }
     }
 }
diff --git a/src/libraries/System.Linq/src/System/Linq/ToCollection.cs b/src/libraries/System.Linq/src/System/Linq/ToCollection.cs
@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Generic;
+using System.Diagnostics;
 
 namespace System.Linq
 {
@@ -255,11 +256,21 @@ private static TSource[] HashSetToArray<TSource>(HashSet<TSource> set)
 
         private static List<TSource> HashSetToList<TSource>(HashSet<TSource> set)
         {
-            var result = new List<TSource>(set.Count);
+            int count = set.Count;
 
-            foreach (TSource item in set)
+            var result = new List<TSource>(count);
+            if (count > 0)
             {
-                result.Add(item);
+                Span<TSource> span = SetCountAndGetSpan(result, count);
+
+                int index = 0;
+                foreach (TSource item in set)
+                {
+                    span[index] = item;
+                    ++index;
+                }
+
+                Debug.Assert(index == span.Length, "All list elements were not initialized.");
             }
 
             return result;

diff --git a/src/libraries/System.Linq/tests/SelectTests.cs b/src/libraries/System.Linq/tests/SelectTests.cs
@@ -1107,6 +1107,13 @@ public void Select_SourceIsArrayTakeTake()
             Assert.Equal(new[] { 2 }, source.Take(10));
         }
 
+        [Fact]
+        public void Select_SourceIsIPartitionToArray()
+        {
+            Assert.Equal(Array.Empty<int>(), new List<int>().Order().Select(i => i * 2).ToArray());
+            Assert.Equal(new[] { 2, 4, 6, 8 }, new List<int> { 1, 2, 3, 4 }.Order().Select(i => i * 2).ToArray());
+        }
+
         [Fact]
         public void Select_SourceIsListSkipTakeCount()
         {
@@ -1134,6 +1141,13 @@ public void Select_SourceIsListSkipTakeToList()
             Assert.Empty(new List<int> { 1, 2, 3, 4 }.Select(i => i * 2).Skip(8).ToList());
         }
 
+        [Fact]
+        public void Select_SourceIsIPartitionToList()
+        {
+            Assert.Equal(Array.Empty<int>(), new List<int>().Order().Select(i => i * 2).ToList());
+            Assert.Equal(new[] { 2, 4, 6, 8 }, new List<int> { 1, 2, 3, 4 }.Order().Select(i => i * 2).ToList());
+        }
+
         [Theory]
         [MemberData(nameof(MoveNextAfterDisposeData))]
         public void MoveNextAfterDispose(IEnumerable<int> source)