Skip to content

Commit

Permalink
wip: upsert nodes with the same key
Browse files Browse the repository at this point in the history
  • Loading branch information
@Nuhvi
Nuhvi committed Dec 20, 2023
1 parent bbd6652 commit 5d2f61c
Show file tree
Hide file tree
Showing 2 changed files with 97 additions and 25 deletions.
16 changes: 11 additions & 5 deletions mast/src/node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,8 +115,6 @@ impl Node {
new_child: Option<Hash>,
table: &mut Table<&[u8], (u64, &[u8])>,
) {
let old_hash = self.hash();

let old_child = match branch {
Branch::Left => self.left,
Branch::Right => self.right,
Expand Down Expand Up @@ -194,16 +192,17 @@ fn hash(bytes: &[u8]) -> Hash {
hasher.finalize()
}

#[derive(Debug)]
enum RefCountDiff {
Increment,
Decrement,
}

fn increment_ref_count(child: Option<Hash>, table: &mut Table<&[u8], (u64, &[u8])>) {
pub(crate) fn increment_ref_count(child: Option<Hash>, table: &mut Table<&[u8], (u64, &[u8])>) {
update_ref_count(child, RefCountDiff::Increment, table);
}

fn decrement_ref_count(child: Option<Hash>, table: &mut Table<&[u8], (u64, &[u8])>) {
pub(crate) fn decrement_ref_count(child: Option<Hash>, table: &mut Table<&[u8], (u64, &[u8])>) {
update_ref_count(child, RefCountDiff::Decrement, table);
}

Expand All @@ -213,16 +212,23 @@ fn update_ref_count(
table: &mut Table<&[u8], (u64, &[u8])>,
) {
if let Some(hash) = child {
dbg!("should update child ref", &child);
let mut existing = table
.get(hash.as_bytes().as_slice())
.unwrap()
.expect("Child shouldn't be messing!");

let (ref_count, bytes) = {
let (r, v) = existing.value();
(r + 1, v.to_vec())
(r, v.to_vec())
};
drop(existing);
dbg!((
"\n\n decrmenting blah blah blah child",
&child,
&ref_count,
&ref_diff
));

let ref_count = match ref_diff {
RefCountDiff::Increment => ref_count + 1,
Expand Down
106 changes: 86 additions & 20 deletions mast/src/treap.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
use blake3::{Hash, Hasher};
use redb::{Database, ReadableTable, Table, TableDefinition};
// use redb::{Database, ite, ReadableTable, Table, TableDefinition};
use redb::*;

use crate::node::{Branch, Node};
use crate::node::{decrement_ref_count, increment_ref_count, Branch, Node};

// TODO: remove unused
// TODO: remove unwrap
Expand Down Expand Up @@ -49,7 +50,6 @@ impl<'a> HashTreap<'a> {
if self.root.is_none() {
// We are done.
self.update_root(&node, &mut nodes_table);

break 'transaction;
}

Expand Down Expand Up @@ -94,7 +94,16 @@ impl<'a> HashTreap<'a> {

// Top down traversal of the binary search path.
while let Some(current) = self.get_node(&next) {
let should_zip = node.rank().as_bytes() > current.rank().as_bytes();
let node_rank = node.rank();
let curr_rank = current.rank();

if node_rank.as_bytes() == curr_rank.as_bytes() {
// Same key, we should update the value and return.
self.update_root(&node, &mut nodes_table);
break 'transaction;
}

let should_zip = node_rank.as_bytes() > curr_rank.as_bytes();

// Traverse left or right.
if key < current.key() {
Expand Down Expand Up @@ -137,18 +146,18 @@ impl<'a> HashTreap<'a> {
// \ |
// I | M

dbg!((
"unzipping left",
String::from_utf8(node.key().to_vec()).unwrap(),
&left_unzip_path
.iter()
.map(|n| String::from_utf8(n.key().to_vec()).unwrap())
.collect::<Vec<_>>(),
&right_unzip_path
.iter()
.map(|n| String::from_utf8(n.key().to_vec()).unwrap())
.collect::<Vec<_>>(),
));
// dbg!((
// "unzipping left",
// String::from_utf8(node.key().to_vec()).unwrap(),
// &left_unzip_path
// .iter()
// .map(|n| String::from_utf8(n.key().to_vec()).unwrap())
// .collect::<Vec<_>>(),
// &right_unzip_path
// .iter()
// .map(|n| String::from_utf8(n.key().to_vec()).unwrap())
// .collect::<Vec<_>>(),
// ));

let left_unzip_path_len = left_unzip_path.len();
for i in 0..left_unzip_path_len {
Expand All @@ -175,7 +184,9 @@ impl<'a> HashTreap<'a> {
// The last node in the path is special, since we need to clear its right
// child from older versions.
let child = right_unzip_path.get_mut(i).unwrap();
dbg!(("clearing the left child fuckin please", &child));
child.set_child(&Branch::Left, None, &mut nodes_table);
dbg!(("clearing the left child fuckin please", &child));

// Skip the last element for the first iterator
break;
Expand Down Expand Up @@ -244,12 +255,16 @@ impl<'a> HashTreap<'a> {
// === Private Methods ===

fn update_root(&mut self, node: &Node, table: &mut Table<&[u8], (u64, &[u8])>) {
// decrement_ref_count(self.root.clone().map(|n| n.hash()), table);

node.save(table);

// The tree is empty, the incoming node has to be the root, and we are done.
self.root = Some(node.clone());

// TODO: we need to persist the root change too to the storage.
// TODO: add a tag to persist snapshots.
increment_ref_count(self.root.clone().map(|n| n.hash()), table);
}

pub(crate) fn get_node(&self, hash: &Option<Hash>) -> Option<Node> {
Expand Down Expand Up @@ -277,9 +292,11 @@ impl<'a> HashTreap<'a> {
match node {
Some(n) => {
let left_check = self.get_node(n.left()).map_or(true, |left| {
dbg!(("left", &left));
n.rank().as_bytes() > left.rank().as_bytes() && self.check_rank(Some(left))
});
let right_check = self.get_node(n.right()).map_or(true, |right| {
dbg!(("right", &right));
n.rank().as_bytes() > right.rank().as_bytes() && self.check_rank(Some(right))
});

Expand All @@ -288,6 +305,31 @@ impl<'a> HashTreap<'a> {
None => true,
}
}

#[cfg(test)]
fn list_all_nodes(&self) {
let read_txn = self.db.begin_read().unwrap();
let nodes_table = read_txn.open_table(NODES_TABLE).unwrap();

let mut iter = nodes_table.iter().unwrap();

while let Some(existing) = iter.next() {
let key;
let data;
let existing = existing.unwrap();
{
key = existing.0.value();
data = existing.1.value();
}

println!(
"HEre is a node key:{:?} ref_count:{:?} node:{:?}",
Hash::from_bytes(key.try_into().unwrap()),
data.0,
Node::decode(data)
);
}
}
}

#[cfg(test)]
Expand Down Expand Up @@ -332,9 +374,10 @@ mod test {
let mut keys = [
// "D", "N", "P",
"X", // "F", "Z", "Y",
"A", "G", //
"C",
//"M", "H", "I", "J",
"A", //
// "G", //
// "C", //
//"M", "H", "I", "J",
];

// TODO: fix without sort.
Expand All @@ -344,7 +387,30 @@ mod test {
treap.insert(key.as_bytes(), b"0");
}

assert!(treap.verify_ranks());
assert!(treap.verify_ranks(), "Ranks are not correct");

treap.list_all_nodes();

println!("{}", treap.as_mermaid_graph())
}

#[test]
fn upsert() {
let file = tempfile::NamedTempFile::new().unwrap();
let db = Database::create(file.path()).unwrap();

let mut treap = HashTreap::new(&db);

let mut keys = ["X", "X"];

for key in keys.iter() {
treap.insert(key.as_bytes(), b"0");
}

assert!(treap.verify_ranks(), "Ranks are not correct");

// TODO: check the value.

println!("{}", treap.as_mermaid_graph())
}
}

0 comments on commit 5d2f61c

Please sign in to comment.