[Feat]: Split MemoryBTree key-value pairs into separate regions

Author: tomijaga

mops.toml

@@ -1,6 +1,6 @@
 [package]
 name = "memory-collection"
-version = "0.0.1"
+version = "0.0.2"
 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"]
@@ -11,13 +11,15 @@ base = "0.12.1"
 itertools = "0.2.1"
 memory-region = "1.2.4"
 buffer-deque = "0.1.0"
+memory-collection = "0.0.1"
 
 [dev-dependencies]
 test = "2.0.0"
 bench = "1.0.0"
 augmented-btrees = "0.5.2"
 fuzz = "0.2.1"
 MotokoStableBTree = "https://github.com/sardariuss/MotokoStableBTree#master@b590ede4489c2d4b2189299c3a2cc35dc4faa3fc"
+memory-collection-btree-v0-v0_0_1_migration = "https://github.com/NatLabs/memory-collection#master@bdba9e30ad95048634f23fe9c45c1e93c7787ce0"
 
 [toolchain]
 wasmtime = "14.0.4"

src/MemoryBTree/Base.mo

@@ -34,9 +34,6 @@ module {
     public type MemoryCmp<A> = MemoryCmp.MemoryCmp<A>;
 
     public type MemoryBTree = Migrations.MemoryBTree;
-    public type Node = Migrations.Node;
-    public type Leaf = Migrations.Leaf;
-    public type Branch = Migrations.Branch;
     public type VersionedMemoryBTree = Migrations.VersionedMemoryBTree;
 
     public type MemoryBlock = T.MemoryBlock;
@@ -60,13 +57,10 @@ module {
             var depth = 0;
             var is_root_a_leaf = true;
 
-            metadata = MemoryRegion.new();
-            blocks = MemoryRegion.new();
-            blobs = MemoryRegion.new();
-
             leaves = MemoryRegion.new();
             branches = MemoryRegion.new();
             data = MemoryRegion.new();
+            values = MemoryRegion.new();
 
         };
 
@@ -95,7 +89,7 @@ module {
     public let POINTER_SIZE = 12;
     public let LAYOUT_VERSION = 0;
 
-    let MC = {
+    public let MC = {
 
         REGION_HEADER_SIZE = 64;
 
@@ -110,12 +104,24 @@ module {
             COUNT_ADDRESS = 22; // 8 bytes
             DEPTH_ADDRESS = 30; // 1 byte
             IS_ROOT_A_LEAF_ADDRESS = 31; // 1 byte
+            VALUES_REGION_ID_ADDRESS = 32; // 4 bytes
 
             // values
             MAGIC : Blob = "BTR";
             LAYOUT_VERSION : Nat8 = 0;
         };
 
+        VALUES = {
+            // addresses
+            MAGIC_ADDRESS = 0;
+            LAYOUT_VERSION_ADDRESS = 3;
+            DATA_REGION_ID_ADDRESS = 4; // 4 bytes
+
+            // values
+            MAGIC : Blob = "VLS";
+            LAYOUT_VERSION : Nat8 = 0;
+        };
+
         BRANCHES = {
             // addresses
             MAGIC_ADDRESS = 0;
@@ -154,8 +160,16 @@ module {
         MemoryRegion.storeNat64(btree.data, MC.DATA.ROOT_ADDRESS, 0); // set to default value, will be updated once a node is created
         MemoryRegion.storeNat64(btree.data, MC.DATA.COUNT_ADDRESS, 0);
         MemoryRegion.storeNat8(btree.data, MC.DATA.DEPTH_ADDRESS, 0);
+        MemoryRegion.storeNat8(btree.data, MC.DATA.IS_ROOT_A_LEAF_ADDRESS, 0);
+        MemoryRegion.storeNat32(btree.data, MC.DATA.VALUES_REGION_ID_ADDRESS, Nat32.fromNat(MemoryRegion.id(btree.values)));
         assert MemoryRegion.allocated(btree.data) == MC.REGION_HEADER_SIZE;
 
+        ignore MemoryRegion.allocate(btree.values, MC.REGION_HEADER_SIZE);
+        MemoryRegion.storeBlob(btree.values, MC.VALUES.MAGIC_ADDRESS, MC.VALUES.MAGIC);
+        MemoryRegion.storeNat8(btree.values, MC.VALUES.LAYOUT_VERSION_ADDRESS, MC.VALUES.LAYOUT_VERSION);
+        MemoryRegion.storeNat32(btree.values, MC.VALUES.DATA_REGION_ID_ADDRESS, Nat32.fromNat(MemoryRegion.id(btree.data)));
+        assert MemoryRegion.allocated(btree.values) == MC.REGION_HEADER_SIZE;
+
         ignore MemoryRegion.allocate(btree.branches, MC.REGION_HEADER_SIZE);
         MemoryRegion.storeBlob(btree.branches, MC.BRANCHES.MAGIC_ADDRESS, MC.BRANCHES.MAGIC);
         MemoryRegion.storeNat8(btree.branches, MC.BRANCHES.LAYOUT_VERSION_ADDRESS, MC.BRANCHES.LAYOUT_VERSION);
@@ -187,7 +201,7 @@ module {
     };
 
     public func toVersioned(btree : MemoryBTree) : VersionedMemoryBTree {
-        #v0(btree);
+        Migrations.addVersion(btree);
     };
 
     public func bytes(btree : MemoryBTree) : Nat {

src/MemoryBTree/Migrations/V0_0_1.mo

@@ -0,0 +1,61 @@
+import Nat "mo:base/Nat";
+
+import MemoryRegion "mo:memory-region/MemoryRegion";
+import RevIter "mo:itertools/RevIter";
+// import Branch "mo:augmented-btrees/BpTree/Branch";
+
+import Blobify "../../TypeUtils/Blobify";
+import MemoryCmp "../../TypeUtils/MemoryCmp";
+
+module V0_0_1 {
+    // thinking the version should match the mops version where the layout or heap structure was last changed
+
+    public type Address = Nat;
+    type Size = Nat;
+
+    public type MemoryBlock = (Address, Size);
+
+    type MemoryRegionV1 = MemoryRegion.MemoryRegionV1;
+    type Blobify<A> = Blobify.Blobify<A>;
+    type RevIter<A> = RevIter.RevIter<A>;
+
+    public type MemoryCmp<A> = MemoryCmp.MemoryCmp<A>;
+
+    public type MemoryBTree = {
+        is_set : Bool; // is true, only keys are stored
+        node_capacity : Nat;
+        var count : Nat;
+        var root : Nat;
+        var branch_count : Nat; // number of branch nodes
+        var leaf_count : Nat; // number of leaf nodes
+        var depth : Nat;
+        var is_root_a_leaf : Bool;
+
+        leaves : MemoryRegionV1;
+        branches : MemoryRegionV1;
+        data : MemoryRegionV1;
+        values : MemoryRegionV1;
+
+    };
+
+    public type Leaf = (
+        nats : [var Nat], // [address, index, count]
+        adjacent_nodes : [var ?Nat], // [parent, prev, next] (is_root if parent is null)
+        key_blocks : [var ?(MemoryBlock)], // [... ((key address, key size), key blob)]
+        val_blocks : [var ?(MemoryBlock)],
+        kv_blobs : [var ?(Blob, Blob)],
+        _branch_children_nodes : [var ?Nat], // [... child address]
+        _branch_keys_blobs : [var ?Blob],
+    );
+
+    public type Branch = (
+        nats : [var Nat], // [address, index, count, subtree_size]
+        parent : [var ?Nat], // parent
+        key_blocks : [var ?(MemoryBlock)], // [... ((key address, key size), key blob)]
+        _leaf_val_blocks : [var ?(MemoryBlock)],
+        _leaf_kv_blobs : [var ?(Blob, Blob)],
+        children_nodes : [var ?Nat], // [... child address]
+        keys_blobs : [var ?Blob],
+    );
+
+};

src/MemoryBTree/Migrations/lib.mo

@@ -1,30 +1,42 @@
 import Debug "mo:base/Debug";
+import Nat32 "mo:base/Nat32";
+
+import MemoryRegion "mo:memory-region/MemoryRegion";
 
 import V0 "V0";
+import V0_0_1 "V0_0_1";
 
 module Migrations {
 
-    public type MemoryBTree = V0.MemoryBTree;
-    public type Leaf = V0.Leaf;
-    public type Node = V0.Node;
-    public type Branch = V0.Branch;
+    // should update to the latest version
+    public type MemoryBTree = V0_0_1.MemoryBTree;
+    public type Leaf = V0_0_1.Leaf;
+    public type Branch = V0_0_1.Branch;
 
     public type VersionedMemoryBTree = {
         #v0 : V0.MemoryBTree;
+        #v0_0_1 : V0_0_1.MemoryBTree;
     };
 
     public type StableStore = VersionedMemoryBTree;
 
     public func upgrade(versions : VersionedMemoryBTree) : VersionedMemoryBTree {
         switch (versions) {
-            case (#v0(v0)) versions;
+            case (#v0(v0)) {
+                Debug.trap("Migration Error: Migrating from #v0 is not supported");
+            };
+            case (#v0_0_1(v0_0_1)) versions;
         };
     };
 
     public func getCurrentVersion(versions : VersionedMemoryBTree) : MemoryBTree {
         switch (versions) {
-            case (#v0(v0)) v0;
+            case (#v0_0_1(curr)) curr;
             case (_) Debug.trap("Unsupported version. Please upgrade the memory buffer to the latest version.");
         };
     };
+
+    public func addVersion(btree : MemoryBTree) : VersionedMemoryBTree {
+        #v0_0_1(btree);
+    };
 };

src/MemoryBTree/Migrations/upgrades.mo

@@ -0,0 +1,176 @@
+import Debug "mo:base/Debug";
+import Nat32 "mo:base/Nat32";
+import Nat16 "mo:base/Nat16";
+import Nat64 "mo:base/Nat64";
+import Nat8 "mo:base/Nat8";
+import Iter "mo:base/Iter";
+
+import MemoryRegion "mo:memory-region/MemoryRegion";
+import Itertools "mo:itertools/Iter";
+
+import V0 "V0";
+import V0_0_1 "V0_0_1";
+
+module Migrations {
+
+    // should update to the latest version
+    public type MemoryBTree = V0_0_1.MemoryBTree;
+    public type Leaf = V0_0_1.Leaf;
+    public type Branch = V0_0_1.Branch;
+
+    public type VersionedMemoryBTree = {
+        #v0 : V0.MemoryBTree;
+        #v0_0_1 : V0_0_1.MemoryBTree;
+    };
+
+    public type StableStore = VersionedMemoryBTree;
+    type Address = Nat;
+
+    public func get_depth(btree : MemoryBTree, node_address : Nat) : Nat {
+        let depth = MemoryRegion.loadNat8(btree.branches, node_address + 3) |> Nat8.toNat(_);
+
+        depth;
+    };
+
+    public func has_leaves(btree : MemoryBTree, branch_address : Nat) : Bool {
+        let depth = get_depth(btree, branch_address);
+        depth == 2;
+    };
+
+    public func CHILDREN_START(btree : MemoryBTree) : Nat {
+        64 + ((btree.node_capacity - 1) * 8);
+    };
+    public func get_child_offset(btree : MemoryBTree, branch_address : Nat, i : Nat) : Nat {
+        branch_address + CHILDREN_START(btree) + (i * 8);
+    };
+    public func get_child(btree : MemoryBTree, branch_address : Nat, i : Nat) : ?Nat {
+
+        MemoryRegion.loadNat64(btree.branches, get_child_offset(btree, branch_address, i))
+        |> ?Nat64.toNat(_);
+    };
+
+    public func get_min_leaf_address(btree : MemoryBTree) : Nat {
+        var curr = btree.root;
+        var is_address_a_leaf = btree.is_root_a_leaf;
+
+        loop {
+            switch (is_address_a_leaf) {
+                case (false) {
+                    let ?first_child = get_child(btree, curr, 0) else Debug.trap("get_min_leaf: accessed a null value");
+                    is_address_a_leaf := has_leaves(btree, curr);
+                    curr := first_child;
+                };
+                case (true) return curr;
+            };
+        };
+    };
+
+    public func get_kv_address_offset(leaf_address : Nat, i : Nat) : Nat {
+        leaf_address + 64 + (i * 8);
+    };
+
+    public func leaf_get_kv_address(btree : MemoryBTree, address : Nat, i : Nat) : ?Nat {
+        let kv_address_offset = get_kv_address_offset(address, i);
+        let opt_id = MemoryRegion.loadNat64(btree.leaves, kv_address_offset);
+
+        if (opt_id == 0x00) null else ?(Nat64.toNat(opt_id));
+    };
+
+    public func leaf_get_count(btree : MemoryBTree, address : Nat) : Nat {
+        MemoryRegion.loadNat16(btree.leaves, address + 6) |> Nat16.toNat(_);
+    };
+    public func leaf_get_kv_address_offset(leaf_address : Nat, i : Nat) : Nat {
+        leaf_address + 64 + (i * 8);
+    };
+
+    public func leaf_get_next(btree : MemoryBTree, address : Nat) : ?Nat {
+
+        let next = MemoryRegion.loadNat64(btree.leaves, address + 24);
+        if (next == 0x00) return null;
+        ?Nat64.toNat(next);
+    };
+
+    public func kv_block_addresses(btree : MemoryBTree) : Iter.Iter<Address> {
+
+        let min_leaf = get_min_leaf_address(btree);
+        var i = 0;
+        var leaf_count = leaf_get_count(btree, min_leaf);
+        var var_leaf = ?min_leaf;
+
+        object {
+            public func next() : ?Address {
+                let ?leaf = var_leaf else return null;
+
+                if (i >= leaf_count) {
+                    switch (leaf_get_next(btree, leaf)) {
+                        case (null) var_leaf := null;
+                        case (?next_address) {
+                            var_leaf := ?next_address;
+                            leaf_count := leaf_get_count(btree, leaf);
+                        };
+                    };
+
+                    i := 0;
+                    return next();
+                };
+
+                let address = leaf_get_kv_address(btree, leaf, i);
+                i += 1;
+                return address;
+            };
+        };
+
+    };
+
+    public func upgrade(versions : VersionedMemoryBTree) : VersionedMemoryBTree {
+        switch (versions) {
+            case (#v0(v0)) {
+                let values = MemoryRegion.new();
+
+                let VALUES_REGION_ID_ADDRESS = 32;
+                MemoryRegion.storeNat32(v0.data, VALUES_REGION_ID_ADDRESS, Nat32.fromNat(MemoryRegion.id(values)));
+
+                // move all values from data region to values region
+                // might not upgrade if the data is too large
+
+                let REFERENCE_COUNT_START = 0;
+                let KEY_SIZE_START = 1;
+                let VAL_POINTER_START = 3;
+                let VAL_SIZE_START = 11;
+                let KEY_BLOB_START = 15;
+
+                let BLOCK_ENTRY_SIZE = 15;
+
+                var key_block_address = 64;
+
+                for (i in Itertools.range(0, v0.count)) {
+                    let key_size = MemoryRegion.loadNat16(v0.data, key_block_address + KEY_SIZE_START) |> Nat16.toNat(_);
+                    let val_size = MemoryRegion.loadNat32(v0.data, key_block_address + VAL_SIZE_START) |> Nat32.toNat(_);
+                    let val_address = MemoryRegion.loadNat64(v0.data, key_block_address + VAL_POINTER_START) |> Nat64.toNat(_);
+
+                    let val_blob = MemoryRegion.loadBlob(v0.data, val_address, val_size);
+                    MemoryRegion.deallocate(v0.data, val_address, val_size);
+
+                    let new_val_address = MemoryRegion.addBlob(values, val_blob);
+
+                    MemoryRegion.storeNat64(v0.data, key_block_address + VAL_POINTER_START, Nat64.fromNat(new_val_address));
+
+                    key_block_address += BLOCK_ENTRY_SIZE + key_size + val_size;
+
+                };
+
+                #v0_0_1({
+                    v0 with values;
+                    var count = v0.count;
+                    var root = v0.root;
+                    var branch_count = v0.branch_count;
+                    var leaf_count = v0.leaf_count;
+                    var depth = v0.depth;
+                    var is_root_a_leaf = v0.is_root_a_leaf;
+                });
+            };
+            case (#v0_0_1(v0_0_1)) versions;
+        };
+    };
+
+};