LINQ GroupBy Grouping can use InlineArray for first elements to avoid allocations

src/libraries/System.Linq/src/System/Linq/Grouping.cs

@@ -4,6 +4,7 @@
 using System.Collections;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Runtime.CompilerServices;
 
 namespace System.Linq
 {
@@ -173,8 +174,7 @@ public override bool MoveNext()
             ValidItem:
                 _g = _g._next;
                 Debug.Assert(_g is not null);
-                _g.Trim();
-                _current = _resultSelector(_g.Key, _g._elements);
+                _current = _resultSelector(_g.Key, _g);
                 return true;
             }
         }
@@ -229,8 +229,7 @@ public override bool MoveNext()
             ValidItem:
                 _g = _g._next;
                 Debug.Assert(_g is not null);
-                _g.Trim();
-                _current = _resultSelector(_g.Key, _g._elements);
+                _current = _resultSelector(_g.Key, _g);
                 return true;
             }
         }
@@ -335,7 +334,7 @@ public override bool MoveNext()
                 Dispose();
                 return false;
 
-                ValidItem:
+            ValidItem:
                 _g = _g._next;
                 Debug.Assert(_g is not null);
                 _current = _g;
@@ -355,41 +354,57 @@ internal sealed class Grouping<TKey, TElement> : IGrouping<TKey, TElement>, ILis
     {
         internal readonly TKey _key;
         internal readonly int _hashCode;
-        internal TElement[] _elements;
-        internal int _count;
+        private GroupingElementArray<TElement> _inlineElements;
+        private TElement[]? _elements;
+        private int _count;
         internal Grouping<TKey, TElement>? _hashNext;
         internal Grouping<TKey, TElement>? _next;
 
         internal Grouping(TKey key, int hashCode)
         {
             _key = key;
             _hashCode = hashCode;
-            _elements = new TElement[1];
         }
 
         internal void Add(TElement element)
         {
-            if (_elements.Length == _count)
-            {
-                Array.Resize(ref _elements, checked(_count * 2));
-            }
+            Span<TElement> destination;
 
-            _elements[_count] = element;
-            _count++;
-        }
+            if (_elements is null)
+            {
+                destination = _inlineElements;
 
-        internal void Trim()
-        {
-            if (_elements.Length != _count)
+                if (_count == GroupingElementArray<TElement>.Size)
+                {
+                    _elements = new TElement[checked(_count * 2)];
+                    destination.CopyTo(_elements);
+                    destination = _elements;
+                }
+            }
+            else
             {
-                Array.Resize(ref _elements, _count);
+                if (_elements.Length == _count)
+                {
+                    Array.Resize(ref _elements, checked(_count * 2));
+                }
+
+                destination = _elements;
             }
+
+            destination[_count++] = element;
         }
 
         public IEnumerator<TElement> GetEnumerator()
         {
             Debug.Assert(_count > 0, "A grouping should only have been created if an element was being added to it.");
-            return new PartialArrayEnumerator<TElement>(_elements, _count);
+            if (_elements is null)
+            {
+                return new GroupingElementArrayEnumerator<TElement>(_inlineElements, _count);
+            }
+            else
+            {
+                return new PartialArrayEnumerator<TElement>(_elements, _count);
+            }
         }
 
         IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
@@ -404,14 +419,29 @@ public IEnumerator<TElement> GetEnumerator()
 
         void ICollection<TElement>.Clear() => ThrowHelper.ThrowNotSupportedException();
 
-        bool ICollection<TElement>.Contains(TElement item) => Array.IndexOf(_elements, item, 0, _count) >= 0;
+        bool ICollection<TElement>.Contains(TElement item)
+        {
+            if (_elements is null)
+            {
+                return IndexOfInlineElements(item) >= 0;
+            }
+
+            return Array.IndexOf(_elements, item) >= 0;
+        }
 
-        void ICollection<TElement>.CopyTo(TElement[] array, int arrayIndex) =>
-            Array.Copy(_elements, 0, array, arrayIndex, _count);
+        void ICollection<TElement>.CopyTo(TElement[] array, int arrayIndex) => GetElements().CopyTo(array.AsSpan(arrayIndex));
 
         bool ICollection<TElement>.Remove(TElement item) => ThrowHelper.ThrowNotSupportedException_Boolean();
 
-        int IList<TElement>.IndexOf(TElement item) => Array.IndexOf(_elements, item, 0, _count);
+        int IList<TElement>.IndexOf(TElement item)
+        {
+            if (_elements is null)
+            {
+                return IndexOfInlineElements(item);
+            }
+
+            return Array.IndexOf(_elements, item, 0, _count);
+        }
 
         void IList<TElement>.Insert(int index, TElement item) => ThrowHelper.ThrowNotSupportedException();
 
@@ -426,10 +456,74 @@ TElement IList<TElement>.this[int index]
                     ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.index);
                 }
 
-                return _elements[index];
+                return GetElements()[index];
             }
 
             set => ThrowHelper.ThrowNotSupportedException();
         }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal ReadOnlySpan<TElement> GetElements()
+        {
+            ReadOnlySpan<TElement> buffer = _elements is null ? _inlineElements : _elements;
+            return buffer.Slice(0, _count);
+        }
+
+        private int IndexOfInlineElements(TElement item)
+        {
+            ReadOnlySpan<TElement> inlineElements = ((ReadOnlySpan<TElement>)_inlineElements).Slice(0, _count);
+            for (int i = 0; i < inlineElements.Length; i++)
+            {
+                if (EqualityComparer<TElement>.Default.Equals(inlineElements[i], item))
+                {
+                    return i;
+                }
+            }
+
+            return -1;
+        }
+
+    }
+
+    [InlineArray(GroupingElementArray<int>.Size)]
+    internal struct GroupingElementArray<TElement>
+    {
+        public const int Size = 3;
+        private TElement _element0;
+    }
+
+    internal sealed class GroupingElementArrayEnumerator<TElement> : IEnumerator<TElement>
+    {
+        private GroupingElementArray<TElement> _array;
+        private int _count;
+        private int _index;
+
+        public GroupingElementArrayEnumerator(GroupingElementArray<TElement> array, int count)
+        {
+            Debug.Assert((uint)_count <= GroupingElementArray<TElement>.Size);
+            _array = array;
+            _count = count;
+            _index = -1;
+        }
+
+        public bool MoveNext()
+        {
+            if (_index + 1 < _count)
+            {
+                _index++;
+                return true;
+            }
+
+            return false;
+        }
+
+        public TElement Current => _array[_index];
+
+        object? IEnumerator.Current => Current;
+
+        public void Dispose() { }
+
+        public void Reset() => _index = -1;
     }
+
 }

src/libraries/System.Linq/src/System/Linq/Join.cs

@@ -261,11 +261,10 @@ private static IEnumerable<TResult> JoinIterator<TOuter, TInner, TKey, TResult>(
                         Grouping<TKey, TInner>? g = lookup.GetGrouping(outerKeySelector(item), create: false);
                         if (g is not null)
                         {
-                            int count = g._count;
-                            TInner[] elements = g._elements;
+                            int count = g.GetElements().Length;
                             for (int i = 0; i != count; ++i)
                             {
-                                yield return resultSelector(item, elements[i]);
+                                yield return resultSelector(item, g.GetElements()[i]);
                             }
                         }
                     } while (e.MoveNext());

src/libraries/System.Linq/src/System/Linq/LeftJoin.cs

@@ -261,11 +261,10 @@ private static IEnumerable<TResult> LeftJoinIterator<TOuter, TInner, TKey, TResu
                     }
                     else
                     {
-                        int count = g._count;
-                        TInner[] elements = g._elements;
+                        int count = g.GetElements().Length;
                         for (int i = 0; i != count; ++i)
                         {
-                            yield return resultSelector(item, elements[i]);
+                            yield return resultSelector(item, g.GetElements()[i]);
                         }
                     }
                 }

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

@@ -20,8 +20,7 @@ internal TResult[] ToArray<TResult>(Func<TKey, IEnumerable<TElement>, TResult> r
                     g = g._next;
                     Debug.Assert(g is not null);
 
-                    g.Trim();
-                    array[index] = resultSelector(g._key, g._elements);
+                    array[index] = resultSelector(g._key, g);
                     ++index;
                 }
                 while (g != _lastGrouping);

src/libraries/System.Linq/src/System/Linq/Lookup.cs

@@ -164,8 +164,7 @@ internal List<TResult> ToList<TResult>(Func<TKey, IEnumerable<TElement>, TResult
                     g = g._next;
 
                     Debug.Assert(g is not null);
-                    g.Trim();
-                    span[index] = resultSelector(g._key, g._elements);
+                    span[index] = resultSelector(g._key, g);
                     ++index;
                 }
                 while (g != _lastGrouping);
@@ -186,8 +185,7 @@ public IEnumerable<TResult> ApplyResultSelector<TResult>(Func<TKey, IEnumerable<
                     g = g._next;
 
                     Debug.Assert(g is not null);
-                    g.Trim();
-                    yield return resultSelector(g._key, g._elements);
+                    yield return resultSelector(g._key, g);
                 }
                 while (g != _lastGrouping);
             }

src/libraries/System.Linq/src/System/Linq/RightJoin.cs

@@ -259,11 +259,10 @@ private static IEnumerable<TResult> RightJoinIterator<TOuter, TInner, TKey, TRes
                     }
                     else
                     {
-                        int count = g._count;
-                        TOuter[] elements = g._elements;
+                        int count = g.GetElements().Length;
                         for (int i = 0; i != count; ++i)
                         {
-                            yield return resultSelector(elements[i], item);
+                            yield return resultSelector(g.GetElements()[i], item);
                         }
                     }
                 }