Skip to content

Adding Key Compression to BTree (Head + Tail) #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tomijaga wants to merge 3 commits into
base: main
Choose a base branch
from
+255 −113
Open

Adding Key Compression to BTree (Head + Tail) #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d548252
Update median_key to separator_key
tomijaga Jul 19, 2025
71e7b92
store branch keys independently in their own memory blocks
tomijaga Jul 19, 2025
ac3baa1
Fix memory leaks and add tests to check for leaks
tomijaga Jul 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 29 additions & 14 deletions src/MemoryBTree/Base.mo
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ import MemoryBlock "modules/MemoryBlock";
import Branch "modules/Branch";
import Utils "../Utils";
import Migrations "Migrations";
import Leaf "modules/Leaf";
import LeafModule "modules/Leaf";
import T "modules/Types";
import TypeUtils "../TypeUtils";

Expand All @@ -45,6 +45,7 @@ module {

let CACHE_LIMIT = 50_000;
let DEFAULT_ORDER = 256;
public let Leaf = LeafModule;

public func _new_with_options(node_capacity : ?Nat, opt_cache_size : ?Nat, is_set : Bool) : MemoryBTree {
let cache_size = Option.get(opt_cache_size, CACHE_LIMIT);
Expand Down Expand Up @@ -375,8 +376,8 @@ module {
var opt_parent = Leaf.get_parent(btree, right_node_address);
var right_index = Leaf.get_index(btree, right_node_address);

let ?first_key_address = Leaf.get_kv_address(btree, right_node_address, 0) else Debug.trap("insert: first_key_address accessed a null value");
var median_key_address = first_key_address;
let ?first_key = Leaf.get_key_blob(btree, right_node_address, 0) else Debug.trap("insert: first_key_address accessed a null value");
var separator_key_address = MemoryBlock.Branch.store_key_blob(btree, first_key);

// assert Leaf.get_count(btree, left_node_address) == (btree.node_capacity / 2) + 1;
// assert Leaf.get_count(btree, right_node_address) == (btree.node_capacity / 2);
Expand All @@ -392,7 +393,7 @@ module {
// Debug.print("found branch with enough space");
// Debug.print("parent before insert: " # debug_show Branch.from_memory(btree, parent_address));

Branch.insert(btree, parent_address, right_index, median_key_address, right_node_address);
Branch.insert(btree, parent_address, right_index, separator_key_address, right_node_address);
update_count(btree, btree.count + 1);

// Debug.print("parent after insert: " # debug_show Branch.from_memory(btree, parent_address));
Expand All @@ -402,12 +403,14 @@ module {

// otherwise split parent
left_node_address := parent_address;
right_node_address := Branch.split(btree, left_node_address, right_index, median_key_address, right_node_address);
right_node_address := Branch.split(btree, left_node_address, right_index, separator_key_address, right_node_address);
update_branch_count(btree, btree.branch_count + 1);

// The separator key is temporarily stored in the right node at the last position.
// We need to move it to the left node and update the separator key address.
let ?first_key_address = Branch.get_key_address(btree, right_node_address, btree.node_capacity - 2) else Debug.trap("4. insert: accessed a null value in first key of branch");
Branch.set_key_address_to_null(btree, right_node_address, btree.node_capacity - 2);
median_key_address := first_key_address;
separator_key_address := first_key_address;

right_index := Branch.get_index(btree, right_node_address);
opt_parent := Branch.get_parent(btree, right_node_address);
Expand All @@ -423,7 +426,7 @@ module {
Branch.update_depth(btree, new_root, new_depth);
assert Branch.get_depth(btree, new_root) == new_depth;

Branch.put_key_address(btree, new_root, 0, median_key_address);
Branch.put_key_address(btree, new_root, 0, separator_key_address);

Branch.add_child(btree, new_root, left_node_address);
Branch.add_child(btree, new_root, right_node_address);
Expand Down Expand Up @@ -722,8 +725,8 @@ module {

if (elem_index == 0) {
// if the first element is removed then update the parent key
let ?next_key_address = Leaf.get_kv_address(btree, leaf_address, 0) else Debug.trap("remove: next_key_block is null");
Branch.update_median_key_address(btree, parent, leaf_index, next_key_address);
let ?next_key = Leaf.get_key_blob(btree, leaf_address, 0) else Debug.trap("remove: next_key_block is null");
Branch.update_separator_key(btree, parent, leaf_index, next_key);
};

let min_count = btree.node_capacity / 2;
Expand Down Expand Up @@ -752,21 +755,24 @@ module {

if (Leaf.redistribute(btree, leaf_address, neighbour)) {

let ?key_address = Leaf.get_kv_address(btree, right, 0) else Debug.trap("remove: key_block is null");
Branch.put_key_address(btree, parent, right_index - 1, key_address);
let ?key_blob = Leaf.get_key_blob(btree, right, 0) else Debug.trap("remove: key_block is null");
Branch.replace_key(btree, parent, right_index - 1, key_blob);

return ?prev_val;
};

// Debug.print("merging leaf");

// remove merged leaf from parent
// Debug.print("remove merged index: " # debug_show right_index);
// Debug.print("remove merged index: " # debug_show right_index # ", address: " # debug_show right);
// Debug.print("parent: " # debug_show Branch.from_memory(btree, parent));

// merge leaf with neighbour
Leaf.merge(btree, left, right);
Branch.remove(btree, parent, right_index);
let removed_key_address = Branch.remove(btree, parent, right_index);

// Deallocate the key that was separating the merged leaves
MemoryBlock.Branch.remove_key_blob(btree, removed_key_address);

// deallocate right leaf that was merged into left
Leaf.deallocate(btree, right);
Expand All @@ -790,6 +796,10 @@ module {
update_root(btree, child);
update_is_root_a_leaf(btree, child_is_leaf);
update_depth(btree, btree.depth - 1);

Branch.deallocate(btree, parent);
update_branch_count(btree, btree.branch_count - 1);

return ?prev_val;

} else {
Expand Down Expand Up @@ -822,7 +832,12 @@ module {

let merged_branch = Branch.merge(btree, branch, neighbour);
let merged_branch_index = Branch.get_index(btree, merged_branch);
Branch.remove(btree, parent, merged_branch_index);
let removed_key_address = Branch.remove(btree, parent, merged_branch_index);

// The separator key is transferred to the merged branch during merge,
// so it should not be deallocated here
// MemoryBlock.Branch.remove_key_blob(btree, removed_key_address);

Branch.deallocate(btree, merged_branch);
update_branch_count(btree, btree.branch_count - 1);

Expand Down
Loading
Loading