Slotted enumeration faster

src/Paprika.Benchmarks/SlottedArrayBenchmarks.cs

@@ -26,7 +26,6 @@ private const int
     private readonly void* _hashCollidingKeys;
     private readonly void* _hashCollidingMap;
 
-
     public SlottedArrayBenchmarks()
     {
         // Create keys
@@ -186,9 +185,25 @@ public int EnumerateAll()
         return length;
     }
 
-    private NibblePath GetKey(byte i, bool odd)
+    [Benchmark]
+    [Arguments((byte)0)]
+    [Arguments((byte)1)]
+    public int EnumerateNibble(byte nibble)
     {
+        var map = new SlottedArray(new Span<byte>(_map, Page.PageSize));
+
+        var length = 0;
+        foreach (var item in map.EnumerateNibble(nibble))
+        {
+            length += item.Key.Length;
+            length += item.RawData.Length;
+        }
+
+        return length;
+    }
 
+    private NibblePath GetKey(byte i, bool odd)
+    {
         var span = new Span<byte>(_keys, BytesPerKey * KeyCount);
         var slice = span.Slice(i * BytesPerKey, BytesPerKey);
 

src/Paprika.Tests/Data/SlottedArrayTests.cs

@@ -60,6 +60,7 @@ public Task Enumerate_all([Values(0, 1)] int odd)
         map.GetAssert(key0, Data0);
         map.GetAssert(key1, Data1);
         map.GetAssert(key2, Data2);
+        map.GetAssert(key3, Data3);
         map.GetAssert(key4, Data4);
 
         using var e = map.EnumerateAll();
@@ -90,6 +91,61 @@ public Task Enumerate_all([Values(0, 1)] int odd)
         return Verify(span.ToArray());
     }
 
+    [Test]
+    public Task Enumerate_nibble([Values(1, 2, 3, 4)] int nibble)
+    {
+        Span<byte> span = stackalloc byte[256];
+        var map = new SlottedArray(span);
+
+        var key0 = NibblePath.Empty;
+        var key1 = NibblePath.FromKey(stackalloc byte[1] { 0x1A });
+        var key2 = NibblePath.FromKey(stackalloc byte[2] { 0x2A, 13 });
+        var key3 = NibblePath.FromKey(stackalloc byte[3] { 0x3A, 13, 31 });
+        var key4 = NibblePath.FromKey(stackalloc byte[4] { 0x4A, 13, 31, 41 });
+
+        map.SetAssert(key0, Data0);
+        map.SetAssert(key1, Data1);
+        map.SetAssert(key2, Data2);
+        map.SetAssert(key3, Data3);
+        map.SetAssert(key4, Data4);
+
+        map.GetAssert(key0, Data0);
+        map.GetAssert(key1, Data1);
+        map.GetAssert(key2, Data2);
+        map.GetAssert(key3, Data3);
+        map.GetAssert(key4, Data4);
+
+        var expected = nibble switch
+        {
+            1 => key1,
+            2 => key2,
+            3 => key3,
+            4 => key4,
+            _ => throw new Exception()
+        };
+
+        var data = nibble switch
+        {
+            1 => Data1,
+            2 => Data2,
+            3 => Data3,
+            4 => Data4,
+            _ => throw new Exception()
+        };
+
+        using var e = map.EnumerateNibble((byte)nibble);
+
+        e.MoveNext().Should().BeTrue();
+
+        e.Current.Key.Equals(expected).Should().BeTrue();
+        e.Current.RawData.SequenceEqual(data).Should().BeTrue();
+
+        e.MoveNext().Should().BeFalse();
+
+        // verify
+        return Verify(span.ToArray());
+    }
+
     [Test]
     public Task Enumerate_long_key([Values(0, 1)] int oddStart, [Values(0, 1)] int lengthCutOff)
     {

src/Paprika/Data/SlottedArray.cs

@@ -158,21 +158,21 @@ private bool TrySetImpl(ushort hash, byte preamble, in NibblePath trimmed, ReadO
 
     public int CapacityLeft => _data.Length - _header.Taken;
 
-    public Enumerator EnumerateAll() =>
-        new(this);
+    public Enumerator EnumerateAll() => new(this);
+    public NibbleEnumerator EnumerateNibble(byte nibble) => new(this, nibble);
 
-    public ref struct Enumerator
+    [StructLayout(LayoutKind.Sequential, Pack = sizeof(byte), Size = Size)]
+    private ref struct Chunk
     {
-        [StructLayout(LayoutKind.Sequential, Pack = sizeof(byte), Size = Size)]
-        private ref struct Chunk
-        {
-            public const int Size = 64;
+        private const int Size = 64;
 
-            private byte _start;
+        private byte _start;
 
-            public Span<byte> Span => MemoryMarshal.CreateSpan(ref _start, Size);
-        }
+        public Span<byte> Span => MemoryMarshal.CreateSpan(ref _start, Size);
+    }
 
+    public ref struct Enumerator
+    {
         /// <summary>The map being enumerated.</summary>
         private readonly SlottedArray _map;
 
@@ -192,7 +192,7 @@ internal Enumerator(SlottedArray map)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public bool MoveNext()
         {
-            int index = _index + 1;
+            var index = _index + 1;
             var to = _map.Count;
 
             ref var slot = ref _map.GetSlotRef(index);
@@ -207,7 +207,7 @@ public bool MoveNext()
             if (index < to)
             {
                 _index = index;
-                Build(out _current);
+                Build(slot, out _current);
                 return true;
             }
 
@@ -216,12 +216,10 @@ public bool MoveNext()
 
         public readonly Item Current => _current;
 
-        private void Build(out Item value)
+        private void Build(Slot slot, out Item value)
         {
-            ref var slot = ref _map.GetSlotRef(_index);
             var hash = _map.GetHashRef(_index);
-
-            var span = _map.GetSlotPayload(_index);
+            var span = _map.GetSlotPayload(_index, slot);
             var key = Slot.UnPrepareKey(hash, slot.KeyPreamble, span, _bytes.Span, out var data);
 
             value = new Item(key, data, _index);
@@ -231,15 +229,84 @@ public readonly void Dispose()
         {
         }
 
-        public readonly ref struct Item(NibblePath key, ReadOnlySpan<byte> rawData, int index)
+        // a shortcut to not allocate, just copy the enumerator
+        public readonly Enumerator GetEnumerator() => this;
+    }
+
+    public ref struct NibbleEnumerator
+    {
+        /// <summary>The map being enumerated.</summary>
+        private readonly SlottedArray _map;
+
+        private readonly byte _nibble;
+
+        /// <summary>The next index to yield.</summary>
+        private int _index;
+
+        private Chunk _bytes;
+        private Item _current;
+
+        internal NibbleEnumerator(SlottedArray map, byte nibble)
+        {
+            _map = map;
+            _nibble = nibble;
+            _index = -1;
+        }
+
+        /// <summary>Advances the enumerator to the next element of the span.</summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public bool MoveNext()
+        {
+            var index = _index + 1;
+            var to = _map.Count;
+
+            var slot = _map.GetSlotRef(index);
+            var hash = _map.GetHashRef(index);
+
+            while (index < to && (slot.IsDeleted || slot.HasAtLeastOneNibble == false || slot.GetNibble0(hash) != _nibble)) // filter out deleted
+            {
+                // move by 1
+                index += 1;
+                slot = _map.GetSlotRef(index);
+                hash = _map.GetHashRef(index);
+            }
+
+            if (index < to)
+            {
+                _index = index;
+                Build(slot, hash, out _current);
+                return true;
+            }
+
+            return false;
+        }
+
+        public readonly Item Current => _current;
+
+        private void Build(Slot slot, ushort hash, out Item value)
+        {
+            var span = _map.GetSlotPayload(_index, slot);
+            var key = Slot.UnPrepareKey(hash, slot.KeyPreamble, span, _bytes.Span, out var data);
+
+            value = new Item(key, data, _index);
+        }
+
+        public readonly void Dispose()
         {
-            public int Index { get; } = index;
-            public NibblePath Key { get; } = key;
-            public ReadOnlySpan<byte> RawData { get; } = rawData;
         }
 
         // a shortcut to not allocate, just copy the enumerator
-        public readonly Enumerator GetEnumerator() => this;
+        public readonly NibbleEnumerator GetEnumerator() => this;
+    }
+
+    /// <summary>
+    /// An enumerator item.
+    /// </summary>
+    public readonly ref struct Item(NibblePath key, ReadOnlySpan<byte> rawData, int index)
+    {
+        public int Index { get; } = index;
+        public NibblePath Key { get; } = key;
+        public ReadOnlySpan<byte> RawData { get; } = rawData;
     }
 
     public void MoveNonEmptyKeysTo(in MapSource destination, bool treatEmptyAsTombstone = false)
@@ -350,7 +417,7 @@ public bool Delete(in NibblePath key)
         return false;
     }
 
-    public void Delete(in Enumerator.Item item) => DeleteImpl(item.Index);
+    public void Delete(in Item item) => DeleteImpl(item.Index);
 
     private void DeleteImpl(int index, bool collectTombstones = true)
     {
@@ -626,9 +693,12 @@ private int TryFind(int at, uint matches, in NibblePath key, byte preamble, out
     /// Gets the payload pointed to by the given slot without the length prefix.
     /// </summary>
     [SkipLocalsInit]
-    private Span<byte> GetSlotPayload(int index)
+    private Span<byte> GetSlotPayload(int index) => GetSlotPayload(index, GetSlotRef(index));
+
+    [SkipLocalsInit]
+    private Span<byte> GetSlotPayload(int index, Slot slot)
     {
-        var addr = GetSlotRef(index).ItemAddress;
+        var addr = slot.ItemAddress;
 
         // If this is the first, just slice of data
         if (index == 0)

src/Paprika/Store/DataPage.cs

@@ -174,16 +174,9 @@ private void TryFlushDownToExistingChildren(in SlottedArray map, IBatchContext b
 
     private static Page FlushDown(in SlottedArray map, byte nibble, Page destination, IBatchContext batch)
     {
-        foreach (var item in map.EnumerateAll())
+        foreach (var item in map.EnumerateNibble(nibble))
         {
-            var key = item.Key;
-            if (key.IsEmpty) // empty keys are left in page
-                continue;
-
-            if (key.FirstNibble != nibble)
-                continue;
-
-            var sliced = key.SliceFrom(ConsumedNibbles);
+            var sliced = item.Key.SliceFrom(ConsumedNibbles);
 
             destination = destination.Header.PageType == PageType.Leaf
                 ? new LeafPage(destination).Set(sliced, item.RawData, batch)