NatLabs · tomijaga · Dec 21, 2024
diff --git a/mops.toml b/mops.toml
@@ -11,7 +11,6 @@ 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"

diff --git a/src/MemoryBTree/Base.mo b/src/MemoryBTree/Base.mo
diff --git a/src/MemoryBTree/Migrations/V0_0_2.mo b/src/MemoryBTree/Migrations/V0_0_2.mo
@@ -0,0 +1,62 @@
+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;
+        var supports_key_compression : 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],
+    );
+
+};
diff --git a/src/MemoryBTree/Migrations/lib.mo b/src/MemoryBTree/Migrations/lib.mo
@@ -5,17 +5,19 @@ import MemoryRegion "mo:memory-region/MemoryRegion";
 
 import V0 "V0";
 import V0_0_1 "V0_0_1";
+import V0_0_2 "V0_0_2";
 
 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 MemoryBTree = V0_0_2.MemoryBTree;
+    public type Leaf = V0_0_2.Leaf;
+    public type Branch = V0_0_2.Branch;
 
     public type VersionedMemoryBTree = {
         #v0 : V0.MemoryBTree;
         #v0_0_1 : V0_0_1.MemoryBTree;
+        #v0_0_2 : V0_0_2.MemoryBTree;
     };
 
     public type StableStore = VersionedMemoryBTree;
@@ -25,18 +27,29 @@ module Migrations {
             case (#v0(v0)) {
                 Debug.trap("Migration Error: Migrating from #v0 is not supported");
             };
-            case (#v0_0_1(v0_0_1)) versions;
+            case (#v0_0_1(v0_0_1)) #v0_0_2({
+                v0_0_1 with
+                var supports_key_compression = false;
+                var branch_count = v0_0_1.branch_count;
+                var count = v0_0_1.count;
+                var root = v0_0_1.root;
+                var depth = v0_0_1.depth;
+                var is_root_a_leaf = v0_0_1.is_root_a_leaf;
+                var leaf_count = v0_0_1.leaf_count;
+
+            });
+            case (#v0_0_2(_)) versions;
         };
     };
 
     public func getCurrentVersion(versions : VersionedMemoryBTree) : MemoryBTree {
         switch (versions) {
-            case (#v0_0_1(curr)) curr;
+            case (#v0_0_2(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);
+        #v0_0_2(btree);
     };
 };
diff --git a/src/MemoryBTree/Stable.mo b/src/MemoryBTree/Stable.mo
@@ -18,17 +18,17 @@ module StableMemoryBTree {
         };
     };
 
-    public func new(order : ?Nat) : StableMemoryBTree {
-        let btree = MemoryBTree.new(order);
+    public func new(opt_options : ?T.InitOptions) : StableMemoryBTree {
+        let btree = MemoryBTree.new(opt_options);
         MemoryBTree.toVersioned(btree);
     };
 
     public func fromArray<K, V>(
         btree_utils : BTreeUtils<K, V>,
         arr : [(K, V)],
-        order : ?Nat,
+        opt_options : ?T.InitOptions,
     ) : StableMemoryBTree {
-        let btree = MemoryBTree.fromArray(btree_utils, arr, order);
+        let btree = MemoryBTree.fromArray(btree_utils, arr, opt_options);
         MemoryBTree.toVersioned(btree);
     };
 

diff --git a/src/MemoryBTree/lib.mo b/src/MemoryBTree/lib.mo
@@ -24,7 +24,7 @@ module {
     public type ExpectedIndex = BaseMemoryBTree.ExpectedIndex;
 
     /// Create a new stable store
-    public func newStableStore(order : ?Nat) : StableStore = StableMemoryBTree.new(order);
+    public func newStableStore(opt_options : ?T.InitOptions) : StableStore = StableMemoryBTree.new(opt_options);
 
     /// Upgrade an older version of the BTree to the latest version
     public func upgrade<K, V>(sstore : StableStore) : StableStore {