Skip to content

Commit 33ac146

Browse files
morph-devmorph-dev
committed
refactor: verkle code cleanup
1 parent 377ed98 commit 33ac146

4 files changed

Lines changed: 100 additions & 88 deletions

File tree

‎verkle/src/nodes/branch.rs‎

Lines changed: 55 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use std::{array, collections::BTreeMap, ops::DerefMut};
1+
use std::collections::BTreeMap;
22

33
use alloy_primitives::B256;
44
use anyhow::Result;
@@ -7,53 +7,73 @@ use ssz::{Decode, Encode};
77

88
use crate::{
99
committer::DEFAULT_COMMITER,
10-
constants::VERKLE_NODE_WIDTH,
1110
utils::{b256_to_fr, fr_to_b256},
1211
Db, TrieKey, TrieValue,
1312
};
1413

15-
use super::{node::NodeTrait, Node};
14+
use super::{node::NodeTrait, LeafNode, Node};
1615

1716
pub struct BranchNode {
18-
values: Box<[Node; VERKLE_NODE_WIDTH]>,
19-
cp: Element,
17+
values: BTreeMap<u8, Node>,
18+
commitment: Element,
2019
}
2120

2221
impl BranchNode {
2322
pub fn new() -> Self {
2423
Self {
25-
values: array::from_fn(|_| Node::Empty).into(),
26-
cp: Element::zero(),
24+
values: BTreeMap::new(),
25+
commitment: Element::zero(),
2726
}
2827
}
2928

30-
pub fn set(&mut self, index: usize, node: Node) {
31-
let node_at_index = &mut self.values[index];
32-
let pre_commitment = node_at_index.commit();
33-
*node_at_index = node;
34-
let post_commitment = node_at_index.commit();
35-
self.cp += DEFAULT_COMMITER.scalar_mul(index, post_commitment - pre_commitment);
29+
pub fn set(&mut self, index: u8, node: Node) {
30+
let old_node = self.values.insert(index, node);
31+
self.update_commitment(
32+
index,
33+
old_node
34+
.map(|node| node.hash_commitment())
35+
.unwrap_or_default(),
36+
);
3637
}
3738

38-
pub(super) fn get_mut(&mut self, index: usize) -> &mut Node {
39-
&mut self.values[index]
39+
pub(super) fn get_mut(&mut self, index: u8) -> Option<&mut Node> {
40+
self.values.get_mut(&index)
4041
}
4142

4243
pub fn insert(&mut self, depth: usize, key: TrieKey, value: TrieValue, db: &Db) -> Result<()> {
43-
let index = key[depth] as usize;
44-
let node = &mut self.values[index];
45-
let pre_commitment = node.commit();
46-
node.insert(depth + 1, key, value, db)?;
47-
let post_commitment = node.commit();
48-
self.cp += DEFAULT_COMMITER.scalar_mul(index, post_commitment - pre_commitment);
44+
let index = key[depth];
45+
let pre_commitment = self.get_child_commit(index);
46+
match self.values.get_mut(&index) {
47+
Some(node) => {
48+
node.insert(depth + 1, key, value, db)?;
49+
node.hash_commitment_mut();
50+
}
51+
None => {
52+
self.values
53+
.insert(index, Node::Leaf(LeafNode::new_for_key_value(&key, value)));
54+
}
55+
};
56+
self.update_commitment(index, pre_commitment);
4957
Ok(())
5058
}
5159

60+
fn get_child_commit(&mut self, index: u8) -> Fr {
61+
self.values
62+
.get_mut(&index)
63+
.map(|node| node.hash_commitment_mut())
64+
.unwrap_or_default()
65+
}
66+
67+
fn update_commitment(&mut self, index: u8, pre_commitment: Fr) {
68+
let post_commitment = self.get_child_commit(index);
69+
self.commitment += DEFAULT_COMMITER.scalar_mul(index as usize, post_commitment - pre_commitment);
70+
}
71+
5272
pub fn write_and_commit(&mut self, db: &mut Db) -> Result<Fr> {
53-
for node in self.values.deref_mut() {
73+
for (_, node) in self.values.iter_mut() {
5474
node.write_and_commit(db)?;
5575
}
56-
Ok(self.commit())
76+
Ok(self.hash_commitment_mut())
5777
}
5878
}
5979

@@ -65,7 +85,7 @@ impl Default for BranchNode {
6585

6686
impl NodeTrait for BranchNode {
6787
fn hash_commitment(&self) -> Fr {
68-
self.cp.map_to_scalar_field()
88+
self.commitment.map_to_scalar_field()
6989
}
7090
}
7191

@@ -78,14 +98,7 @@ impl Encode for BranchNode {
7898
let commitments: BTreeMap<u8, B256> = self
7999
.values
80100
.iter()
81-
.enumerate()
82-
.filter_map(|(index, node)| {
83-
if node.is_empty() {
84-
None
85-
} else {
86-
Some((index as u8, fr_to_b256(&node.hash_commitment())))
87-
}
88-
})
101+
.map(|(index, node)| (*index, fr_to_b256(&node.hash_commitment())))
89102
.collect();
90103
commitments.ssz_append(buf);
91104
}
@@ -106,22 +119,19 @@ impl Decode for BranchNode {
106119

107120
fn from_ssz_bytes(bytes: &[u8]) -> Result<Self, ssz::DecodeError> {
108121
let commitments = BTreeMap::<u8, B256>::from_ssz_bytes(bytes)?;
109-
let commitments: BTreeMap<usize, Fr> = commitments
122+
123+
let values = commitments
110124
.iter()
111-
.map(|(index, commitment)| (*index as usize, b256_to_fr(commitment)))
125+
.map(|(index, c)| (*index, Node::Commitment(b256_to_fr(c))))
112126
.collect();
113127

114-
let values = array::from_fn(|i| {
128+
let cp = DEFAULT_COMMITER.commit_sparse(
115129
commitments
116-
.get(&i)
117-
.map(|c| Node::Commitment(*c))
118-
.unwrap_or_else(|| Node::Empty)
119-
});
120-
let cp = DEFAULT_COMMITER.commit_sparse(commitments.into_iter().collect());
121-
122-
Ok(Self {
123-
values: values.into(),
124-
cp,
125-
})
130+
.iter()
131+
.map(|(index, commitment)| (*index as usize, b256_to_fr(commitment)))
132+
.collect(),
133+
);
134+
135+
Ok(Self { values, commitment: cp })
126136
}
127137
}

‎verkle/src/nodes/leaf.rs‎

Lines changed: 22 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -27,14 +27,14 @@ pub struct LeafNode {
2727
values: BTreeMap<u8, TrieValue>,
2828

2929
#[ssz(skip_serializing)]
30-
cp1: Element,
30+
c1: Element,
3131
#[ssz(skip_serializing)]
32-
cp2: Element,
32+
c2: Element,
3333

3434
#[ssz(skip_serializing)]
35-
const_cp: Element,
35+
const_c: Element,
3636
#[ssz(skip_serializing)]
37-
c: Option<Fr>,
37+
hash_commitment: Option<Fr>,
3838
}
3939

4040
impl LeafNode {
@@ -46,10 +46,10 @@ impl LeafNode {
4646
Self {
4747
stem,
4848
values: BTreeMap::new(),
49-
cp1: Element::zero(),
50-
cp2: Element::zero(),
51-
const_cp: const_c,
52-
c: None,
49+
c1: Element::zero(),
50+
c2: Element::zero(),
51+
const_c,
52+
hash_commitment: None,
5353
}
5454
}
5555

@@ -64,10 +64,10 @@ impl LeafNode {
6464
}
6565

6666
fn calculate_commitment(&self) -> Element {
67-
self.const_cp
67+
self.const_c
6868
+ DEFAULT_COMMITER.commit_sparse(vec![
69-
(2, self.cp1.map_to_scalar_field()),
70-
(3, self.cp2.map_to_scalar_field()),
69+
(2, self.c1.map_to_scalar_field()),
70+
(3, self.c2.map_to_scalar_field()),
7171
])
7272
}
7373

@@ -95,11 +95,11 @@ impl LeafNode {
9595
+ CRS[high_index] * (value_high_16 - old_value_high_16);
9696

9797
if index < VERKLE_NODE_WIDTH / 2 {
98-
self.cp1 += diff;
98+
self.c1 += diff;
9999
} else {
100-
self.cp2 += diff;
100+
self.c2 += diff;
101101
};
102-
self.c = None;
102+
self.hash_commitment = None;
103103
}
104104

105105
pub fn set_all(&mut self, values: impl IntoIterator<Item = (u8, TrieValue)>) {
@@ -119,13 +119,13 @@ impl LeafNode {
119119

120120
impl NodeTrait for LeafNode {
121121
fn hash_commitment(&self) -> Fr {
122-
self.c
122+
self.hash_commitment
123123
.unwrap_or_else(|| self.calculate_commitment().map_to_scalar_field())
124124
}
125125

126-
fn commit(&mut self) -> Fr {
127-
self.c = Some(self.hash_commitment());
128-
self.c.expect("Value must be present")
126+
fn hash_commitment_mut(&mut self) -> Fr {
127+
self.hash_commitment = Some(self.hash_commitment());
128+
self.hash_commitment.expect("Value must be present")
129129
}
130130
}
131131

@@ -163,7 +163,7 @@ mod tests {
163163
let mut leaf = LeafNode::new_for_key_value(&key, TrieValue::ZERO);
164164

165165
assert_eq!(
166-
fr_to_b256(&leaf.commit()).to_string(),
166+
fr_to_b256(&leaf.hash_commitment_mut()).to_string(),
167167
"0x1c0727f0c6c9887189f75a9d08b804aba20892a238e147750767eac22a830d08"
168168
);
169169
}
@@ -174,7 +174,7 @@ mod tests {
174174
let mut leaf = LeafNode::new_for_key_value(&key, TrieValue::from(1));
175175

176176
assert_eq!(
177-
fr_to_b256(&leaf.commit()).to_string(),
177+
fr_to_b256(&leaf.hash_commitment_mut()).to_string(),
178178
"0x6ef020caaeda01ff573afe6df6460d4aae14b4987e02ea39074f270ce62dfc14"
179179
);
180180
}
@@ -189,7 +189,7 @@ mod tests {
189189
let mut leaf = LeafNode::new_for_key_value(&key, TrieValue::from_le_bytes(bytes));
190190

191191
assert_eq!(
192-
fr_to_b256(&leaf.commit()).to_string(),
192+
fr_to_b256(&leaf.hash_commitment_mut()).to_string(),
193193
"0xb897ba52c5317acd75f5f3c3922f461357d4fb8b685fe63f20a3b2adb014370a"
194194
);
195195
}
@@ -231,7 +231,7 @@ mod tests {
231231
);
232232

233233
assert_eq!(
234-
fr_to_b256(&leaf.commit()).to_string(),
234+
fr_to_b256(&leaf.hash_commitment_mut()).to_string(),
235235
"0xcc30be1f0d50eacfacaa3361b8df4d2014a849854a6cf35e6c55e07d6963f519"
236236
);
237237
}

0 commit comments

Comments
 (0)