Add BTree.get() function

canscale · Nov 6, 2022 · 7375098 · 7375098
1 parent e7defa5
commit 7375098
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 0 deletions.
diff --git a/src/BTree.mo b/src/BTree.mo
@@ -60,6 +60,15 @@ module {
   };
 
 
+  /// Retrieves the value corresponding to the key of BTree if it exists
+  public func get<K, V>(tree: BTree<K, V>, compare: (K, K) -> O.Order, key: K): ?V {
+    switch(tree.root) {
+      case (#internal(internalNode)) { getFromInternal(internalNode, compare, key) };
+      case (#leaf(leafNode)) { getFromLeaf(leafNode, compare, key) }
+    }
+  };
+
+
   /// Inserts an element into a BTree
   public func insert<K, V>(tree: BTree<K, V>, compare: (K, K) -> O.Order, key: K, value: V): ?V {
     switch(tree.root) {
@@ -114,6 +123,38 @@ module {
   };
 
 
+  // get helper if internal node
+  func getFromInternal<K, V>(internalNode: Internal<K, V>, compare: (K, K) -> O.Order, key: K): ?V { 
+    switch(getKeyIndex<K, V>(internalNode.data, compare, key)) {
+      case (#keyFound(index)) { getExistingValueFromIndex(internalNode.data, index) };
+      case (#notFound(index)) {
+        switch(internalNode.children[index]) {
+          // expects the child to be there, otherwise there's a bug in binary search or the tree is invalid
+          case null { assert false; null };
+          case (?#leaf(leafNode)) { getFromLeaf(leafNode, compare, key)};
+          case (?#internal(internalNode)) { getFromInternal(internalNode, compare, key)}
+        }
+      }
+    }
+  };
+
+  // get function helper if leaf node
+  func getFromLeaf<K, V>(leafNode: Leaf<K, V>, compare: (K, K) -> O.Order, key: K): ?V { 
+    switch(getKeyIndex<K, V>(leafNode.data, compare, key)) {
+      case (#keyFound(index)) { getExistingValueFromIndex(leafNode.data, index) };
+      case _ null;
+    }
+  };
+
+  // get function helper that retrieves an existing value in the case that the key is found
+  func getExistingValueFromIndex<K, V>(data: Data<K, V>, index: Nat): ?V {
+    switch(data.kvs[index]) {
+      case null { null };
+      case (?ov) { ?ov.1 }
+    }
+  };
+
+
   // This type is used to signal to the parent calling context what happened in the level below
   type IntermediateInsertResult<K, V> = {
     // element was inserted or replaced, returning the old value (?value or null)

diff --git a/test/BTreeTest.mo b/test/BTreeTest.mo
@@ -49,6 +49,45 @@ let initSuite = S.suite("init", [
   )
 ]);
 
+let getSuite = S.suite("get", [
+  S.test("returns null on an empty BTree",
+    BT.get<Nat, Nat>(BT.init<Nat, Nat>(4), Nat.compare, 5),
+    M.equals(T.optional<Nat>(T.natTestable, null))
+  ),
+  S.test("returns null on a BTree leaf node that does not contain the key",
+    BT.get<Nat, Nat>(quickCreateBTreeWithKVPairs(4, [3, 7]), Nat.compare, 5),
+    M.equals(T.optional<Nat>(T.natTestable, null))
+  ),
+  S.test("returns null on a multi-level BTree that does not contain the key",
+    BT.get<Nat, Nat>(
+      quickCreateBTreeWithKVPairs(4, [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150, 160]),
+      Nat.compare,
+      21
+    ),
+    M.equals(T.optional<Nat>(T.natTestable, null))
+  ),
+  S.test("returns null on a multi-level BTree that does not contain the key, if the key is greater than all elements in the tree",
+    BT.get<Nat, Nat>(
+      quickCreateBTreeWithKVPairs(4, [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150, 160]),
+      Nat.compare,
+      200
+    ),
+    M.equals(T.optional<Nat>(T.natTestable, null))
+  ),
+  S.test("returns the value if a BTree leaf node contains the key",
+    BT.get<Nat, Nat>(quickCreateBTreeWithKVPairs(4, [3, 7, 10]), Nat.compare, 10),
+    M.equals(T.optional<Nat>(T.natTestable, ?10))
+  ),
+  S.test("returns the value if a BTree internal node contains the key",
+    BT.get<Nat, Nat>(
+      quickCreateBTreeWithKVPairs(4, [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150, 160]),
+      Nat.compare,
+      120
+    ),
+    M.equals(T.optional<Nat>(T.natTestable, ?120))
+  ),
+]);
+
 
 let insertSuite = S.suite("insert", [
   S.suite("root as leaf tests", [
@@ -761,6 +800,7 @@ let insertSuite = S.suite("insert", [
 S.run(S.suite("BTree",
   [
     initSuite,
+    getSuite,
     insertSuite
   ]
 ));