Skip to content

Commit 3922062

Browse files
Reimplement LINQ ToList using SegmentedArrayBuilder to reduce allocations (#104365)
* Use SegmentedArrayBuilder in ToList * Implement SegmentedArrayBuilder for Select iterators * Implement SegmentedArrayBuilder for OfType and SelectMany iterators * Implement SegmentedArrayBuilder for SkipTake and Where iterators * Revert change to ToList for non-specialized cases as requested in review * Better naming for ToArray/ToList in Select enumerator when size is unknown --------- Co-authored-by: Eirik Tsarpalis <[email protected]>
1 parent 25a5085 commit 3922062

File tree

6 files changed

+104
-33
lines changed

6 files changed

+104
-33
lines changed

src/libraries/System.Linq/src/System/Linq/OfType.SpeedOpt.cs

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,17 +52,21 @@ public override TResult[] ToArray()
5252

5353
public override List<TResult> ToList()
5454
{
55-
var list = new List<TResult>();
55+
SegmentedArrayBuilder<TResult>.ScratchBuffer scratch = default;
56+
SegmentedArrayBuilder<TResult> builder = new(scratch);
5657

5758
foreach (object? item in _source)
5859
{
5960
if (item is TResult castItem)
6061
{
61-
list.Add(castItem);
62+
builder.Add(castItem);
6263
}
6364
}
6465

65-
return list;
66+
List<TResult> result = builder.ToList();
67+
builder.Dispose();
68+
69+
return result;
6670
}
6771

6872
public override TResult? TryGetFirst(out bool found)

src/libraries/System.Linq/src/System/Linq/SegmentedArrayBuilder.cs

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
using System.Diagnostics;
66
using System.Linq;
77
using System.Runtime.CompilerServices;
8+
using System.Runtime.InteropServices;
89

910
namespace System.Collections.Generic
1011
{
@@ -251,6 +252,27 @@ public readonly T[] ToArray()
251252
return result;
252253
}
253254

255+
/// <summary>Creates a list containing all of the elements in the builder.</summary>
256+
public readonly List<T> ToList()
257+
{
258+
List<T> result;
259+
int count = Count;
260+
261+
if (count != 0)
262+
{
263+
result = new List<T>(count);
264+
265+
CollectionsMarshal.SetCount(result, count);
266+
ToSpanInlined(CollectionsMarshal.AsSpan(result));
267+
}
268+
else
269+
{
270+
result = [];
271+
}
272+
273+
return result;
274+
}
275+
254276
/// <summary>Creates an array containing all of the elements in the builder.</summary>
255277
/// <param name="additionalLength">The number of extra elements of room to allocate in the resulting array.</param>
256278
public readonly T[] ToArray(int additionalLength)

src/libraries/System.Linq/src/System/Linq/Select.SpeedOpt.cs

Lines changed: 30 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
using System.Collections.Generic;
55
using System.Diagnostics;
6+
using System.Runtime.CompilerServices;
67
using System.Runtime.InteropServices;
78
using static System.Linq.Utilities;
89

@@ -31,15 +32,19 @@ public override TResult[] ToArray()
3132

3233
public override List<TResult> ToList()
3334
{
34-
var list = new List<TResult>();
35+
SegmentedArrayBuilder<TResult>.ScratchBuffer scratch = default;
36+
SegmentedArrayBuilder<TResult> builder = new(scratch);
3537

3638
Func<TSource, TResult> selector = _selector;
3739
foreach (TSource item in _source)
3840
{
39-
list.Add(selector(item));
41+
builder.Add(selector(item));
4042
}
4143

42-
return list;
44+
List<TResult> result = builder.ToList();
45+
builder.Dispose();
46+
47+
return result;
4348
}
4449

4550
public override int GetCount(bool onlyIfCheap)
@@ -657,7 +662,7 @@ public override IEnumerable<TResult2> Select<TResult2>(Func<TResult, TResult2> s
657662
return sourceFound ? _selector(input!) : default!;
658663
}
659664

660-
private TResult[] LazyToArray()
665+
private TResult[] ToArrayNoPresizing()
661666
{
662667
Debug.Assert(_source.GetCount(onlyIfCheap: true) == -1);
663668

@@ -691,24 +696,39 @@ public override TResult[] ToArray()
691696
int count = _source.GetCount(onlyIfCheap: true);
692697
return count switch
693698
{
694-
-1 => LazyToArray(),
699+
-1 => ToArrayNoPresizing(),
695700
0 => [],
696701
_ => PreallocatingToArray(count),
697702
};
698703
}
699704

705+
private List<TResult> ToListNoPresizing()
706+
{
707+
Debug.Assert(_source.GetCount(onlyIfCheap: true) == -1);
708+
709+
SegmentedArrayBuilder<TResult>.ScratchBuffer scratch = default;
710+
SegmentedArrayBuilder<TResult> builder = new(scratch);
711+
712+
Func<TSource, TResult> selector = _selector;
713+
foreach (TSource input in _source)
714+
{
715+
builder.Add(selector(input));
716+
}
717+
718+
List<TResult> result = builder.ToList();
719+
builder.Dispose();
720+
721+
return result;
722+
}
723+
700724
public override List<TResult> ToList()
701725
{
702726
int count = _source.GetCount(onlyIfCheap: true);
703727
List<TResult> list;
704728
switch (count)
705729
{
706730
case -1:
707-
list = new List<TResult>();
708-
foreach (TSource input in _source)
709-
{
710-
list.Add(_selector(input));
711-
}
731+
list = ToListNoPresizing();
712732
break;
713733
case 0:
714734
list = new List<TResult>();

src/libraries/System.Linq/src/System/Linq/SelectMany.SpeedOpt.cs

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,15 +48,19 @@ public override TResult[] ToArray()
4848

4949
public override List<TResult> ToList()
5050
{
51-
var list = new List<TResult>();
51+
SegmentedArrayBuilder<TResult>.ScratchBuffer scratch = default;
52+
SegmentedArrayBuilder<TResult> builder = new(scratch);
5253

5354
Func<TSource, IEnumerable<TResult>> selector = _selector;
54-
foreach (TSource element in _source)
55+
foreach (TSource item in _source)
5556
{
56-
list.AddRange(selector(element));
57+
builder.AddRange(selector(item));
5758
}
5859

59-
return list;
60+
List<TResult> result = builder.ToList();
61+
builder.Dispose();
62+
63+
return result;
6064
}
6165
}
6266
}

src/libraries/System.Linq/src/System/Linq/SkipTake.SpeedOpt.cs

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -491,25 +491,30 @@ public override TSource[] ToArray()
491491

492492
public override List<TSource> ToList()
493493
{
494-
var list = new List<TSource>();
495-
496494
using (IEnumerator<TSource> en = _source.GetEnumerator())
497495
{
498496
if (SkipBeforeFirst(en) && en.MoveNext())
499497
{
500498
int remaining = Limit - 1; // Max number of items left, not counting the current element.
501499
int comparand = HasLimit ? 0 : int.MinValue; // If we don't have an upper bound, have the comparison always return true.
502500

501+
SegmentedArrayBuilder<TSource>.ScratchBuffer scratch = default;
502+
SegmentedArrayBuilder<TSource> builder = new(scratch);
503503
do
504504
{
505505
remaining--;
506-
list.Add(en.Current);
506+
builder.Add(en.Current);
507507
}
508508
while (remaining >= comparand && en.MoveNext());
509+
510+
List<TSource> result = builder.ToList();
511+
builder.Dispose();
512+
513+
return result;
509514
}
510515
}
511516

512-
return list;
517+
return [];
513518
}
514519

515520
private bool SkipBeforeFirst(IEnumerator<TSource> en) => SkipBefore(_minIndexInclusive, en);

src/libraries/System.Linq/src/System/Linq/Where.SpeedOpt.cs

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -55,18 +55,22 @@ public override TSource[] ToArray()
5555

5656
public override List<TSource> ToList()
5757
{
58-
var list = new List<TSource>();
58+
SegmentedArrayBuilder<TSource>.ScratchBuffer scratch = default;
59+
SegmentedArrayBuilder<TSource> builder = new(scratch);
5960

6061
Func<TSource, bool> predicate = _predicate;
6162
foreach (TSource item in _source)
6263
{
6364
if (predicate(item))
6465
{
65-
list.Add(item);
66+
builder.Add(item);
6667
}
6768
}
6869

69-
return list;
70+
List<TSource> result = builder.ToList();
71+
builder.Dispose();
72+
73+
return result;
7074
}
7175

7276
public override TSource? TryGetFirst(out bool found)
@@ -199,17 +203,21 @@ public static TSource[] ToArray(ReadOnlySpan<TSource> source, Func<TSource, bool
199203

200204
public static List<TSource> ToList(ReadOnlySpan<TSource> source, Func<TSource, bool> predicate)
201205
{
202-
var list = new List<TSource>();
206+
SegmentedArrayBuilder<TSource>.ScratchBuffer scratch = default;
207+
SegmentedArrayBuilder<TSource> builder = new(scratch);
203208

204209
foreach (TSource item in source)
205210
{
206211
if (predicate(item))
207212
{
208-
list.Add(item);
213+
builder.Add(item);
209214
}
210215
}
211216

212-
return list;
217+
List<TSource> result = builder.ToList();
218+
builder.Dispose();
219+
220+
return result;
213221
}
214222

215223
public override TSource? TryGetFirst(out bool found)
@@ -398,17 +406,21 @@ public static TResult[] ToArray(ReadOnlySpan<TSource> source, Func<TSource, bool
398406

399407
public static List<TResult> ToList(ReadOnlySpan<TSource> source, Func<TSource, bool> predicate, Func<TSource, TResult> selector)
400408
{
401-
var list = new List<TResult>();
409+
SegmentedArrayBuilder<TResult>.ScratchBuffer scratch = default;
410+
SegmentedArrayBuilder<TResult> builder = new(scratch);
402411

403412
foreach (TSource item in source)
404413
{
405414
if (predicate(item))
406415
{
407-
list.Add(selector(item));
416+
builder.Add(selector(item));
408417
}
409418
}
410419

411-
return list;
420+
List<TResult> result = builder.ToList();
421+
builder.Dispose();
422+
423+
return result;
412424
}
413425

414426
public override TResult? TryGetFirst(out bool found) => TryGetFirst(_source, _predicate, _selector, out found);
@@ -538,19 +550,23 @@ public override TResult[] ToArray()
538550

539551
public override List<TResult> ToList()
540552
{
541-
var list = new List<TResult>();
553+
SegmentedArrayBuilder<TResult>.ScratchBuffer scratch = default;
554+
SegmentedArrayBuilder<TResult> builder = new(scratch);
542555

543556
Func<TSource, bool> predicate = _predicate;
544557
Func<TSource, TResult> selector = _selector;
545558
foreach (TSource item in _source)
546559
{
547560
if (predicate(item))
548561
{
549-
list.Add(selector(item));
562+
builder.Add(selector(item));
550563
}
551564
}
552565

553-
return list;
566+
List<TResult> result = builder.ToList();
567+
builder.Dispose();
568+
569+
return result;
554570
}
555571

556572
public override TResult? TryGetFirst(out bool found)

0 commit comments

Comments
 (0)