Skip to content

Commit

Permalink
[Feat]: Split MemoryBTree key-value pairs into separate regions
Browse files Browse the repository at this point in the history
  • Loading branch information
tomijaga committed Aug 24, 2024
1 parent f846d68 commit d099846
Show file tree
Hide file tree
Showing 12 changed files with 365 additions and 55 deletions.
4 changes: 3 additions & 1 deletion mops.toml
Original file line number Diff line number Diff line change
@@ -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"]
Expand All @@ -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"
Expand Down
32 changes: 23 additions & 9 deletions src/MemoryBTree/Base.mo
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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();

};

Expand Down Expand Up @@ -95,7 +89,7 @@ module {
public let POINTER_SIZE = 12;
public let LAYOUT_VERSION = 0;

let MC = {
public let MC = {

REGION_HEADER_SIZE = 64;

Expand All @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -187,7 +201,7 @@ module {
};

public func toVersioned(btree : MemoryBTree) : VersionedMemoryBTree {
#v0(btree);
Migrations.addVersion(btree);
};

public func bytes(btree : MemoryBTree) : Nat {
Expand Down
61 changes: 61 additions & 0 deletions src/MemoryBTree/Migrations/V0_0_1.mo
Original file line number Diff line number Diff line change
@@ -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],
);

};
24 changes: 18 additions & 6 deletions src/MemoryBTree/Migrations/lib.mo
Original file line number Diff line number Diff line change
@@ -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);
};
};
176 changes: 176 additions & 0 deletions src/MemoryBTree/Migrations/upgrades.mo
Original file line number Diff line number Diff line change
@@ -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;
};
};

};
Loading

0 comments on commit d099846

Please sign in to comment.