Revert getIndex() and use getExpectedIndex() to handle cases where the key is not in the tree

Author: tomijaga

src/MemoryBTree/Base.mo

@@ -11,6 +11,7 @@ import Blob "mo:base/Blob";
 
 import MemoryRegion "mo:memory-region/MemoryRegion";
 import RevIter "mo:itertools/RevIter";
+import Find "mo:map/Map/modules/find";
 
 import MemoryCmp "../TypeUtils/MemoryCmp";
 import Blobify "../TypeUtils/Blobify";
@@ -900,9 +901,35 @@ module {
         };
 
         if (int_index < 0) {
-            index_pos + Int.abs(int_index + 1);
+            Debug.trap("getIndex(): Key not found. Use getExpectedIndex() instead to get keys that might not be in the tree");
+        };
+
+        index_pos + Int.abs(int_index);
+    };
+
+    /// Returns the index of the given key in the tree.
+    /// Used to indicate whether the key is in the tree or not.
+    public type ExpectedIndex = {
+        #Found : Nat;
+        #NotFound : Nat;
+    };
+
+    public func getExpectedIndex<K, V>(btree : MemoryBTree, btree_utils : BTreeUtils<K, V>, key : K) : ExpectedIndex {
+        let key_blob = btree_utils.key.blobify.to_blob(key);
+        let (leaf_address, index_pos) = Methods.get_leaf_node_and_index(btree, btree_utils, key_blob);
+
+        let count = Leaf.get_count(btree, leaf_address);
+        let int_index = switch (btree_utils.key.cmp) {
+            case (#GenCmp(cmp)) Leaf.binary_search<K, V>(btree, btree_utils, leaf_address, cmp, key, count);
+            case (#BlobCmp(cmp)) {
+                Leaf.binary_search_blob_seq(btree, leaf_address, cmp, key_blob, count);
+            };
+        };
+
+        if (int_index < 0) {
+            #NotFound(index_pos + Int.abs(int_index + 1));
         } else {
-            index_pos + Int.abs(int_index);
+            #Found(index_pos + Int.abs(int_index));
         };
 
     };

src/MemoryBTree/Stable.mo

@@ -106,6 +106,13 @@ module StableMemoryBTree {
         MemoryBTree.getIndex(state, btree_utils, key);
     };
 
+    public type ExpectedIndex = MemoryBTree.ExpectedIndex;
+
+    public func getExpectedIndex<K, V>(btree : StableMemoryBTree, btree_utils : BTreeUtils<K, V>, key : K) : ExpectedIndex {
+        let state = Migrations.getCurrentVersion(btree);
+        MemoryBTree.getExpectedIndex(state, btree_utils, key);
+    };
+
     public func clear(btree : StableMemoryBTree) {
         let state = Migrations.getCurrentVersion(btree);
         MemoryBTree.clear(state);

src/MemoryBTree/lib.mo

@@ -21,6 +21,8 @@ module {
     public type ValueUtils<V> = T.ValueUtils<V>;
     public type MemoryBTreeStats = BaseMemoryBTree.MemoryBTreeStats;
 
+    public type ExpectedIndex = BaseMemoryBTree.ExpectedIndex;
+
     /// Create a new stable store
     public func newStableStore(order : ?Nat) : StableStore = StableMemoryBTree.new(order);
 
@@ -90,9 +92,14 @@ module {
         public func getFromIndex(i : Nat) : (K, V) = BaseMemoryBTree.getFromIndex<K, V>(state, btree_utils, i);
 
         /// Get the index (sorted position) of the given key in the btree
-        /// If the key is not in the BTree, the expected index is returned
+        /// > Throws an error if the key is not found.
+        /// > Use `getExpectedIndex()` if you want to get the index without throwing an error.
         public func getIndex(key : K) : Nat = BaseMemoryBTree.getIndex<K, V>(state, btree_utils, key);
 
+        /// Get the index (sorted position) of the given key in the btree
+        /// Returns a `ExpectedIndex` variant that returns `#Found(index)` if the key exists or `#NotFound(index)` if it does not.
+        public func getExpectedIndex(key : K) : BaseMemoryBTree.ExpectedIndex = BaseMemoryBTree.getExpectedIndex<K, V>(state, btree_utils, key);
+
         /// Insert a new key-value pair into the BTree
         public func insert(key : K, val : V) : ?V = BaseMemoryBTree.insert<K, V>(state, btree_utils, key, val);