From a6e8ced7f7402c43e3bf79881ab7392b58176619 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Linus=20Bj=C3=B6rnstam?= <linus.internet@fastmail.se>
Date: Tue, 28 Apr 2026 21:05:43 +0200
Subject: [PATCH] Added int avx dispatch in internal nodes

---
 PersistentMap/BTreeFunctions.cs           | 12 +++-
 PersistentMap/KeyStrategies/IntScanner.cs | 78 +++++++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/PersistentMap/BTreeFunctions.cs b/PersistentMap/BTreeFunctions.cs
index 51e4b08..ca199fc 100644
--- a/PersistentMap/BTreeFunctions.cs
+++ b/PersistentMap/BTreeFunctions.cs
@@ -210,7 +210,17 @@ namespace PersistentMap
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal static int FindRoutingIndex<K, TStrategy>(InternalNode<K> node, K key, long keyPrefix, TStrategy strategy)
             where TStrategy : IKeyStrategy<K>
-        {
+        {   
+
+if (typeof(K) == typeof(int))
+            {
+                Span<K> keys = node.GetKeys();
+                ref K firstKeyRef = ref MemoryMarshal.GetReference(keys);
+                ref int firstIntRef = ref Unsafe.As<K, int>(ref firstKeyRef);
+                ReadOnlySpan<int> intKeys = MemoryMarshal.CreateReadOnlySpan(ref firstIntRef, keys.Length);
+                int intKey = Unsafe.As<K, int>(ref key);
+                return IntScanner.FindFirstGreater(intKeys, intKey);
+            }
             if (!strategy.UsesPrefixes)
             {
                 return FallbackRoutingKeys(node.GetKeys(), key, strategy);
diff --git a/PersistentMap/KeyStrategies/IntScanner.cs b/PersistentMap/KeyStrategies/IntScanner.cs
index de7b691..3c19856 100644
--- a/PersistentMap/KeyStrategies/IntScanner.cs
+++ b/PersistentMap/KeyStrategies/IntScanner.cs
@@ -87,4 +87,82 @@ public static class IntScanner
 
         return LinearScan(keys.Slice(i), target) + i;
     }
+
+ [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static int FindFirstGreater(ReadOnlySpan<int> keys, int target)
+    {
+        if (!Avx2.IsSupported || keys.Length < 8)
+            return LinearScanGreater(keys, target);
+
+        return Avx512F.IsSupported
+            ? ScanAvx512Greater(keys, target)
+            : ScanAvx2Greater(keys, target);
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private static int LinearScanGreater(ReadOnlySpan<int> keys, int target)
+    {
+        for (var i = 0; i < keys.Length; i++)
+            if (keys[i] > target)
+                return i;
+        return keys.Length;
+    }
+ 
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private static unsafe int ScanAvx2Greater(ReadOnlySpan<int> keys, int target)
+    {
+        // For > target, AVX2 CompareGreaterThan works directly without the (target - 1) offset
+        var vTarget = Vector256.Create(target);
+        var i = 0;
+        var len = keys.Length;
+
+        for (; i <= len - 8; i += 8)
+        {
+            fixed (int* ptr = keys)
+            {
+                var vData = Avx2.LoadVector256(ptr + i);
+                var vResult = Avx2.CompareGreaterThan(vData, vTarget);
+                
+                var mask = (uint)Avx2.MoveMask(vResult.AsByte());
+
+                if (mask != 0)
+                {
+                    return i + (BitOperations.TrailingZeroCount(mask) / 4);
+                }
+            }
+        }
+
+        return LinearScanGreater(keys.Slice(i), target) + i;
+    }
+
+[MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private static unsafe int ScanAvx512Greater(ReadOnlySpan<int> keys, int target)
+    {
+        var vTarget = Vector512.Create(target);
+        var i = 0;
+        var len = keys.Length;
+
+        for (; i <= len - 16; i += 16)
+        {
+            fixed (int* ptr = keys)
+            {
+                var vData = Avx512F.LoadVector512(ptr + i);
+                
+                // Use GreaterThan instead of GreaterThanOrEqual
+                var mask = Vector512.GreaterThan(vData, vTarget);
+
+                if (mask != Vector512<int>.Zero)
+                {
+                    uint m = (uint)mask.ExtractMostSignificantBits();
+                    return i + BitOperations.TrailingZeroCount(m);
+                }
+            }
+        }
+
+        return LinearScanGreater(keys.Slice(i), target) + i;
+    }
 }
+
+
+
+