[MemoryBTree]: Add rangeKeys(), rangeVals(), scanKeys(), scanVals() for more efficient range request

Author: tomijaga

mops.toml

@@ -1,6 +1,6 @@
 [package]
 name = "memory-collection"
-version = "0.0.3"
+version = "0.0.4"
 description = "A collection of data structures that store their data in stable memory"
 repository = "https://github.com/NatLabs/memory-collection"
 keywords = ["region", "sorted", "persistent", "map", "encoding"]

src/MemoryBTree/Base.mo

@@ -934,21 +934,64 @@ module {
 
     };
 
-    public func range<K, V>(btree : MemoryBTree, btree_utils : BTreeUtils<K, V>, start : Nat, end : Nat) : RevIter<(K, V)> {
+    func process_range_request<ReturnValue>(
+        btree : MemoryBTree,
+        start : Nat,
+        end : Nat,
+        kv_address_to_result_mapper : (Address) -> ReturnValue,
+    ) : RevIter<ReturnValue> {
+
         let (start_node, start_node_index) = Methods.get_leaf_node_by_index(btree, start);
         let (end_node, end_node_index) = Methods.get_leaf_node_by_index(btree, end);
 
         let start_index = start_node_index : Nat;
         let end_index = end_node_index : Nat;
 
-        RevIter.map<(Blob, Blob), (K, V)>(
-            Methods.new_blobs_iterator(btree, start_node, start_index, end_node, end_index),
-            func((key_blob, val_blob) : (Blob, Blob)) : (K, V) {
+        RevIter.map<Address, ReturnValue>(
+            Methods.new_kv_block_address_iterator(btree, start_node, start_index, end_node, end_index),
+            kv_address_to_result_mapper,
+        );
+    };
+
+    public func range<K, V>(btree : MemoryBTree, btree_utils : BTreeUtils<K, V>, start : Nat, end : Nat) : RevIter<(K, V)> {
+        process_range_request<(K, V)>(
+            btree,
+            start,
+            end,
+            func(kv_address : Address) : (K, V) {
+                let key_blob = MemoryBlock.get_key_blob(btree, kv_address);
+                let val_blob = MemoryBlock.get_val_blob(btree, kv_address);
                 Methods.deserialize_kv_blobs<K, V>(btree_utils, key_blob, val_blob);
             },
         );
     };
 
+    public func rangeKeys<K, V>(btree : MemoryBTree, btree_utils : BTreeUtils<K, V>, start : Nat, end : Nat) : RevIter<K> {
+
+        process_range_request(
+            btree,
+            start,
+            end,
+            func(kv_address : Address) : K {
+                let key_blob = MemoryBlock.get_key_blob(btree, kv_address);
+                btree_utils.key.blobify.from_blob(key_blob);
+            },
+        );
+
+    };
+
+    public func rangeVals<K, V>(btree : MemoryBTree, btree_utils : BTreeUtils<K, V>, start : Nat, end : Nat) : RevIter<V> {
+        process_range_request(
+            btree,
+            start,
+            end,
+            func(kv_address : Address) : V {
+                let val_blob = MemoryBlock.get_val_blob(btree, kv_address);
+                btree_utils.value.blobify.from_blob(val_blob);
+            },
+        );
+    };
+
     public func getInterval<K, V>(btree : MemoryBTree, btree_utils : BTreeUtils<K, V>, start : ?K, end : ?K) : (Nat, Nat) {
         let start_rank = switch (start) {
             case (?key) getIndex(btree, btree_utils, key);
@@ -963,7 +1006,14 @@ module {
         (start_rank, end_rank);
     };
 
-    public func scan<K, V>(btree : MemoryBTree, btree_utils : BTreeUtils<K, V>, start : ?K, end : ?K) : RevIter<(K, V)> {
+    func process_scan_request<K, V, ReturnValue>(
+        btree : MemoryBTree,
+        btree_utils : BTreeUtils<K, V>,
+        start : ?K,
+        end : ?K,
+        kv_address_to_result_mapper : (Address) -> ReturnValue,
+    ) : RevIter<ReturnValue> {
+
         let start_address = switch (start) {
             case (?key) {
                 let key_blob = btree_utils.key.blobify.to_blob(key);
@@ -1008,11 +1058,51 @@ module {
 
         var j = if (end_index >= 0) Int.abs(end_index) + 1 else Int.abs(end_index) - 1 : Nat;
 
-        RevIter.map<(Blob, Blob), (K, V)>(
-            Methods.new_blobs_iterator(btree, start_address, i, end_address, j),
-            func((key_blob, val_blob) : (Blob, Blob)) : (K, V) {
+        RevIter.map<Address, ReturnValue>(
+            Methods.new_kv_block_address_iterator(btree, start_address, i, end_address, j),
+            kv_address_to_result_mapper,
+        );
+
+    };
+
+    public func scan<K, V>(btree : MemoryBTree, btree_utils : BTreeUtils<K, V>, start : ?K, end : ?K) : RevIter<(K, V)> {
+        process_scan_request<K, V, (K, V)>(
+            btree,
+            btree_utils,
+            start,
+            end,
+            func(kv_block_address : Address) : (K, V) {
+                let key_blob = MemoryBlock.get_key_blob(btree, kv_block_address);
+                let val_blob = MemoryBlock.get_val_blob(btree, kv_block_address);
+
                 Methods.deserialize_kv_blobs<K, V>(btree_utils, key_blob, val_blob);
             },
         );
     };
+
+    public func scanKeys<K, V>(btree : MemoryBTree, btree_utils : BTreeUtils<K, V>, start : ?K, end : ?K) : RevIter<K> {
+        process_scan_request<K, V, K>(
+            btree,
+            btree_utils,
+            start,
+            end,
+            func(kv_block_address : Address) : K {
+                let key_blob = MemoryBlock.get_key_blob(btree, kv_block_address);
+                btree_utils.key.blobify.from_blob(key_blob);
+            },
+        );
+    };
+
+    public func scanVals<K, V>(btree : MemoryBTree, btree_utils : BTreeUtils<K, V>, start : ?K, end : ?K) : RevIter<V> {
+        process_scan_request<K, V, V>(
+            btree,
+            btree_utils,
+            start,
+            end,
+            func(kv_block_address : Address) : V {
+                let val_blob = MemoryBlock.get_val_blob(btree, kv_block_address);
+                btree_utils.value.blobify.from_blob(val_blob);
+            },
+        );
+    };
 };

src/MemoryBTree/Stable.mo

@@ -138,11 +138,31 @@ module StableMemoryBTree {
         MemoryBTree.scan(state, btree_utils, start, end);
     };
 
+    public func scanKeys<K, V>(btree : StableMemoryBTree, btree_utils : BTreeUtils<K, V>, start : ?K, end : ?K) : RevIter<K> {
+        let state = Migrations.getCurrentVersion(btree);
+        MemoryBTree.scanKeys(state, btree_utils, start, end);
+    };
+
+    public func scanVals<K, V>(btree : StableMemoryBTree, btree_utils : BTreeUtils<K, V>, start : ?K, end : ?K) : RevIter<V> {
+        let state = Migrations.getCurrentVersion(btree);
+        MemoryBTree.scanVals(state, btree_utils, start, end);
+    };
+
     public func range<K, V>(btree : StableMemoryBTree, btree_utils : BTreeUtils<K, V>, start : Nat, end : Nat) : RevIter<(K, V)> {
         let state = Migrations.getCurrentVersion(btree);
         MemoryBTree.range(state, btree_utils, start, end);
     };
 
+    public func rangeKeys<K, V>(btree : StableMemoryBTree, btree_utils : BTreeUtils<K, V>, start : Nat, end : Nat) : RevIter<K> {
+        let state = Migrations.getCurrentVersion(btree);
+        MemoryBTree.rangeKeys(state, btree_utils, start, end);
+    };
+
+    public func rangeVals<K, V>(btree : StableMemoryBTree, btree_utils : BTreeUtils<K, V>, start : Nat, end : Nat) : RevIter<V> {
+        let state = Migrations.getCurrentVersion(btree);
+        MemoryBTree.rangeVals(state, btree_utils, start, end);
+    };
+
     public func size(btree : StableMemoryBTree) : Nat {
         let state = Migrations.getCurrentVersion(btree);
         MemoryBTree.size(state);

src/MemoryBTree/lib.mo

@@ -139,9 +139,21 @@ module {
         /// Returns a reversible iterator over the entries in the given range
         public func range(i : Nat, j : Nat) : RevIter<(K, V)> = BaseMemoryBTree.range(state, btree_utils, i, j);
 
-        /// Returns a reversible iterator over the entries in the given range
+        /// Returns a reversible iterator over the keys in the given range
+        public func rangeKeys(i : Nat, j : Nat) : RevIter<K> = BaseMemoryBTree.rangeKeys(state, btree_utils, i, j);
+
+        /// Returns a reversible iterator over the values in the given range
+        public func rangeVals(i : Nat, j : Nat) : RevIter<V> = BaseMemoryBTree.rangeVals(state, btree_utils, i, j);
+
+        /// Returns a reversible iterator over the entries between the given keys
         public func scan(start : ?K, end : ?K) : RevIter<(K, V)> = BaseMemoryBTree.scan(state, btree_utils, start, end);
 
+        /// Returns a reversible iterator over the keys between the given ones
+        public func scanKeys(start : ?K, end : ?K) : RevIter<K> = BaseMemoryBTree.scanKeys(state, btree_utils, start, end);
+
+        /// Returns a reversible iterator over the values between the given keys
+        public func scanVals(start : ?K, end : ?K) : RevIter<V> = BaseMemoryBTree.scanVals(state, btree_utils, start, end);
+
         /// Returns the number of entries in the BTree
         public func size() : Nat = BaseMemoryBTree.size(state);
 

src/MemoryBTree/modules/Methods.mo

@@ -14,8 +14,9 @@ import T "Types";
 import Leaf "Leaf";
 import Branch "Branch";
 import Migrations "../Migrations";
+import MemoryBlock "MemoryBlock";
 
-module {
+module Methods {
     type MemoryBTree = Migrations.MemoryBTree;
     type MemoryBlock = T.MemoryBlock;
 
@@ -312,6 +313,110 @@ module {
         (get_node(root), search_index);
     };
 
+    public func new_kv_block_address_iterator(
+        btree : MemoryBTree,
+        start_leaf : Nat,
+        start_index : Nat,
+        end_leaf : Nat,
+        end_index : Nat // exclusive
+    ) : RevIter<Address> {
+
+        var start = start_leaf;
+        var i = start_index;
+        var start_count = Leaf.get_count(btree, start_leaf);
+
+        var end = end_leaf;
+        var j = end_index;
+
+        var terminate = false;
+
+        func next() : ?Address {
+            if (terminate) return null;
+
+            if (start == end and i >= j) {
+                return null;
+            };
+
+            if (i >= start_count) {
+                switch (Leaf.get_next(btree, start)) {
+                    case (null) {
+                        terminate := true;
+                    };
+                    case (?next_address) {
+                        start := next_address;
+                        start_count := Leaf.get_count(btree, next_address);
+                    };
+                };
+
+                i := 0;
+                return next();
+            };
+
+            let opt_kv_block_address = Leaf.get_kv_address(btree, start, i);
+
+            i += 1;
+            return opt_kv_block_address;
+        };
+
+        func nextFromEnd() : ?Address {
+            if (terminate) return null;
+
+            if (start == end and i >= j) return null;
+
+            if (j == 0) {
+                switch (Leaf.get_prev(btree, end)) {
+                    case (null) terminate := true;
+                    case (?prev_address) {
+                        end := prev_address;
+                        j := Leaf.get_count(btree, prev_address);
+                    };
+                };
+
+                return nextFromEnd();
+            };
+
+            let opt_kv_block_address = Leaf.get_kv_address(btree, end, j - 1);
+
+            j -= 1;
+
+            return opt_kv_block_address;
+        };
+
+        RevIter.new(next, nextFromEnd);
+    };
+
+    public func KeyBlobIterator(
+        btree : MemoryBTree,
+        start_leaf : Nat,
+        start_index : Nat,
+        end_leaf : Nat,
+        end_index : Nat,
+    ) : RevIter<Blob> {
+        RevIter.map(
+            new_kv_block_address_iterator(btree, start_leaf, start_index, end_leaf, end_index),
+            func(kv_block_address : Address) : Blob {
+                MemoryBlock.get_key_blob(btree, kv_block_address);
+            },
+        );
+    };
+
+    public func ValueBlobIterator(
+        btree : MemoryBTree,
+        start_leaf : Nat,
+        start_index : Nat,
+        end_leaf : Nat,
+        end_index : Nat,
+    ) : RevIter<Blob> {
+
+        RevIter.map(
+            Methods.new_kv_block_address_iterator(btree, start_leaf, start_index, end_leaf, end_index),
+            func(kv_block_address : Address) : Blob {
+                MemoryBlock.get_val_blob(btree, kv_block_address);
+            },
+        )
+
+    };
+
     public func new_blobs_iterator(
         btree : MemoryBTree,
         start_leaf : Nat,
@@ -424,6 +529,14 @@ module {
 
     };
 
+    public func deserialize_key_blob<K>(btree_utils : BTreeUtils<K, Nat>, key_blob : Blob) : K {
+        btree_utils.key.blobify.from_blob(key_blob);
+    };
+
+    public func deserialize_val_blob<V>(btree_utils : BTreeUtils<Nat, V>, val_blob : Blob) : V {
+        btree_utils.value.blobify.from_blob(val_blob);
+    };
+
     public func deserialize_kv_blobs<K, V>(btree_utils : BTreeUtils<K, V>, key_blob : Blob, val_blob : Blob) : (K, V) {
         let key = btree_utils.key.blobify.from_blob(key_blob);
         let value = btree_utils.value.blobify.from_blob(val_blob);