Lost in Linq source code - Part 4
In this part we will focus on GroupBy, the grouping operation. We can often find ourselves in a situation where we have multiple items, but rather than filtering, we want to group them instead to process further.
Using a dictionary in this case can resolve this problem, but we will end up dealing with Dictionary<TKey, IEnumerable<TValue>>. This approach can resolve the problem, but it feels like too much, especially when the result won’t be modified.
That’s why Lookup<TKey, TElement> was introduced.
Lookup is a data structure similar to the Dictionary, where a collection of values is mapped to a key. Furthermore, a dictionary is mutable, whereas a lookup is not.
How to group books by colour?
Let’s sit back and imagine that we have a collection of books and we want to group them by colour (red to reds, green to greens, etc.). As our books haven’t arrived yet, we can’t exactly prepare groups in advance (what if there isn’t a single green book?). We can’t predict if one of them will be damaged in delivery or borrowed by someone. So, we can’t exactly prepare here; all we know is the fact that we will receive books (elements) and our books will be grouped by colour (key).
Time passes by and our books (elements) finally arrive; they’re unordered in a very long box that allows us to see only the first book. We take the book, look at it — it’s blue—and we don’t have a place for blue books. So, we take a sticky note, write BLUE on it, and place it on the ground. We’ve created a new group here. Then, we place a book behind our sticky note. This activity created a new group of books with the BLUE key and one blue book. We also place another sticky note with an arrow pointing to where the next group will be. As we don’t expect a new group (what if all books are blue?), this arrow points to the blues. Remember where to place the next group.
We look for the second book—it’s green—so we search through our group keys. There is one group; this group points to itself. We don’t have the GREEN key, only BLUE, so we create the new GREEN group and rotate the arrow in the BLUE group, making it point to green. The new group, on the other hand, points to BLUE. Again, remember where to place the next group.
We take a third book, and it’s green again. We browse our keys—blue, green… aha! It’s the place for our new book.
Skipping ahead to the end of our work, we could get results as follows: 20 green books, 12 blue books, no red, 30 white books, 22 black, 4 violet.
Similarities to the HashSet
Like before, we will use hashes, buckets, and entries, but our stored element will be more complex. Additionally, we can improve the bucket index calculation mechanism by introducing the modulo mask.
As we resize our bucket array by the power of 2, we can easily calculate the next size using a mask and the AND operator.
Taking our book grouping story from the previous paragraph, we were pointing to the next group by insertion order, but we will also need to point to the next bucket (like we did in HashSet).
Lastly, our stored element will be more complex, because we will need to store a unique key and non-unique elements.
Implementation
After going through a real-world example (with additional steps) and discussing similarities and differences, we can step into the implementation.
Base structs to work with
Knowing that we will store not only a unique key but also elements matching this key, we need to prepare structs that allow us to do such a thing. One of them also needs to satisfy the immutability requirement.
public struct Grouping<TKey, TElement> : IDisposable
{
public uint HashCode;
public TKey Key;
public int HashNext;
public int Next;
public TElement[]? Elements;
public int ElementCount;
public int ElementCapacity => Elements.Length;
public TElement this[int index] =>
index > Elements.Length - 1 ? throw new IndexOutOfRangeException() : Elements[index];
public void Dispose()
{
ArrayPool<TElement>.Shared.Return(Elements);
}
}
This is our base grouping that we will work with most of the time. It stores a hash code, a key, and an array of elements. ElementCount and ElementCapacity allow us to assign and track when we need to resize the element’s array. As an element’s array can grow enormous (imagine that we have 50,000 green books), it needs to be rented from the array pool and then disposed. HashNext and Next, on the other hand, represent the next group in the same bucket and the next group in insertion order.
public readonly struct ReadOnlyGrouping<TKey, TElement>
{
private readonly Grouping<TKey, TElement> _grouping;
public ReadOnlyGrouping(in Grouping<TKey, TElement> grouping)
{
_grouping = grouping;
}
public TKey Key => _grouping.Key;
public int Count => _grouping.ElementCount;
public ReadOnlySpan<TElement> Elements => _grouping.Elements != null
? _grouping.Elements.AsSpan(0, _grouping.ElementCount)
: ReadOnlySpan<TElement>.Empty;
}
Aside from the struct that will contain all necessary information to work with, we also need to expose Grouping safely to the user. Here, we have a readonly grouping variant exposing only Key, Count of elements, and Span of the Element’s array. The cost of this struct is negligible, as we’re passing it by reference.
Lookup
When we’re picking the first book from our story, we didn’t know exactly what’s inside and if someone didn’t want to take out some books. Approaching this problem will lead us to a requirement disallowing us to load the whole box (input) at the beginning.
public ref struct Lookup<TKey, TElement> : IDisposable
{
private const double LoadFactor = 0.75;
private const int InitialCapacity = 16;
private const int InitialElementCapacity = 4;
private Grouping<TKey, TElement>[] _entries;
private int[] _buckets;
private readonly IEqualityComparer<TKey> _comparer;
private int _count;
private int _threshold;
private int _lastGroupIndex;
private uint _bucketMask;
Code follows the idea that we discussed before, also we used the same approach as in MinimalHashSet with initial capacity and threshold.
First difference is in _entries field, where we’ll use the Grouping<TKey, TElement>.
Second, bucket mask to calculate index.
private Lookup(IEqualityComparer<TKey>? comparer = null)
{
_comparer = comparer ?? EqualityComparer<TKey>.Default;
_bucketMask = InitialCapacity - 1;
_threshold = (int)(InitialCapacity * LoadFactor);
_buckets = ArrayPool<int>.Shared.Rent(InitialCapacity);
_entries = ArrayPool<Grouping<TKey, TElement>>.Shared.Rent(InitialCapacity);
_buckets.AsSpan(0, InitialCapacity).Clear();
_count = 0;
_lastGroupIndex = 0;
}
Again, following MinimalHashSet’s Construct method, we assign fields similarly.
So far, our start of Lookup struct should look like this:
public ref struct Lookup<TKey, TElement> : IDisposable
{
private const double LoadFactor = 0.75;
private const int InitialCapacity = 16;
private const int InitialElementCapacity = 4;
private Grouping<TKey, TElement>[] _entries;
private int[] _buckets;
private readonly IEqualityComparer<TKey> _comparer;
private int _count;
private int _threshold;
private int _lastGroupIndex;
private uint _bucketMask;
private Lookup(IEqualityComparer<TKey>? comparer = null)
{
_comparer = comparer ?? EqualityComparer<TKey>.Default;
_bucketMask = InitialCapacity - 1;
_threshold = (int)(InitialCapacity * LoadFactor);
_buckets = ArrayPool<int>.Shared.Rent(InitialCapacity);
_entries = ArrayPool<Grouping<TKey, TElement>>.Shared.Rent(InitialCapacity);
_buckets.AsSpan(0, InitialCapacity).Clear();
_count = 0;
_lastGroupIndex = 0;
}
Let’s move on to essential methods.
private uint InternalGetHashCode(ref TKey item)
{
return (uint)(_comparer.GetHashCode(item) & 0x7FFFFFFF);
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private int GetIndex(ref uint hashCode)
{
return (int)(hashCode & _bucketMask);
}
Both of them are taken straight from HashSet, but for GetIndex we’ll use our mask, for quick calculation using bitwise AND.
//The same as we used in MinimalHashSet
private void Resize()
{
uint newSize = System.Numerics.BitOperations.RoundUpToPowerOf2((uint)_entries.Length * 2);
var newBucket = ArrayPool<int>.Shared.Rent((int)newSize);
var newEntries = ArrayPool<Grouping<TKey, TElement>>.Shared.Rent((int)newSize);
newBucket.AsSpan(0, (int)newSize).Clear();
_entries.AsSpan(0, _count).CopyTo(newEntries);
_bucketMask = newSize - 1;
for (int i = 0; i < _count; i++)
{
ref Grouping<TKey, TElement> entry = ref newEntries[i];
int bucketIndex = GetIndex(ref entry.HashCode);
entry.HashNext = newBucket[bucketIndex];
newBucket[bucketIndex] = i + 1;
}
ArrayPool<int>.Shared.Return(_buckets, clearArray: true);
ArrayPool<Grouping<TKey, TElement>>.Shared.Return(_entries, clearArray: true);
_threshold = (int)(newSize * LoadFactor);
_buckets = newBucket;
_entries = newEntries;
}
In this method we need to use Grouping and remember to change _bucketMask to a new size.
public void Add(TKey key, TElement value)
{
if (_count >= _threshold)
{
//Because we need to resize when threshold hit
Resize();
}
uint hashCode = InternalGetHashCode(ref key);
var bucketIdx = GetIndex(ref hashCode);
int currentEntry = _buckets[bucketIdx];
//Because we need to check if we already have the key in the bucket
while (currentEntry > 0)
{
ref Grouping<TKey, TElement> entry = ref _entries[currentEntry - 1];
//In hashset we would rather stop on the first match and notify the caller
if (entry.HashCode == hashCode && _comparer.Equals(entry.Key, key))
{
AddToGroup(ref entry, value);
return;
}
currentEntry = entry.HashNext;
}
//Because we checked that there is no matching key
CreateGroup(key, hashCode, value, bucketIdx);
}
Add is where we need to include our Grouping and think about assignment to group (AddToGroup), and creating a new one (CreateGroup).
Typically, when inserting to a Hashset, we want to detect collision before we resize our _buckets and _entries, but this time we don’t care, as we want to allow duplicates to be included in one of groups.
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private void AddToGroup(ref Grouping<TKey, TElement> entry, TElement value)
{
//Because every Key has its own group, and we want to resize it if needed
if (entry.ElementCount >= entry.ElementCapacity)
{
int newCapacity = entry.ElementCapacity == 0 ? InitialElementCapacity : entry.ElementCapacity * 2;
var newArray = ArrayPool<TElement>.Shared.Rent(newCapacity);
if (entry.Elements != null)
{
entry.Elements.AsSpan(0, entry.ElementCount).CopyTo(newArray);
ArrayPool<TElement>.Shared.Return(entry.Elements,
clearArray: RuntimeHelpers.IsReferenceOrContainsReferences<TElement>());
}
entry.Elements = newArray;
}
//Because element count changed
entry.Elements[entry.ElementCount++] = value;
}
If Add did find a collision, we know there is a group associated with element that we want to add.
Currently, we need to check if there is a place for new element (and resize Grouping elements if not), assign it under next free index, then increment ElementCount to keep track.
private void CreateGroup(TKey key, uint hashCode, TElement value, int bucketIndex)
{
if (_count >= _entries.Length)
{
Resize();
//Because we resized the array, we need to get the new bucket index
bucketIndex = GetIndex(ref hashCode);
}
int newIndex = _count + 1;
ref Grouping<TKey, TElement> newEntry = ref _entries[_count];
newEntry.HashCode = hashCode;
newEntry.Key = key;
newEntry.HashNext = _buckets[bucketIndex];
//Because arrays can be huge
newEntry.Elements = ArrayPool<TElement>.Shared.Rent(InitialElementCapacity);
newEntry.Elements[0] = value;
newEntry.ElementCount = 1;
_buckets[bucketIndex] = newIndex;
if (_lastGroupIndex == 0)
{
newEntry.Next = newIndex;
_lastGroupIndex = newIndex;
}
else
{
ref Grouping<TKey, TElement> lastGroup = ref _entries[_lastGroupIndex - 1];
newEntry.Next = lastGroup.Next;
lastGroup.Next = newIndex;
_lastGroupIndex = newIndex;
}
//Because we added a new entry
_count++;
}
When Lookup doesn’t find a collision, that means we’ve got a new element that should be labeled under a new key.
Again, we need to check if we can add a new group to _entries (in AddToGroup, we were checking the Grouping.Elements array only) and resize if necessary (forcing us to recalculate the index again). Then, we can assign values of a new entry under the first free index and assign this index to the calculated bucket index.
Lastly, we remember the index of a future new Grouping and point to the last Grouping (keeping insertion order).
public static Lookup<TKey, TElement> Create<TEnumerator>(TEnumerator source, Func<TElement, TKey> selector,
IEqualityComparer<TKey>? comparer = null) where TEnumerator : IStructEnumerator<TElement>, allows ref struct
{
var lookup = new Lookup<TKey, TElement>(comparer);
ReadOnlySpan<TElement> span = default;
if (source.GetUnderlying(ref span))
{
foreach (var element in span)
{
var key = selector(element);
if (key != null)
{
lookup.Add(key, element);
}
}
}
else
{
TElement element = default;
while (source.Next(ref element))
{
var key = selector(element);
if (key != null)
{
lookup.Add(key, element);
}
}
}
return lookup;
}
We’re introducing the Create method due to the fact that we don’t know what the source enumerator looks like. From the Lookup perspective, it is unpredictable, and we can retrieve elements using a span (from GetUnderlying) or by iterating through the source enumerator. This method will be responsible for creating every new Lookup instance.
public readonly void Dispose()
{
if (_buckets != null)
{
ArrayPool<int>.Shared.Return(_buckets);
foreach (var entry in _entries)
{
ArrayPool<TElement>.Shared.Return(entry.Elements);
}
ArrayPool<Grouping<TKey, TElement>>.Shared.Return(_entries);
}
}
Don’t forget to dispose resources we rented.
LookupIterator
Creation could be done here to jump into InternalGroupBy like we did in previous parts, exposing all Groupings as an array or span using one method.
This step can be done smarter than iterating over them in the lookup and storing it in a field in the GroupBy iterator.
We will introduce LookupIterator to iterate over every grouping. To do such a thing, we need to create our struct first (in my case, inside the Lookup<TKey, TElement> struct):
//Inside lookup
public ref struct LookupIterator
{
private int _index;
private readonly int _count;
private readonly ReadOnlySpan<Grouping<TKey, TElement>> _entries;
public LookupIterator(ReadOnlySpan<Grouping<TKey, TElement>> groups)
{
_entries = groups;
_count = groups.Length;
_index = 0;
}
It’s a straightforward iterator, without any oddities. To make it faster, it’ll be a ref struct, and we want to use ReadOnlySpan of our grouping to avoid copying the whole _entries array.
//Inside lookup, but outside LookupIterator
public int Count => _count;
public LookupIterator GetIterator()
{
return new LookupIterator(_entries.AsSpan(0, _count));
}
Of course we need to expose our LookupIterator in Lookup to be able to iterate.
//In LookupIterator
public bool Next()
{
while (++_index <= _count)
{
ref readonly Grouping<TKey, TElement> entry = ref _entries[_index - 1];
if (entry.Elements != null && entry.ElementCount > 0)
{
return true;
}
}
return false;
}
public ReadOnlyGrouping<TKey, TElement> Current
{
get
{
if(_index < 0 || _index > _count)
{
throw new InvalidOperationException();
}
ref readonly var entry = ref _entries[_index - 1];
return new ReadOnlyGrouping<TKey, TElement>(in entry);
}
}
public int Remaining => _count - _index;
Heading back to LookupIterator, we will rely on our _index as we don’t need to calculate buckets. Then we will allow InternalGroupBy to advance using Next() and receive ReadOnlyGrouping<TKey, TElement> using Current, keeping in mind the immutability requirement and hiding unnecessary information.
InternalGroupBy
Finally, we can lay our hands on implementing the InternalGroupBy struct.
Our effort will be used there, as we will use our Lookup and LookupIterator.
We need to remember the fact that we shouldn’t initialize everything during the constructor.
public ref struct InternalGroupBy<TEnumerator, TKey, TElement> : IStructEnumerator<ReadOnlyGrouping<TKey, TElement>>
where TEnumerator : struct, IStructEnumerator<TElement>, allows ref struct
{
private TEnumerator _source;
private readonly Func<TElement, TKey> _selector;
private readonly IEqualityComparer<TKey> _comparer;
private Base.Lookup<TKey, TElement> _lookup;
private Base.Lookup<TKey,TElement>.LookupIterator _lookupIterator;
private bool _initialized;
public InternalGroupBy(TEnumerator source, Func<TElement, TKey> selector, IEqualityComparer<TKey>? comparer = null)
{
_source = source;
_selector = selector;
_comparer = comparer ?? EqualityComparer<TKey>.Default;
_lookup = default;
_initialized = false;
}
public void Dispose()
{
_source.Dispose();
_lookup.Dispose();
}
In part 3, we were selecting a key using _selector, and here it will find its usage too, as we need to select a key from an element.
To avoid early materialization of data before iterating, we introduce the _initialized flag and set _lookup as default; later, we will check if the iterator initialized the data.
public bool Next(ref ReadOnlyGrouping<TKey, TElement> current)
{
Initialize();
if (_lookupIterator.Next())
{
current = _lookupIterator.Current;
return true;
}
current = default;
return false;
}
public bool GetCountToLeftEnumerate(out int count)
{
Initialize();
count = _lookupIterator.Remaining;
return true;
}
public bool GetUnderlying(ref ReadOnlySpan<ReadOnlyGrouping<TKey, TElement>> span)
{
span = default;
return false;
}
public bool TryCopy(ref Span<ReadOnlyGrouping<TKey, TElement>> destination, int offset)
{
return false;
}
private void Initialize()
{
if (!_initialized)
{
_lookup = Base.Lookup<TKey, TElement>.Create(_source, _selector, _comparer);
_lookupIterator = _lookup.GetIterator();
_initialized = true;
}
}
Thanks to the Lookup and LookupIterator design, there isn’t much work to do in methods defined by the IStructEnumerator interface.
The major change compared to the previous iterator is checking if the lookup got initialized. We can notice that the Next method differs, as we rely on _lookupIterator.Next() and _lookupIterator.Current to move through groupings.
Lastly, the Initialize() method creates a new instance of Lookup, then we can get the LookupIterator from GetIterator to work with our groupings.
Don’t forget to set _initialized to true.
public static class GroupByStruct
{
public static StructEnumerable<InternalGroupBy<TEnumerator, TKey, TElement>, ReadOnlyGrouping<TKey, TElement>>
StructGroupBy<TEnumerator, TElement, TKey>(
this StructEnumerable<TEnumerator, TElement> source,
Func<TElement, TKey> keySelector,
IEqualityComparer<TKey>? comparer = null)
where TEnumerator : struct, IStructEnumerator<TElement>, allows ref struct
{
return new(new InternalGroupBy<TEnumerator, TKey, TElement>(source.Enumerator, keySelector, comparer));
}
public static StructEnumerable<InternalGroupBy<ArrayStructEnumerator<TElement>, TKey, TElement>,
ReadOnlyGrouping<TKey, TElement>>
StructGroupBy<TElement, TKey>(
this TElement[] array,
Func<TElement, TKey> keySelector,
IEqualityComparer<TKey>? comparer = null)
{
return new(new InternalGroupBy<ArrayStructEnumerator<TElement>, TKey, TElement>(
new ArrayStructEnumerator<TElement>(array),
keySelector, comparer));
}
}
Finishing our implementation, we shouldn’t forget about extension methods to connect our InternalGroupBy to the rest of our Linq implementation.
ToArrayHelper
We shouldn’t forget about the possibility of retrieving our groupings as an array.
public static ReadOnlyGrouping<TKey, TElement>[] ToArray<TEnumerator, TKey, TElement>(
this StructEnumerable<InternalGroupBy<TEnumerator, TKey, TElement>, ReadOnlyGrouping<TKey, TElement>> enumerator)
where TEnumerator : struct, IStructEnumerator<TElement>, allows ref struct
{
ReadOnlyGrouping<TKey, TElement>[] result = null;
if (enumerator.Enumerator.GetCountToLeftEnumerate(out int count))
{
result = GC.AllocateUninitializedArray<ReadOnlyGrouping<TKey, TElement>>(count);
ReadOnlyGrouping<TKey, TElement> current = default;
int idx = 0;
while (enumerator.Enumerator.Next(ref current))
{
result[idx++] = current;
}
return result;
}
return [];
}
Final look at our code
Grouping.cs
public struct Grouping<TKey, TElement> : IDisposable
{
public uint HashCode;
public TKey Key;
public int HashNext;
public int Next;
public TElement[]? Elements;
public int ElementCount;
public int ElementCapacity => Elements.Length;
public TElement this[int index] =>
index > Elements.Length - 1 ? throw new IndexOutOfRangeException() : Elements[index];
public void Dispose()
{
ArrayPool<TElement>.Shared.Return(Elements);
}
}
public readonly struct ReadOnlyGrouping<TKey, TElement>
{
private readonly Grouping<TKey, TElement> _grouping;
public ReadOnlyGrouping(in Grouping<TKey, TElement> grouping)
{
_grouping = grouping;
}
public TKey Key => _grouping.Key;
public int Count => _grouping.ElementCount;
public ReadOnlySpan<TElement> Elements => _grouping.Elements != null
? _grouping.Elements.AsSpan(0, _grouping.ElementCount)
: ReadOnlySpan<TElement>.Empty;
}
Lookup.cs
public ref struct Lookup<TKey, TElement> : IDisposable
{
private const double LoadFactor = 0.75;
private const int InitialCapacity = 16;
private const int InitialElementCapacity = 4;
private Grouping<TKey, TElement>[] _entries;
private int[] _buckets;
private readonly IEqualityComparer<TKey> _comparer;
private int _count;
private int _threshold;
private int _lastGroupIndex;
private uint _bucketMask;
private Lookup(IEqualityComparer<TKey>? comparer = null)
{
_comparer = comparer ?? EqualityComparer<TKey>.Default;
_bucketMask = InitialCapacity - 1;
_threshold = (int)(InitialCapacity * LoadFactor);
_buckets = ArrayPool<int>.Shared.Rent(InitialCapacity);
_entries = ArrayPool<Grouping<TKey, TElement>>.Shared.Rent(InitialCapacity);
_buckets.AsSpan(0, InitialCapacity).Clear();
_count = 0;
_lastGroupIndex = 0;
}
public static Lookup<TKey, TElement> Create<TEnumerator>(TEnumerator source, Func<TElement, TKey> selector,
IEqualityComparer<TKey>? comparer = null) where TEnumerator : IStructEnumerator<TElement>, allows ref struct
{
var lookup = new Lookup<TKey, TElement>(comparer);
ReadOnlySpan<TElement> span = default;
if (source.GetUnderlying(ref span))
{
foreach (var element in span)
{
var key = selector(element);
if (key != null)
{
lookup.Add(key, element);
}
}
}
else
{
TElement element = default;
while (source.Next(ref element))
{
var key = selector(element);
if (key != null)
{
lookup.Add(key, element);
}
}
}
return lookup;
}
public int Count => _count;
public LookupIterator GetIterator()
{
return new LookupIterator(_entries.AsSpan(0, _count));
}
public void Add(TKey key, TElement value)
{
if (_count >= _threshold)
{
//Because we need to resize when threshold hit
Resize();
}
uint hashCode = InternalGetHashCode(ref key);
var bucketIdx = GetIndex(ref hashCode);
int currentEntry = _buckets[bucketIdx];
//Because we need to check if we already have the key in the bucket
while (currentEntry > 0)
{
ref Grouping<TKey, TElement> entry = ref _entries[currentEntry - 1];
//In hashset we would rather stop on the first match and notify the caller
if (entry.HashCode == hashCode && _comparer.Equals(entry.Key, key))
{
AddToGroup(ref entry, value);
return;
}
currentEntry = entry.HashNext;
}
//Because we checked that there is no matching key
CreateGroup(key, hashCode, value, bucketIdx);
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private void AddToGroup(ref Grouping<TKey, TElement> entry, TElement value)
{
//Because every Key has its own group, and we want to resize it if needed
if (entry.ElementCount >= entry.ElementCapacity)
{
int newCapacity = entry.ElementCapacity == 0 ? InitialElementCapacity : entry.ElementCapacity * 2;
var newArray = ArrayPool<TElement>.Shared.Rent(newCapacity);
if (entry.Elements != null)
{
entry.Elements.AsSpan(0, entry.ElementCount).CopyTo(newArray);
ArrayPool<TElement>.Shared.Return(entry.Elements,
clearArray: RuntimeHelpers.IsReferenceOrContainsReferences<TElement>());
}
entry.Elements = newArray;
}
//Because element count changed
entry.Elements[entry.ElementCount++] = value;
}
private void CreateGroup(TKey key, uint hashCode, TElement value, int bucketIndex)
{
if (_count >= _entries.Length)
{
Resize();
//Because we resized the array, we need to get the new bucket index
bucketIndex = GetIndex(ref hashCode);
}
int newIndex = _count + 1;
ref Grouping<TKey, TElement> newEntry = ref _entries[_count];
newEntry.HashCode = hashCode;
newEntry.Key = key;
newEntry.HashNext = _buckets[bucketIndex];
//Because arrays can be huge
newEntry.Elements = ArrayPool<TElement>.Shared.Rent(InitialElementCapacity);
newEntry.Elements[0] = value;
newEntry.ElementCount = 1;
_buckets[bucketIndex] = newIndex;
if (_lastGroupIndex == 0)
{
newEntry.Next = newIndex;
_lastGroupIndex = newIndex;
}
else
{
ref Grouping<TKey, TElement> lastGroup = ref _entries[_lastGroupIndex - 1];
newEntry.Next = lastGroup.Next;
lastGroup.Next = newIndex;
_lastGroupIndex = newIndex;
}
//Because we added a new entry
_count++;
}
//The same as we used in MinimalHashSet
private void Resize()
{
uint newSize = System.Numerics.BitOperations.RoundUpToPowerOf2((uint)_entries.Length * 2);
var newBucket = ArrayPool<int>.Shared.Rent((int)newSize);
var newEntries = ArrayPool<Grouping<TKey, TElement>>.Shared.Rent((int)newSize);
newBucket.AsSpan(0, (int)newSize).Clear();
_entries.AsSpan(0, _count).CopyTo(newEntries);
_bucketMask = newSize - 1;
for (int i = 0; i < _count; i++)
{
ref Grouping<TKey, TElement> entry = ref newEntries[i];
int bucketIndex = GetIndex(ref entry.HashCode);
entry.HashNext = newBucket[bucketIndex];
newBucket[bucketIndex] = i + 1;
}
ArrayPool<int>.Shared.Return(_buckets, clearArray: true);
ArrayPool<Grouping<TKey, TElement>>.Shared.Return(_entries, clearArray: true);
_threshold = (int)(newSize * LoadFactor);
_buckets = newBucket;
_entries = newEntries;
}
private uint InternalGetHashCode(ref TKey item)
{
return (uint)(_comparer.GetHashCode(item) & 0x7FFFFFFF);
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private int GetIndex(ref uint hashCode)
{
return (int)(hashCode & _bucketMask);
}
public readonly void Dispose()
{
if (_buckets != null)
{
ArrayPool<int>.Shared.Return(_buckets);
foreach (var entry in _entries)
{
ArrayPool<TElement>.Shared.Return(entry.Elements);
}
ArrayPool<Grouping<TKey, TElement>>.Shared.Return(_entries);
}
}
public ref struct LookupIterator
{
private int _index;
private readonly int _count;
private readonly ReadOnlySpan<Grouping<TKey, TElement>> _entries;
public LookupIterator(ReadOnlySpan<Grouping<TKey, TElement>> groups)
{
_entries = groups;
_count = groups.Length;
_index = 0;
}
public bool Next()
{
while (++_index <= _count)
{
ref readonly Grouping<TKey, TElement> entry = ref _entries[_index - 1];
if (entry.Elements != null && entry.ElementCount > 0)
{
return true;
}
}
return false;
}
public ReadOnlyGrouping<TKey, TElement> Current
{
get
{
if(_index < 0 || _index > _count)
{
throw new InvalidOperationException();
}
ref readonly var entry = ref _entries[_index - 1];
return new ReadOnlyGrouping<TKey, TElement>(in entry);
}
}
public int Remaining => _count - _index;
}
}
GroupByStruct.cs
public static class GroupByStruct
{
public static StructEnumerable<InternalGroupBy<TEnumerator, TKey, TElement>, ReadOnlyGrouping<TKey, TElement>>
StructGroupBy<TEnumerator, TElement, TKey>(
this StructEnumerable<TEnumerator, TElement> source,
Func<TElement, TKey> keySelector,
IEqualityComparer<TKey>? comparer = null)
where TEnumerator : struct, IStructEnumerator<TElement>, allows ref struct
{
return new(new InternalGroupBy<TEnumerator, TKey, TElement>(source.Enumerator, keySelector, comparer));
}
public static StructEnumerable<InternalGroupBy<ArrayStructEnumerator<TElement>, TKey, TElement>,
ReadOnlyGrouping<TKey, TElement>>
StructGroupBy<TElement, TKey>(
this TElement[] array,
Func<TElement, TKey> keySelector,
IEqualityComparer<TKey>? comparer = null)
{
return new(new InternalGroupBy<ArrayStructEnumerator<TElement>, TKey, TElement>(
new ArrayStructEnumerator<TElement>(array),
keySelector, comparer));
}
}
public ref struct InternalGroupBy<TEnumerator, TKey, TElement> : IStructEnumerator<ReadOnlyGrouping<TKey, TElement>>
where TEnumerator : struct, IStructEnumerator<TElement>, allows ref struct
{
private TEnumerator _source;
private readonly Func<TElement, TKey> _selector;
private readonly IEqualityComparer<TKey> _comparer;
private Base.Lookup<TKey, TElement> _lookup;
private Base.Lookup<TKey,TElement>.LookupIterator _lookupIterator;
private bool _initialized;
public InternalGroupBy(TEnumerator source, Func<TElement, TKey> selector, IEqualityComparer<TKey>? comparer = null)
{
_source = source;
_selector = selector;
_comparer = comparer ?? EqualityComparer<TKey>.Default;
_lookup = default;
_initialized = false;
}
public void Dispose()
{
_source.Dispose();
_lookup.Dispose();
}
public bool Next(ref ReadOnlyGrouping<TKey, TElement> current)
{
Initialize();
if (_lookupIterator.Next())
{
current = _lookupIterator.Current;
return true;
}
current = default;
return false;
}
public bool GetCountToLeftEnumerate(out int count)
{
Initialize();
count = _lookupIterator.Remaining;
return true;
}
public bool GetUnderlying(ref ReadOnlySpan<ReadOnlyGrouping<TKey, TElement>> span)
{
span = default;
return false;
}
public bool TryCopy(ref Span<ReadOnlyGrouping<TKey, TElement>> destination, int offset)
{
return false;
}
private void Initialize()
{
if (!_initialized)
{
_lookup = Base.Lookup<TKey, TElement>.Create(_source, _selector, _comparer);
_lookupIterator = _lookup.GetIterator();
_initialized = true;
}
}
}
Benchmark
[MemoryDiagnoser]
public class GroupByBenchmark
{
public IEnumerable<Example[]> Data()
{
yield return Generator.Generate(500_000);
}
[Benchmark]
[ArgumentsSource(nameof(Data))]
public void LinqGroupBy(Example[] input)
{
input.GroupBy(x => x.ExampleNumber % 3).ToArray().Consume(new Consumer());
}
[Benchmark]
[ArgumentsSource(nameof(Data))]
public void ZLinqGroupBy(Example[] input)
{
input.AsValueEnumerable().GroupBy(x => x.ExampleNumber % 3).ToArray().Consume(new Consumer());
}
[Benchmark]
[ArgumentsSource(nameof(Data))]
public void StructGroupBy(Example[] input)
{
input.StructGroupBy(x => x.ExampleNumber % 3).ToArray().Consume(new Consumer());
}
}
Result
| Method | input | Mean | Error | StdDev | Gen0 | Gen1 | Gen2 | Allocated |
|-------------- |---------------- |---------:|----------:|----------:|---------:|---------:|---------:|----------:|
| LinqGroupBy | Example[500000] | 5.283 ms | 0.1022 ms | 0.1399 ms | 210.9375 | 187.5000 | 187.5000 | 12 MB |
| ZLinqGroupBy | Example[500000] | 4.410 ms | 0.0862 ms | 0.0847 ms | 179.6875 | 156.2500 | 156.2500 | 12 MB |
| StructGroupBy | Example[500000] | 4.112 ms | 0.0818 ms | 0.2213 ms | 101.5625 | 101.5625 | 101.5625 | 6 MB |