diff --git a/verkle/src/nodes/branch.rs b/verkle/src/nodes/branch.rs index fcb5b41..e50b882 100644 --- a/verkle/src/nodes/branch.rs +++ b/verkle/src/nodes/branch.rs @@ -1,4 +1,4 @@ -use std::{array, collections::BTreeMap, ops::DerefMut}; +use std::collections::BTreeMap; use alloy_primitives::B256; use anyhow::Result; @@ -7,53 +7,73 @@ use ssz::{Decode, Encode}; use crate::{ committer::DEFAULT_COMMITER, - constants::VERKLE_NODE_WIDTH, utils::{b256_to_fr, fr_to_b256}, Db, TrieKey, TrieValue, }; -use super::{node::NodeTrait, Node}; +use super::{node::NodeTrait, LeafNode, Node}; pub struct BranchNode { - values: Box<[Node; VERKLE_NODE_WIDTH]>, - cp: Element, + values: BTreeMap, + commitment: Element, } impl BranchNode { pub fn new() -> Self { Self { - values: array::from_fn(|_| Node::Empty).into(), - cp: Element::zero(), + values: BTreeMap::new(), + commitment: Element::zero(), } } - pub fn set(&mut self, index: usize, node: Node) { - let node_at_index = &mut self.values[index]; - let pre_commitment = node_at_index.commit(); - *node_at_index = node; - let post_commitment = node_at_index.commit(); - self.cp += DEFAULT_COMMITER.scalar_mul(index, post_commitment - pre_commitment); + pub fn set(&mut self, index: u8, node: Node) { + let old_node = self.values.insert(index, node); + self.update_commitment( + index, + old_node + .map(|node| node.hash_commitment()) + .unwrap_or_default(), + ); } - pub(super) fn get_mut(&mut self, index: usize) -> &mut Node { - &mut self.values[index] + pub(super) fn get_mut(&mut self, index: u8) -> Option<&mut Node> { + self.values.get_mut(&index) } pub fn insert(&mut self, depth: usize, key: TrieKey, value: TrieValue, db: &Db) -> Result<()> { - let index = key[depth] as usize; - let node = &mut self.values[index]; - let pre_commitment = node.commit(); - node.insert(depth + 1, key, value, db)?; - let post_commitment = node.commit(); - self.cp += DEFAULT_COMMITER.scalar_mul(index, post_commitment - pre_commitment); + let index = key[depth]; + let pre_commitment = self.get_child_commit(index); + match self.values.get_mut(&index) { + Some(node) => { + node.insert(depth + 1, key, value, db)?; + node.hash_commitment_mut(); + } + None => { + self.values + .insert(index, Node::Leaf(LeafNode::new_for_key_value(&key, value))); + } + }; + self.update_commitment(index, pre_commitment); Ok(()) } + fn get_child_commit(&mut self, index: u8) -> Fr { + self.values + .get_mut(&index) + .map(|node| node.hash_commitment_mut()) + .unwrap_or_default() + } + + fn update_commitment(&mut self, index: u8, pre_commitment: Fr) { + let post_commitment = self.get_child_commit(index); + self.commitment += DEFAULT_COMMITER.scalar_mul(index as usize, post_commitment - pre_commitment); + } + pub fn write_and_commit(&mut self, db: &mut Db) -> Result { - for node in self.values.deref_mut() { + for (_, node) in self.values.iter_mut() { node.write_and_commit(db)?; } - Ok(self.commit()) + Ok(self.hash_commitment_mut()) } } @@ -65,7 +85,7 @@ impl Default for BranchNode { impl NodeTrait for BranchNode { fn hash_commitment(&self) -> Fr { - self.cp.map_to_scalar_field() + self.commitment.map_to_scalar_field() } } @@ -78,14 +98,7 @@ impl Encode for BranchNode { let commitments: BTreeMap = self .values .iter() - .enumerate() - .filter_map(|(index, node)| { - if node.is_empty() { - None - } else { - Some((index as u8, fr_to_b256(&node.hash_commitment()))) - } - }) + .map(|(index, node)| (*index, fr_to_b256(&node.hash_commitment()))) .collect(); commitments.ssz_append(buf); } @@ -106,22 +119,19 @@ impl Decode for BranchNode { fn from_ssz_bytes(bytes: &[u8]) -> Result { let commitments = BTreeMap::::from_ssz_bytes(bytes)?; - let commitments: BTreeMap = commitments + + let values = commitments .iter() - .map(|(index, commitment)| (*index as usize, b256_to_fr(commitment))) + .map(|(index, c)| (*index, Node::Commitment(b256_to_fr(c)))) .collect(); - let values = array::from_fn(|i| { + let cp = DEFAULT_COMMITER.commit_sparse( commitments - .get(&i) - .map(|c| Node::Commitment(*c)) - .unwrap_or_else(|| Node::Empty) - }); - let cp = DEFAULT_COMMITER.commit_sparse(commitments.into_iter().collect()); - - Ok(Self { - values: values.into(), - cp, - }) + .iter() + .map(|(index, commitment)| (*index as usize, b256_to_fr(commitment))) + .collect(), + ); + + Ok(Self { values, commitment: cp }) } } diff --git a/verkle/src/nodes/leaf.rs b/verkle/src/nodes/leaf.rs index 9f938d5..d77a9b5 100644 --- a/verkle/src/nodes/leaf.rs +++ b/verkle/src/nodes/leaf.rs @@ -27,14 +27,14 @@ pub struct LeafNode { values: BTreeMap, #[ssz(skip_serializing)] - cp1: Element, + c1: Element, #[ssz(skip_serializing)] - cp2: Element, + c2: Element, #[ssz(skip_serializing)] - const_cp: Element, + const_c: Element, #[ssz(skip_serializing)] - c: Option, + hash_commitment: Option, } impl LeafNode { @@ -46,10 +46,10 @@ impl LeafNode { Self { stem, values: BTreeMap::new(), - cp1: Element::zero(), - cp2: Element::zero(), - const_cp: const_c, - c: None, + c1: Element::zero(), + c2: Element::zero(), + const_c, + hash_commitment: None, } } @@ -64,10 +64,10 @@ impl LeafNode { } fn calculate_commitment(&self) -> Element { - self.const_cp + self.const_c + DEFAULT_COMMITER.commit_sparse(vec![ - (2, self.cp1.map_to_scalar_field()), - (3, self.cp2.map_to_scalar_field()), + (2, self.c1.map_to_scalar_field()), + (3, self.c2.map_to_scalar_field()), ]) } @@ -95,11 +95,11 @@ impl LeafNode { + CRS[high_index] * (value_high_16 - old_value_high_16); if index < VERKLE_NODE_WIDTH / 2 { - self.cp1 += diff; + self.c1 += diff; } else { - self.cp2 += diff; + self.c2 += diff; }; - self.c = None; + self.hash_commitment = None; } pub fn set_all(&mut self, values: impl IntoIterator) { @@ -119,13 +119,13 @@ impl LeafNode { impl NodeTrait for LeafNode { fn hash_commitment(&self) -> Fr { - self.c + self.hash_commitment .unwrap_or_else(|| self.calculate_commitment().map_to_scalar_field()) } - fn commit(&mut self) -> Fr { - self.c = Some(self.hash_commitment()); - self.c.expect("Value must be present") + fn hash_commitment_mut(&mut self) -> Fr { + self.hash_commitment = Some(self.hash_commitment()); + self.hash_commitment.expect("Value must be present") } } @@ -163,7 +163,7 @@ mod tests { let mut leaf = LeafNode::new_for_key_value(&key, TrieValue::ZERO); assert_eq!( - fr_to_b256(&leaf.commit()).to_string(), + fr_to_b256(&leaf.hash_commitment_mut()).to_string(), "0x1c0727f0c6c9887189f75a9d08b804aba20892a238e147750767eac22a830d08" ); } @@ -174,7 +174,7 @@ mod tests { let mut leaf = LeafNode::new_for_key_value(&key, TrieValue::from(1)); assert_eq!( - fr_to_b256(&leaf.commit()).to_string(), + fr_to_b256(&leaf.hash_commitment_mut()).to_string(), "0x6ef020caaeda01ff573afe6df6460d4aae14b4987e02ea39074f270ce62dfc14" ); } @@ -189,7 +189,7 @@ mod tests { let mut leaf = LeafNode::new_for_key_value(&key, TrieValue::from_le_bytes(bytes)); assert_eq!( - fr_to_b256(&leaf.commit()).to_string(), + fr_to_b256(&leaf.hash_commitment_mut()).to_string(), "0xb897ba52c5317acd75f5f3c3922f461357d4fb8b685fe63f20a3b2adb014370a" ); } @@ -231,7 +231,7 @@ mod tests { ); assert_eq!( - fr_to_b256(&leaf.commit()).to_string(), + fr_to_b256(&leaf.hash_commitment_mut()).to_string(), "0xcc30be1f0d50eacfacaa3361b8df4d2014a849854a6cf35e6c55e07d6963f519" ); } diff --git a/verkle/src/nodes/node.rs b/verkle/src/nodes/node.rs index 3702bcd..c0bd206 100644 --- a/verkle/src/nodes/node.rs +++ b/verkle/src/nodes/node.rs @@ -2,7 +2,7 @@ use std::mem; use alloy_primitives::B256; use anyhow::{anyhow, bail, Result}; -use banderwagon::{Fr, Zero}; +use banderwagon::Fr; use ssz::{Decode, Encode}; use crate::{Db, TrieKey, TrieValue}; @@ -12,13 +12,12 @@ use super::{BranchNode, LeafNode}; pub trait NodeTrait { fn hash_commitment(&self) -> Fr; - fn commit(&mut self) -> Fr { + fn hash_commitment_mut(&mut self) -> Fr { self.hash_commitment() } } pub enum Node { - Empty, Branch(BranchNode), Leaf(LeafNode), Commitment(Fr), @@ -27,26 +26,24 @@ pub enum Node { impl NodeTrait for Node { fn hash_commitment(&self) -> Fr { match self { - Node::Empty => Fr::zero(), Node::Branch(branch_node) => branch_node.hash_commitment(), Node::Leaf(leaf_node) => leaf_node.hash_commitment(), Node::Commitment(c) => *c, } } - fn commit(&mut self) -> Fr { + fn hash_commitment_mut(&mut self) -> Fr { match self { - Node::Empty => Fr::zero(), - Node::Branch(branch_node) => branch_node.commit(), - Node::Leaf(leaf_node) => leaf_node.commit(), + Node::Branch(branch_node) => branch_node.hash_commitment_mut(), + Node::Leaf(leaf_node) => leaf_node.hash_commitment_mut(), Node::Commitment(c) => *c, } } } impl Node { - pub fn is_empty(&self) -> bool { - matches!(self, Node::Empty) + pub fn new() -> Self { + Self::Branch(BranchNode::new()) } pub fn check(&self, commitment: &Fr) -> Result<()> { @@ -65,9 +62,11 @@ impl Node { let mut node = self; loop { match node { - Node::Empty => return Ok(None), Node::Branch(branch_node) => { - node = branch_node.get_mut(key[depth] as usize); + node = match branch_node.get_mut(key[depth]) { + Some(node) => node, + None => return Ok(None), + }; depth += 1; } Node::Leaf(leaf_node) => { @@ -92,7 +91,6 @@ impl Node { pub fn insert(&mut self, depth: usize, key: TrieKey, value: TrieValue, db: &Db) -> Result<()> { match self { - Node::Empty => *self = Node::Leaf(LeafNode::new_for_key_value(&key, value)), Node::Branch(branch_node) => branch_node.insert(depth, key, value, db)?, Node::Leaf(leaf_node) => { if leaf_node.stem() == &key.stem() { @@ -100,7 +98,7 @@ impl Node { } else { let mut branch_node = BranchNode::new(); branch_node.set( - leaf_node.stem()[depth] as usize, + leaf_node.stem()[depth], Node::Leaf(mem::replace( leaf_node, LeafNode::new(TrieKey(B256::ZERO).stem()), @@ -129,22 +127,28 @@ impl Node { match self { Node::Branch(branch_node) => { branch_node.write_and_commit(db)?; - let c = branch_node.commit(); + let c = branch_node.hash_commitment_mut(); db.write(c, self.as_ssz_bytes())?; *self = Node::Commitment(c); Ok(c) } Node::Leaf(leaf_node) => { - let c = leaf_node.commit(); + let c = leaf_node.hash_commitment_mut(); db.write(c, self.as_ssz_bytes())?; *self = Node::Commitment(c); Ok(c) } - _ => Ok(self.commit()), + _ => Ok(self.hash_commitment_mut()), } } } +impl Default for Node { + fn default() -> Self { + Self::new() + } +} + const SSZ_TAG_BRANCH: u8 = 1; const SSZ_TAG_LEAF: u8 = 0; @@ -163,7 +167,6 @@ impl Encode for Node { buf.push(SSZ_TAG_LEAF); leaf_node.ssz_append(buf); } - Node::Empty => panic!("Can't encode Empty node"), Node::Commitment(_) => panic!("Can't encode Commitment node"), } } @@ -172,7 +175,6 @@ impl Encode for Node { match self { Node::Branch(branch_node) => 1 + branch_node.ssz_bytes_len(), Node::Leaf(leaf_node) => 1 + leaf_node.ssz_bytes_len(), - Node::Empty => panic!("Can't encode Empty node"), Node::Commitment(_) => panic!("Can't encode Commitment node"), } } diff --git a/verkle/src/trie.rs b/verkle/src/trie.rs index be601dc..cfdf969 100644 --- a/verkle/src/trie.rs +++ b/verkle/src/trie.rs @@ -4,7 +4,7 @@ use banderwagon::Fr; use crate::{ account::AccountStorageLayout, - nodes::{BranchNode, Node}, + nodes::Node, utils::{b256_to_fr, fr_to_b256}, Db, TrieKey, TrieValue, }; @@ -17,7 +17,7 @@ pub struct Trie { impl Trie { pub fn new(db: Box) -> Self { Self { - root: Node::Branch(BranchNode::new()), + root: Node::new(), db, } }