|
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
using System.Buffers;
using System.Diagnostics;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
namespace System.Collections.Generic
{
/// <summary>Provides a helper for efficiently building arrays and lists.</summary>
/// <remarks>This is implemented as an inline array of rented arrays.</remarks>
/// <typeparam name="T">Specifies the element type of the collection being built.</typeparam>
internal ref struct SegmentedArrayBuilder<T>
{
/// <summary>The size to use for the first segment that's stack allocated by the caller.</summary>
/// <remarks>
/// This value needs to be small enough that we don't need to be overly concerned about really large
/// value types. It's not unreasonable for a method to say it has 8 locals of a T, and that's effectively
/// what this is.
/// </remarks>
private const int ScratchBufferSize = 8;
/// <summary>Minimum size to request renting from the pool.</summary>
private const int MinimumRentSize = 16;
/// <summary>The array of segments.</summary>
/// <remarks><see cref="_segmentsCount"/> is how many of the segments are valid in <see cref="_segments"/>, not including <see cref="_firstSegment"/>.</remarks>
private Arrays _segments;
/// <summary>The scratch buffer provided by the caller.</summary>
/// <remarks>This is treated as the initial segment, before anything in <see cref="_segments"/>.</remarks>
private Span<T> _firstSegment;
/// <summary>The current span. This points either to <see cref="_firstSegment"/> or to <see cref="_segments"/>[<see cref="_segmentsCount"/> - 1].</summary>
private Span<T> _currentSegment;
/// <summary>The count of segments in <see cref="_segments"/> that are valid.</summary>
/// <remarks>All but the last are known to be fully populated.</remarks>
private int _segmentsCount;
/// <summary>The total number of elements in all but the current/last segment.</summary>
private int _countInFinishedSegments;
/// <summary>The number of elements in the current/last segment.</summary>
private int _countInCurrentSegment;
/// <summary>Initialize the builder.</summary>
/// <param name="scratchBuffer">A buffer that can be used as part of the builder.</param>
public SegmentedArrayBuilder(Span<T> scratchBuffer)
{
_currentSegment = _firstSegment = scratchBuffer;
}
/// <summary>Clean up the resources used by the builder.</summary>
public void Dispose()
{
int segmentsCount = _segmentsCount;
if (segmentsCount != 0)
{
ReturnArrays(segmentsCount);
}
}
private void ReturnArrays(int segmentsCount)
{
Debug.Assert(segmentsCount > 0);
ReadOnlySpan<T[]> segments = _segments;
// We need to return all rented arrays to the pool, and if the arrays contain any references,
// we want to clear them first so that the pool doesn't artificially root contained objects.
if (RuntimeHelpers.IsReferenceOrContainsReferences<T>())
{
// Return all but the last segment. All of these are full and need to be entirely cleared.
segmentsCount--;
foreach (T[] segment in segments.Slice(0, segmentsCount))
{
Array.Clear(segment);
ArrayPool<T>.Shared.Return(segment);
}
// For the last segment, we can clear only what we know was used.
T[] currentSegment = segments[segmentsCount];
Array.Clear(currentSegment, 0, _countInCurrentSegment);
ArrayPool<T>.Shared.Return(currentSegment);
}
else
{
// Return every rented array without clearing.
for (int i = 0; i < segments.Length; i++)
{
T[] segment = segments[i];
if (segment is null)
{
break;
}
ArrayPool<T>.Shared.Return(segment);
}
}
}
/// <summary>Gets the number of elements in the builder.</summary>
public readonly int Count => checked(_countInFinishedSegments + _countInCurrentSegment);
/// <summary>Adds an item into the builder.</summary>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void Add(T item)
{
Span<T> currentSegment = _currentSegment;
int countInCurrentSegment = _countInCurrentSegment;
if ((uint)countInCurrentSegment < (uint)currentSegment.Length)
{
currentSegment[countInCurrentSegment] = item;
_countInCurrentSegment++;
}
else
{
AddSlow(item);
}
}
[MethodImpl(MethodImplOptions.NoInlining)]
private void AddSlow(T item)
{
Expand();
_currentSegment[0] = item;
_countInCurrentSegment = 1;
}
/// <summary>Adds a collection of items into the builder.</summary>
public void AddRange(IEnumerable<T> source)
{
if (source is ICollection<T> collection)
{
int collectionCount = collection.Count;
// If the source is empty, there's nothing to add.
if (collectionCount == 0)
{
return;
}
// If the source is something from which we can get a span, e.g. a T[] or a List<T>,
// we can do one or two copies to handle copying everything, even if we need to split
// across segments.
if (Enumerable.TryGetSpan(source, out ReadOnlySpan<T> sourceSpan))
{
int availableSpaceInCurrentSpan = _currentSegment.Length - _countInCurrentSegment;
ReadOnlySpan<T> sourceSlice = sourceSpan.Slice(0, Math.Min(availableSpaceInCurrentSpan, sourceSpan.Length));
sourceSlice.CopyTo(_currentSegment.Slice(_countInCurrentSegment));
_countInCurrentSegment += sourceSlice.Length;
sourceSlice = sourceSpan.Slice(sourceSlice.Length);
if (!sourceSlice.IsEmpty)
{
Expand(sourceSlice.Length);
sourceSlice.CopyTo(_currentSegment);
_countInCurrentSegment = sourceSlice.Length;
}
return;
}
// Otherwise, since we have an ICollection, we'd like to use ICollection.CopyTo, but it
// requires targeting an array, so we can't use it if we're using a scratch buffer.
bool currentSegmentIsScratchBufferWithRemainingSpace = _segmentsCount == 0 && _countInCurrentSegment < _currentSegment.Length;
if (!currentSegmentIsScratchBufferWithRemainingSpace)
{
// It also only works if we can copy the whole collection in one go, which
// means we need enough space in the current segment to hold the whole collection
// or we need to be at the end of the current segment so we can allocate a
// new one without leaving any holes.
int remainingSpaceInCurrentSegment = _currentSegment.Length - _countInCurrentSegment;
// If there's no space remaining in the current segment, we can just expand
// into a new segment that we ensure is large enough.
if (remainingSpaceInCurrentSegment == 0)
{
Expand(collectionCount);
collection.CopyTo(_segments[_segmentsCount - 1], 0);
_countInCurrentSegment = collectionCount;
return;
}
// If there's enough space remaining in the current segment, we can also just copy into it.
if (collectionCount <= remainingSpaceInCurrentSegment)
{
collection.CopyTo(_segments[_segmentsCount - 1], _countInCurrentSegment);
_countInCurrentSegment += collectionCount;
return;
}
// Otherwise, we're forced to fall back to enumeration.
}
}
// Fall back to enumerating and adding each element individually.
AddNonICollectionRangeInlined(source);
}
/// <summary>Adds a collection of items into the builder.</summary>
/// <remarks>
/// The implementation assumes the caller has already ruled out the source being
/// and ICollection and thus doesn't bother checking to see if it is.
/// </remarks>
[MethodImpl(MethodImplOptions.NoInlining)]
public void AddNonICollectionRange(IEnumerable<T> source) =>
AddNonICollectionRangeInlined(source);
/// <summary>Adds a collection of items into the builder.</summary>
/// <remarks>
/// The implementation assumes the caller has already ruled out the source being
/// and ICollection and thus doesn't bother checking to see if it is.
/// </remarks>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
internal void AddNonICollectionRangeInlined(IEnumerable<T> source)
{
Span<T> currentSegment = _currentSegment;
int countInCurrentSegment = _countInCurrentSegment;
foreach (T item in source)
{
if ((uint)countInCurrentSegment < (uint)currentSegment.Length)
{
currentSegment[countInCurrentSegment] = item;
countInCurrentSegment++;
}
else
{
Expand();
currentSegment = _currentSegment;
currentSegment[0] = item;
countInCurrentSegment = 1;
}
}
_countInCurrentSegment = countInCurrentSegment;
}
/// <summary>Creates an array containing all of the elements in the builder.</summary>
public readonly T[] ToArray()
{
T[] result;
int count = Count;
if (count != 0)
{
result = GC.AllocateUninitializedArray<T>(count);
ToSpanInlined(result);
}
else
{
result = [];
}
return result;
}
/// <summary>Creates a list containing all of the elements in the builder.</summary>
public readonly List<T> ToList()
{
List<T> result;
int count = Count;
if (count != 0)
{
result = new List<T>(count);
CollectionsMarshal.SetCount(result, count);
ToSpanInlined(CollectionsMarshal.AsSpan(result));
}
else
{
result = [];
}
return result;
}
/// <summary>Creates an array containing all of the elements in the builder.</summary>
/// <param name="additionalLength">The number of extra elements of room to allocate in the resulting array.</param>
public readonly T[] ToArray(int additionalLength)
{
T[] result;
int count = checked(Count + additionalLength);
if (count != 0)
{
result = GC.AllocateUninitializedArray<T>(count);
ToSpanInlined(result);
}
else
{
result = [];
}
return result;
}
/// <summary>Populates the destination span with all of the elements in the builder.</summary>
/// <param name="destination">The destination span.</param>
[MethodImpl(MethodImplOptions.NoInlining)]
public readonly void ToSpan(Span<T> destination) => ToSpanInlined(destination);
/// <summary>Populates the destination span with all of the elements in the builder.</summary>
/// <param name="destination">The destination span.</param>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private readonly void ToSpanInlined(Span<T> destination)
{
int segmentsCount = _segmentsCount;
if (segmentsCount != 0)
{
// Copy the first segment
ReadOnlySpan<T> firstSegment = _firstSegment;
firstSegment.CopyTo(destination);
destination = destination.Slice(firstSegment.Length);
// Copy the 0..N-1 segments
segmentsCount--;
if (segmentsCount != 0)
{
foreach (T[] arr in ((ReadOnlySpan<T[]>)_segments).Slice(0, segmentsCount))
{
ReadOnlySpan<T> segment = arr;
segment.CopyTo(destination);
destination = destination.Slice(segment.Length);
}
}
}
// Copy the last segment
_currentSegment.Slice(0, _countInCurrentSegment).CopyTo(destination);
}
/// <summary>Appends a new segment onto the builder.</summary>
/// <param name="minimumRequired">The minimum amount of space to allocate in a new segment being appended.</param>
private void Expand(int minimumRequired = MinimumRentSize)
{
if (minimumRequired < MinimumRentSize)
{
minimumRequired = MinimumRentSize;
}
// Update our count of the number of elements in the arrays.
// If we know we're exceeding the maximum allowed array length, throw.
int currentSegmentLength = _currentSegment.Length;
checked { _countInFinishedSegments += currentSegmentLength; }
if (_countInFinishedSegments > Array.MaxLength)
{
throw new OutOfMemoryException();
}
// Use a typical doubling algorithm to decide the length of the next array
// and allocate it. We want to double the current array length, but if the
// minimum required is larger than that, use the minimum required. And if
// doubling would result in going above the max array length, only use the
// max array length, as List<T> does.
int newSegmentLength = (int)Math.Min(Math.Max(minimumRequired, currentSegmentLength * 2L), Array.MaxLength);
_currentSegment = _segments[_segmentsCount] = ArrayPool<T>.Shared.Rent(newSegmentLength);
_segmentsCount++;
}
#pragma warning disable IDE0044 // Add readonly modifier
#pragma warning disable IDE0051 // Remove unused private members
/// <summary>A struct to hold all of the T[]s that compose the full result set.</summary>
/// <remarks>
/// Starting at the minimum size of <see cref="MinimumRentSize"/>, and with a minimum of doubling
/// on every growth, this is large enough to hold the maximum number arrays that could result
/// until the total length has exceeded Array.MaxLength.
/// </remarks>
[InlineArray(27)]
private struct Arrays
{
private T[] _values;
}
/// <summary>Provides a stack-allocatable buffer for use as an argument to the builder.</summary>
[InlineArray(ScratchBufferSize)]
public struct ScratchBuffer
{
private T _item;
}
#pragma warning restore IDE0051
#pragma warning restore IDE0044
}
}
|