From e36b312aaa1c7ddc5a099938b0115d3dc071150d Mon Sep 17 00:00:00 2001 From: Richard Pringle Date: Wed, 29 Nov 2023 18:00:41 -0500 Subject: [PATCH] Implement Storable for LeafNode (#376) --- firewood/src/merkle/node.rs | 79 +++------------------- firewood/src/merkle/node/extension.rs | 53 ++++++++++++++- firewood/src/merkle/node/leaf.rs | 95 ++++++++++++++++++++++++++- 3 files changed, 154 insertions(+), 73 deletions(-) diff --git a/firewood/src/merkle/node.rs b/firewood/src/merkle/node.rs index 789a772b7..ea290318a 100644 --- a/firewood/src/merkle/node.rs +++ b/firewood/src/merkle/node.rs @@ -30,7 +30,7 @@ pub use partial_path::PartialPath; use crate::merkle::to_nibble_array; use crate::nibbles::Nibbles; -use super::{from_nibbles, TrieHash, TRIE_HASH_LEN}; +use super::{TrieHash, TRIE_HASH_LEN}; bitflags! { // should only ever be the size of a nibble @@ -462,48 +462,8 @@ impl Storable for Node { } NodeTypeId::Leaf => { - let leaf_header_size = 1 + 4; - let node_raw = mem.get_view(addr + Meta::SIZE, leaf_header_size).ok_or( - ShaleError::InvalidCacheView { - offset: addr + Meta::SIZE, - size: leaf_header_size, - }, - )?; - - let mut cur = Cursor::new(node_raw.as_deref()); - let mut buff = [0; 4]; - cur.read_exact(&mut buff[..1])?; - - let path_len = buff[0] as u64; - - cur.read_exact(&mut buff)?; - - let data_len = u32::from_le_bytes(buff) as u64; - let remainder = mem - .get_view( - addr + Meta::SIZE + leaf_header_size as usize, - path_len + data_len, - ) - .ok_or(ShaleError::InvalidCacheView { - offset: addr + Meta::SIZE + leaf_header_size as usize, - size: path_len + data_len, - })?; - - let nibbles: Vec<_> = remainder - .as_deref() - .into_iter() - .take(path_len as usize) - .flat_map(to_nibble_array) - .collect(); - - let (path, _) = PartialPath::decode(&nibbles); - let data = Data(remainder.as_deref()[path_len as usize..].to_vec()); - - let node = Self::new_from_hash( - root_hash, - is_encoded_longer_than_hash_len, - NodeType::Leaf(LeafNode { path, data }), - ); + let inner = NodeType::Leaf(LeafNode::deserialize(offset, mem)?); + let node = Self::new_from_hash(root_hash, is_encoded_longer_than_hash_len, inner); Ok(node) } @@ -517,15 +477,8 @@ impl Storable for Node { // TODO: add path n.serialized_len() } - NodeType::Extension(n) => { - 1 + 8 - + n.path.serialized_len() - + match n.chd_encoded() { - Some(v) => 1 + v.len() as u64, - None => 1, - } - } - NodeType::Leaf(n) => 1 + 4 + n.path.serialized_len() + n.data.len() as u64, + NodeType::Extension(n) => n.serialized_len(), + NodeType::Leaf(n) => n.serialized_len(), } } @@ -556,7 +509,7 @@ impl Storable for Node { match &self.inner { NodeType::Branch(n) => { // TODO: add path - cur.write_all(&[type_id::NodeTypeId::Branch as u8]).unwrap(); + cur.write_all(&[type_id::NodeTypeId::Branch as u8])?; let pos = cur.position() as usize; @@ -566,29 +519,17 @@ impl Storable for Node { NodeType::Extension(n) => { cur.write_all(&[type_id::NodeTypeId::Extension as u8])?; - let path: Vec = from_nibbles(&n.path.encode(false)).collect(); - - cur.write_all(&[path.len() as u8])?; - cur.write_all(&n.child.to_le_bytes())?; - cur.write_all(&path)?; - - if let Some(encoded) = n.chd_encoded() { - cur.write_all(&[encoded.len() as u8])?; - cur.write_all(encoded)?; - } + let pos = cur.position() as usize; - Ok(()) + n.serialize(&mut cur.get_mut()[pos..]) } NodeType::Leaf(n) => { cur.write_all(&[type_id::NodeTypeId::Leaf as u8])?; - let path: Vec = from_nibbles(&n.path.encode(true)).collect(); + let pos = cur.position() as usize; - cur.write_all(&[path.len() as u8])?; - cur.write_all(&(n.data.len() as u32).to_le_bytes())?; - cur.write_all(&path)?; - cur.write_all(&n.data).map_err(ShaleError::Io) + n.serialize(&mut cur.get_mut()[pos..]) } } } diff --git a/firewood/src/merkle/node/extension.rs b/firewood/src/merkle/node/extension.rs index fcbcc4e74..d321268ff 100644 --- a/firewood/src/merkle/node/extension.rs +++ b/firewood/src/merkle/node/extension.rs @@ -6,9 +6,15 @@ use bincode::Options; use super::{Encoded, Node}; use crate::{ merkle::{from_nibbles, PartialPath, TRIE_HASH_LEN}, - shale::{DiskAddress, ShaleStore}, + shale::{DiskAddress, ShaleStore, Storable}, }; -use std::fmt::{Debug, Error as FmtError, Formatter}; +use std::{ + fmt::{Debug, Error as FmtError, Formatter}, + io::{Cursor, Write}, + mem::size_of, +}; + +type DataLen = u8; #[derive(PartialEq, Eq, Clone)] pub struct ExtNode { @@ -87,3 +93,46 @@ impl ExtNode { &mut self.child_encoded } } + +impl Storable for ExtNode { + fn serialized_len(&self) -> u64 { + let path_len_size = size_of::() as u64; + let path_len = self.path.serialized_len(); + let child_len = DiskAddress::MSIZE; + let encoded_len_size = size_of::() as u64; + let encoded_len = self + .child_encoded + .as_ref() + .map(|v| v.len() as u64) + .unwrap_or(0); + + path_len_size + path_len + child_len + encoded_len_size + encoded_len + } + + fn serialize(&self, to: &mut [u8]) -> Result<(), crate::shale::ShaleError> { + let mut cursor = Cursor::new(to); + + let path: Vec = from_nibbles(&self.path.encode(false)).collect(); + + cursor.write_all(&[path.len() as DataLen])?; + cursor.write_all(&self.child.to_le_bytes())?; + cursor.write_all(&path)?; + + if let Some(encoded) = self.chd_encoded() { + cursor.write_all(&[encoded.len() as DataLen])?; + cursor.write_all(encoded)?; + } + + Ok(()) + } + + fn deserialize( + _addr: usize, + _mem: &T, + ) -> Result + where + Self: Sized, + { + todo!() + } +} diff --git a/firewood/src/merkle/node/leaf.rs b/firewood/src/merkle/node/leaf.rs index 475f62a00..6235b2e2e 100644 --- a/firewood/src/merkle/node/leaf.rs +++ b/firewood/src/merkle/node/leaf.rs @@ -1,15 +1,25 @@ // Copyright (C) 2023, Ava Labs, Inc. All rights reserved. // See the file LICENSE.md for licensing terms. -use std::fmt::{Debug, Error as FmtError, Formatter}; +use std::{ + fmt::{Debug, Error as FmtError, Formatter}, + io::{Cursor, Read, Write}, + mem::size_of, +}; use bincode::Options; use super::{Data, Encoded}; -use crate::merkle::{from_nibbles, PartialPath}; +use crate::{ + merkle::{from_nibbles, to_nibble_array, PartialPath}, + shale::{ShaleError::InvalidCacheView, Storable}, +}; pub const SIZE: usize = 2; +type PathLen = u8; +type DataLen = u32; + #[derive(PartialEq, Eq, Clone)] pub struct LeafNode { pub(crate) path: PartialPath, @@ -23,6 +33,9 @@ impl Debug for LeafNode { } impl LeafNode { + const PATH_LEN_SIZE: u64 = size_of::() as u64; + const DATA_LEN_SIZE: u64 = size_of::() as u64; + pub fn new, D: Into>(path: P, data: D) -> Self { Self { path: path.into(), @@ -51,6 +64,84 @@ impl LeafNode { } } +impl Storable for LeafNode { + fn serialized_len(&self) -> u64 { + let path_len_size = size_of::() as u64; + let path_len = self.path.serialized_len(); + let data_len_size = size_of::() as u64; + let data_len = self.data.len() as u64; + + path_len_size + path_len + data_len_size + data_len + } + + fn serialize(&self, to: &mut [u8]) -> Result<(), crate::shale::ShaleError> { + let mut cursor = Cursor::new(to); + + let path: Vec = from_nibbles(&self.path.encode(true)).collect(); + + cursor.write_all(&[path.len() as PathLen])?; + + let data_len = self.data.len() as DataLen; + cursor.write_all(&data_len.to_le_bytes())?; + + cursor.write_all(&path)?; + cursor.write_all(&self.data)?; + + Ok(()) + } + + fn deserialize( + mut offset: usize, + mem: &T, + ) -> Result + where + Self: Sized, + { + let header_size = Self::PATH_LEN_SIZE + Self::DATA_LEN_SIZE; + + let node_header_raw = mem + .get_view(offset, header_size) + .ok_or(InvalidCacheView { + offset, + size: header_size, + })? + .as_deref(); + + offset += header_size as usize; + + let mut cursor = Cursor::new(node_header_raw); + let mut buf = [0u8; Self::DATA_LEN_SIZE as usize]; + + let path_len = { + let buf = &mut buf[..Self::PATH_LEN_SIZE as usize]; + cursor.read_exact(buf)?; + buf[0] as u64 + }; + + let data_len = { + cursor.read_exact(buf.as_mut())?; + DataLen::from_le_bytes(buf) as u64 + }; + + let size = path_len + data_len; + let remainder = mem + .get_view(offset, size) + .ok_or(InvalidCacheView { offset, size })? + .as_deref(); + + let (path, data) = remainder.split_at(path_len as usize); + + let path = { + let nibbles: Vec = path.iter().copied().flat_map(to_nibble_array).collect(); + PartialPath::decode(&nibbles).0 + }; + + let data = Data(data.to_vec()); + + Ok(Self::new(path, data)) + } +} + #[cfg(test)] mod tests { use super::*;