Skip to content

Commit

Permalink
Delegate branch-storage logic to the branch mod
Browse files Browse the repository at this point in the history
  • Loading branch information
@richardpringle
richardpringle committed Nov 23, 2023
1 parent b2b4da8 commit 0dceb44
Show file tree
Hide file tree
Showing 4 changed files with 129 additions and 80 deletions.
12 changes: 6 additions & 6 deletions firewood/src/merkle.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use thiserror::Error;
mod node;
mod trie_hash;

pub use node::{BranchNode, Data, ExtNode, LeafNode, Node, NodeType, PartialPath, MAX_CHILDREN};
pub use node::{BranchNode, Data, ExtNode, LeafNode, Node, NodeType, PartialPath};
pub use trie_hash::{TrieHash, TRIE_HASH_LEN};

type ObjRef<'a> = shale::ObjRef<'a, Node>;
Expand Down Expand Up @@ -81,7 +81,7 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
.put_item(
Node::from_branch(BranchNode {
// path: vec![].into(),
children: [None; MAX_CHILDREN],
children: [None; BranchNode::MAX_CHILDREN],
value: None,
children_encoded: Default::default(),
}),
Expand Down Expand Up @@ -199,7 +199,7 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
));
let leaf_address = self.put_node(new_node)?.as_ptr();

let mut chd = [None; MAX_CHILDREN];
let mut chd = [None; BranchNode::MAX_CHILDREN];

let last_matching_nibble = matching_path[idx];
chd[last_matching_nibble as usize] = Some(leaf_address);
Expand Down Expand Up @@ -340,7 +340,7 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
};

// [parent] (-> [ExtNode]) -> [branch with v] -> [Leaf]
let mut children = [None; MAX_CHILDREN];
let mut children = [None; BranchNode::MAX_CHILDREN];

children[idx] = leaf_address.into();

Expand Down Expand Up @@ -561,7 +561,7 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
};

if let Some((idx, more, ext, val)) = info {
let mut chd = [None; MAX_CHILDREN];
let mut chd = [None; BranchNode::MAX_CHILDREN];

let c_ptr = if more {
u_ptr
Expand Down Expand Up @@ -1691,7 +1691,7 @@ mod tests {
fn branch(value: Vec<u8>, encoded_child: Option<Vec<u8>>) -> Node {
let children = Default::default();
let value = Some(Data(value));
let mut children_encoded = <[Option<Vec<u8>>; MAX_CHILDREN]>::default();
let mut children_encoded = <[Option<Vec<u8>>; BranchNode::MAX_CHILDREN]>::default();

if let Some(child) = encoded_child {
children_encoded[0] = Some(child);
Expand Down
74 changes: 18 additions & 56 deletions firewood/src/merkle/node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ mod extension;
mod leaf;
mod partial_path;

pub use branch::{BranchNode, MAX_CHILDREN, SIZE as BRANCH_NODE_SIZE};
pub use branch::BranchNode;
pub use extension::ExtNode;
pub use leaf::{LeafNode, SIZE as LEAF_NODE_SIZE};
pub use partial_path::PartialPath;
Expand Down Expand Up @@ -114,7 +114,7 @@ impl NodeType {
}
}
// TODO: add path
BRANCH_NODE_SIZE => Ok(NodeType::Branch(BranchNode::decode(buf)?.into())),
BranchNode::MSIZE => Ok(NodeType::Branch(BranchNode::decode(buf)?.into())),
size => Err(Box::new(bincode::ErrorKind::Custom(format!(
"invalid size: {size}"
)))),
Expand Down Expand Up @@ -216,7 +216,7 @@ impl Node {
inner: NodeType::Branch(
BranchNode {
// path: vec![].into(),
children: [Some(DiskAddress::null()); MAX_CHILDREN],
children: [Some(DiskAddress::null()); BranchNode::MAX_CHILDREN],
value: Some(Data(Vec::new())),
children_encoded: Default::default(),
}
Expand Down Expand Up @@ -370,7 +370,7 @@ impl Storable for Node {
NodeTypeId::Branch => {
// TODO: add path
// TODO: figure out what this size is?
let branch_header_size = MAX_CHILDREN as u64 * 8 + 4;
let branch_header_size = BranchNode::MAX_CHILDREN as u64 * 8 + 4;
let node_raw = mem.get_view(addr + Meta::SIZE, branch_header_size).ok_or(
ShaleError::InvalidCacheView {
offset: addr + Meta::SIZE,
Expand All @@ -379,7 +379,7 @@ impl Storable for Node {
)?;

let mut cur = Cursor::new(node_raw.as_deref());
let mut chd = [None; MAX_CHILDREN];
let mut chd = [None; BranchNode::MAX_CHILDREN];
let mut buff = [0; 8];

for chd in chd.iter_mut() {
Expand All @@ -392,12 +392,13 @@ impl Storable for Node {

cur.read_exact(&mut buff[..4])?;

let raw_len =
u32::from_le_bytes(buff[..4].try_into().expect("invalid slice")) as u64;
let raw_len = u32::from_le_bytes(buff[..4].try_into().expect("invalid slice"));

let value = if raw_len == u32::MAX as u64 {
let value = if raw_len == u32::MAX {
None
} else {
let raw_len = raw_len as u64;

Some(Data(
mem.get_view(addr + Meta::SIZE + branch_header_size as usize, raw_len)
.ok_or(ShaleError::InvalidCacheView {
Expand All @@ -408,9 +409,10 @@ impl Storable for Node {
))
};

let mut chd_encoded: [Option<Vec<u8>>; MAX_CHILDREN] = Default::default();
let mut chd_encoded: [Option<Vec<u8>>; BranchNode::MAX_CHILDREN] =
Default::default();

let offset = if raw_len == u32::MAX as u64 {
let offset = if raw_len == u32::MAX {
addr + Meta::SIZE + branch_header_size as usize
} else {
addr + Meta::SIZE + branch_header_size as usize + raw_len as usize
Expand Down Expand Up @@ -598,20 +600,7 @@ impl Storable for Node {
+ match &self.inner {
NodeType::Branch(n) => {
// TODO: add path
let mut encoded_len = 0;
for emcoded in n.children_encoded.iter() {
encoded_len += match emcoded {
Some(v) => 1 + v.len() as u64,
None => 1,
}
}
MAX_CHILDREN as u64 * 8
+ 4
+ match &n.value {
Some(val) => val.len() as u64,
None => 0,
}
+ encoded_len
n.serialized_len()
}
NodeType::Extension(n) => {
1 + 8
Expand Down Expand Up @@ -654,36 +643,9 @@ impl Storable for Node {
// TODO: add path
cur.write_all(&[type_id::NodeTypeId::Branch as u8]).unwrap();

for c in n.children.iter() {
cur.write_all(&match c {
Some(p) => p.to_le_bytes(),
None => 0u64.to_le_bytes(),
})?;
}
let pos = cur.position() as usize;

match &n.value {
Some(val) => {
cur.write_all(&(val.len() as u32).to_le_bytes())?;
cur.write_all(val)?
}
None => {
cur.write_all(&u32::MAX.to_le_bytes())?;
}
}

// Since child encoding will only be unset after initialization (only used for range proof),
// it is fine to encode its value adjacent to other fields. Same for extention node.
for encoded in n.children_encoded.iter() {
match encoded {
Some(v) => {
cur.write_all(&[v.len() as u8])?;
cur.write_all(v)?
}
None => cur.write_all(&0u8.to_le_bytes())?,
}
}

Ok(())
n.serialize(&mut cur.get_mut()[pos..])
}

NodeType::Extension(n) => {
Expand Down Expand Up @@ -734,8 +696,8 @@ pub(super) mod tests {
value: Option<Vec<u8>>,
repeated_encoded_child: Option<Vec<u8>>,
) -> Node {
let children: [Option<DiskAddress>; MAX_CHILDREN] = from_fn(|i| {
if i < MAX_CHILDREN / 2 {
let children: [Option<DiskAddress>; BranchNode::MAX_CHILDREN] = from_fn(|i| {
if i < BranchNode::MAX_CHILDREN / 2 {
DiskAddress::from(repeated_disk_address).into()
} else {
None
Expand All @@ -745,7 +707,7 @@ pub(super) mod tests {
let children_encoded = repeated_encoded_child
.map(|child| {
from_fn(|i| {
if i < MAX_CHILDREN / 2 {
if i < BranchNode::MAX_CHILDREN / 2 {
child.clone().into()
} else {
None
Expand Down
Loading

0 comments on commit 0dceb44

Please sign in to comment.