simple tree iterator now works

src/karaxdb/btree.nim

@@ -73,17 +73,17 @@ proc `=~`(i: int; k: CmpKind): bool =
   of CmpKind.gt: i > 0
   of CmpKind.neq: i != 0
 
-proc canPrune(i: int; k: CmpKind): bool =
+proc follow(i: int; k: CmpKind): bool =
   case k
   of CmpKind.eq:
-    # we demand equality so if bigger, we can prune:
-    i > 0
+    # we demand equality so if less than, follow it
+    i < 0
   of CmpKind.le:
-    i > 0
-  of CmpKind.lt: i >= 0
-  of CmpKind.ge: i < 0
-  of CmpKind.gt: i <= 0
-  of CmpKind.neq: i != 0
+    i <= 0
+  of CmpKind.lt: i <= 0
+  of CmpKind.ge: i >= 0
+  of CmpKind.gt: i >= 0
+  of CmpKind.neq: true
 
 proc startingPoint(x: Node; key: Key; kind: CmpKind): Cursor =
   if x.isInternal:
@@ -95,22 +95,79 @@ proc startingPoint(x: Node; key: Key; kind: CmpKind): Cursor =
       let cmpRes = cmp(key, x.keys[j])
       if cmpRes =~ kind: return Cursor(n: x, i: j)
 
-proc dos(x: Node; key: Key; kind: CmpKind; withKey: proc(k: Key; v: Val)) =
+proc dos(x: Node; kind: CmpKind; key: Key; withKey: proc(k: Key; v: Val)) =
   if not x.isInternal:
-    echo "came here"
     for j in 0 ..< x.m:
-      if cmp(key, x.keys[j]) =~ kind:
-        withKey(x.keys[j+1], x.vals[j])
+      if cmp(x.keys[j], key) =~ kind:
+        withKey(x.keys[j], x.vals[j])
   else:
-    for j in 0 ..< x.m-1:
-      let cmpRes = cmp(key, x.keys[j+1])
-      echo "came here B ", canPrune(cmpRes, kind)
-      if not canPrune(cmpRes, kind): #cmpRes =~ kind:
-        dos(x.links[j], key, kind, withKey)
-        #if canPrune(cmpRes, kind): return
-      # XXX j+1==x.m case?
+    # we compute the range of links to follow first, before
+    # recursing:
+    var followA = 0
+    var followB = -1
+    case kind
+    of CmpKind.eq:
+      # want: key == 10
+      # keys: 0 3  4 5  10 20
+      # keys: 20 30 40
+      for j in 1..x.m:
+        if j == x.m or cmp(key, x.keys[j]) < 0:
+          followA = j-1
+          followB = j-1
+          break
+    of CmpKind.le, CmpKind.lt:
+      # want: key <= 10  or   key < 10
+      # keys: 0 3  4 5  10 20
+      # keys: 20 30 40
 
+      # Case A: all keys are bigger:
+      if cmp(key, x.keys[1]) < 0:
+        # --> use the very first branch
+        followA = 0
+        followB = 0
+      else:
+        # Case B: all keys are smaller --> use all branches is covered too
+        # by this loop.
+        for j in 1..<x.m:
+          let cmpRes = cmp(key, x.keys[j])
+          if cmpRes >= 0:
+            if followB < 0: followA = j-1
+            # if the keys are identical and we require 'lt', we know
+            # only the left branch is required:
+            followB = j - ord(kind == CmpKind.lt and cmpRes == 0)
+          else:
+            # it's already greater, all others are greater too:
+            break
+    of CmpKind.ge, CmpKind.gt:
+      # want: key >= 10  or  key > 10
+      # keys: 0 3  4 5  10 20
+      # keys: 20 30 40
 
+      # Case A:  all keys are smaller:
+      if cmp(key, x.keys[x.m-1]) >= 0:
+        # --> use the very last branch
+        followA = x.m-1
+        followB = x.m-1
+      else:
+        # also covers case B: all keys are bigger --> use all branches
+        # we find the key that is bigger or equal to ours and from
+        # then on, follow every branch:
+        for j in 1..<x.m:
+          let cmpRes = cmp(key, x.keys[j])
+          if cmpRes <= 0:
+            # if the keys are identical and we need 'ge', we don't have
+            # to consider the 'j-1' branch:
+            followA = j - ord(kind != CmpKind.ge or cmpRes != 0)
+            # we know everything else is even bigger:
+            followB = x.m-1
+            break
+    of CmpKind.neq:
+      # neq: just follow all for now:
+      followA = 0
+      followB = x.m-1
+    # now recurse into the branches that hold candidates we're interested in:
+    for i in followA..followB:
+      dos(x.links[i], kind, key, withKey)
 
 proc init(x: Node): Cursor =
   result.up = @[]
@@ -333,7 +390,7 @@ when isMainModule:
     when true:
       var b1 = newBTree()
       var b2 = newBTree()
-      const iters = 10 #60_000
+      const iters = 9 #60_000
       for i in 1..iters:
         b2 = b2.putPs($i, $(iters - i))
         b1.put($i, $(iters - i))
@@ -343,6 +400,16 @@ when isMainModule:
           echo i, "th iteration; got ", x, ", but expected ", iters - i
       echo b2.n, " = ", b1.n
       echo b2.height, " = ", b1.height
-      dos(b1.root, "5", CmpKind.eq, proc(k: Key; v: Val) = echo("k ", k, " = ", v))
+      echo " >= 5"
+      dos(b1.root, CmpKind.ge, "5", proc(k: Key; v: Val) = echo("k ", k, " = ", v))
+      echo " <= 5"
+      dos(b1.root, CmpKind.le, "5", proc(k: Key; v: Val) = echo("k ", k, " = ", v))
+
+      echo " == 5"
+      dos(b1.root, CmpKind.eq, "5", proc(k: Key; v: Val) = echo("k ", k, " = ", v))
+      echo " < 5"
+      dos(b1.root, CmpKind.lt, "5", proc(k: Key; v: Val) = echo("k ", k, " = ", v))
+      echo " > 5"
+      dos(b1.root, CmpKind.gt, "5", proc(k: Key; v: Val) = echo("k ", k, " = ", v))
 
   main()