BTree: restructured for memory layout efficiency

src/karaxdb/btree.nim

@@ -2,25 +2,26 @@
 # max children per B-tree node = M-1
 # (must be even and greater than 2)
 const
-   M = 4
+  M = 4
-   Mhalf = M div 2
+  Mhalf = M div 2
 
 type
   Key = string
   Val = string
-  Entry = ref object
-    key: Key
-    val: Val # external nodes only
-    next: Node   # internal nodes only; helper field to iterate over array entries
   Node = ref object
     m: int
-    children: array[M, Entry]
+    keys: array[M, Key]
+    case isInternal: bool
+    of false:
+      vals: array[M, Val]
+    of true:
+      links: array[M, Node]
   BTree = ref object
     root: Node
     height: int ## height
     n: int      ## number of key-value pairs
 
-proc newBTree(): BTree = BTree(root: Node(m: 0))
+proc newBTree(): BTree = BTree(root: Node(m: 0, isInternal: false))
 
 proc less(a, b: Key): bool = cmp(a, b) < 0
 
@@ -28,63 +29,61 @@ proc eq(a, b: Key): bool = cmp(a, b) == 0
 
 proc search(x: Node, key: Key, ht: int): Val =
   if ht == 0:
-    # external node
+    assert(not x.isInternal)
     for j in 0 ..< x.m:
-      if eq(key, x.children[j].key): return x.children[j].val
+      if eq(key, x.keys[j]): return x.vals[j]
   else:
-    # internal node
+    assert(x.isInternal)
     for j in 0 ..< x.m:
-      if j+1 == x.m or less(key, x.children[j+1].key):
+      if j+1 == x.m or less(key, x.keys[j+1]):
-        return search(x.children[j].next, key, ht-1)
+        return search(x.links[j], key, ht-1)
-  return nil
-
-proc candidates(x: Node, key: Key, ht: int): Node =
-  if ht == 0:
-    return x
-  else:
-    # internal node
-    for j in 0 ..< x.m:
-      if j+1 == x.m or less(key, x.children[j+1].key):
-        return candidates(x.children[j].next, key, ht-1)
-  return nil
-
-iterator allValues(t: BTree; key: Key): Val =
-  let x = candidates(t.root, key, t.height)
-  if x != nil:
-    for j in 0 ..< x.m:
-      if eq(key, x.children[j].key): yield x.children[j].val
 
 proc get(t: BTree; key: Key): Val = search(t.root, key, t.height)
 
 proc split(h: Node): Node =
   ## split node in half
-  result = Node(m: Mhalf)
+  result = Node(m: Mhalf, isInternal: h.isInternal)
   h.m = Mhalf
   for j in 0 ..< Mhalf:
-    result.children[j] = h.children[Mhalf + j]
+    result.keys[j] = h.keys[Mhalf + j]
+  if h.isInternal:
+    for j in 0 ..< Mhalf:
+      result.links[j] = h.links[Mhalf + j]
+  else:
+    for j in 0 ..< Mhalf:
+      shallowCopy(result.vals[j], h.vals[Mhalf + j])
 
 proc insert(h: Node, key: Key, val: Val, ht: int): Node =
-  var t = Entry(key: key, val: val, next: nil)
+  #var t = Entry(key: key, val: val, next: nil)
+  var newKey = key
   var j = 0
   if ht == 0:
-    # external node:
+    assert(not h.isInternal)
     while j < h.m:
-      if less(key, h.children[j].key): break
+      if less(key, h.keys[j]): break
       inc j
+    for i in countdown(h.m, j+1):
+      shallowCopy(h.vals[i], h.vals[i-1])
+    h.vals[j] = val
   else:
-    # internal node
+    assert h.isInternal
+    var newLink: Node = nil
     while j < h.m:
-      if j+1 == h.m or less(key, h.children[j+1].key):
+      if j+1 == h.m or less(key, h.keys[j+1]):
-        let u = insert(h.children[j].next, key, val, ht-1)
+        let u = insert(h.links[j], key, val, ht-1)
         inc j
         if u == nil: return nil
-        t.key = u.children[0].key
+        newKey = u.keys[0]
-        t.next = u
+        newLink = u
         break
       inc j
+    for i in countdown(h.m, j+1):
+      h.links[i] = h.links[i-1]
+    h.links[j] = newLink
+
   for i in countdown(h.m, j+1):
-    h.children[i] = h.children[i-1]
+    h.keys[i] = h.keys[i-1]
-  h.children[j] = t
+  h.keys[j] = newKey
   inc h.m
   return if h.m < M: nil else: split(h)
 
@@ -94,21 +93,25 @@ proc put(b: BTree; key: Key; val: Val) =
   if u == nil: return
 
   # need to split root
-  let t = Node(m: 2)
+  let t = Node(m: 2, isInternal: true)
-  t.children[0] = Entry(key: b.root.children[0].key, val: nil, next: b.root)
+  t.keys[0] = b.root.keys[0]
-  t.children[1] = Entry(key: u.children[0].key, val: nil, next: u)
+  t.links[0] = b.root
+  t.keys[1] = u.keys[0]
+  t.links[1] = u
   b.root = t
   inc b.height
 
 proc toString(h: Node, ht: int, indent: string; result: var string) =
   if ht == 0:
+    assert(not h.isInternal)
     for j in 0..<h.m:
       result.add(indent)
-      result.add($h.children[j].key & " " & $h.children[j].val & "\n")
+      result.add($h.keys[j] & " " & $h.vals[j] & "\n")
   else:
+    assert(h.isInternal)
     for j in 0..<h.m:
-      if j > 0: result.add(indent & "(" & $h.children[j].key & ")\n")
+      if j > 0: result.add(indent & "(" & $h.keys[j] & ")\n")
-      toString(h.children[j].next, ht-1, indent & "   ", result)
+      toString(h.links[j], ht-1, indent & "   ", result)
 
 proc `$`(b: BTree): string =
   result = ""
@@ -116,8 +119,8 @@ proc `$`(b: BTree): string =
 
 proc main =
   var st = newBTree()
-  st.put("www.cs.princeton.edu", "128.112.136.12")
+  st.put("www.cs.princeton.edu", "abc")
-  st.put("www.cs.princeton.edu", "128.112.136.11")
+  st.put("www.cs.princeton.edu", "xyz")
   st.put("www.princeton.edu",    "128.112.128.15")
   st.put("www.yale.edu",         "130.132.143.21")
   st.put("www.simpsons.com",     "209.052.165.60")
@@ -134,23 +137,26 @@ proc main =
   st.put("www.weather.com",      "63.111.66.11")
   st.put("www.yahoo.com",        "216.109.118.65")
 
-  echo("cs.princeton.edu:  ", st.get("www.cs.princeton.edu"))
+  assert st.get("www.cs.princeton.edu") == "abc"
-  echo("hardvardsucks.com: ", st.get("www.harvardsucks.com"))
+  assert st.get("www.harvardsucks.com") == nil
-  echo("simpsons.com:      ", st.get("www.simpsons.com"))
-  echo("apple.com:         ", st.get("www.apple.com"))
-  echo("ebay.com:          ", st.get("www.ebay.com"))
-  echo("dell.com:          ", st.get("www.dell.com"))
-  echo()
-  echo("size:    ", st.n)
-  echo("height:  ", st.height)
-  echo(st)
 
-  var dups = newBTree()
+  assert st.get("www.simpsons.com") == "209.052.165.60"
-  for i in 0..20:
+  assert st.get("www.apple.com") == "17.112.152.32"
-    dups.put("testme", $i)
+  assert st.get("www.ebay.com") == "66.135.192.87"
-  echo dups
+  assert st.get("www.dell.com") == "143.166.224.230"
-  for v in allValues(dups, "testme"):
+  assert(st.n == 17)
-    echo v
+  assert(st.height == 2)
-  echo dups.get("testme")
+
+  when true:
+    var b2 = newBTree()
+    const iters = 10_000
+    for i in 1..iters:
+      b2.put($i, $(iters - i))
+    for i in 1..iters:
+      let x = b2.get($i)
+      if x != $(iters - i):
+        echo "got ", x, ", but expected ", iters - i
+    echo b2.n
+    echo b2.height
 
 main()