Skip to content

Commit 2c4357a

Browse files
cratelyncratelyn
committed
tree: 🍉 JellyfishMerkleTree::put_value_sets() uses PutValueSetError 
1 parent 6b05cc9 commit 2c4357a

File tree

1 file changed

+90
-65
lines changed

1 file changed

+90
-65
lines changed

‎src/tree.rs

+90-65
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::{
22
node_type::{Child, Children, InternalNode, LeafNode, Node, NodeKey, NodeType},
33
proof::{definition::UpdateMerkleProof, SparseMerkleLeafNode, SparseMerkleNode},
44
storage::{Node::Leaf, TreeReader, TreeReaderExt, TreeUpdateBatch},
5-
tree_cache::{NodeAlreadyExists, TreeCache},
5+
tree_cache::{FreezeError, NodeAlreadyExists, TreeCache},
66
types::{
77
nibble::{
88
nibble_path::{skip_common_prefix, NibbleIterator, NibblePath},
@@ -13,8 +13,8 @@ use crate::{
1313
},
1414
Bytes32Ext, KeyHash, MissingRootError, OwnedValue, RootHash, SimpleHasher, ValueHash,
1515
};
16-
use alloc::{collections::BTreeMap, format, vec, vec::Vec};
17-
use anyhow::{bail, format_err, Context};
16+
use alloc::{collections::BTreeMap, vec, vec::Vec};
17+
use anyhow::{bail, format_err};
1818
use core::{cmp::Ordering, convert::TryInto, marker::PhantomData};
1919

2020
#[cfg(not(feature = "std"))]
@@ -61,6 +61,24 @@ pub enum LookupError<E> {
6161
CyclicGraphDetected,
6262
}
6363

64+
/// Errors that can occur when a value set is [applied] to a [`JellyfishMerkleTree<'a, R, H>`].
65+
///
66+
/// [applied]: JellyfishMerkleTree::put_value_sets
67+
#[derive(Debug, thiserror::Error)]
68+
#[non_exhaustive]
69+
pub enum PutValueSetError<E> {
70+
#[error(transparent)]
71+
Read(E),
72+
#[error(transparent)]
73+
NodeAlreadyExists(#[from] NodeAlreadyExists),
74+
#[error("Non-root null node exists with node key {key:?}")]
75+
NonRootNullNodeExists { key: NodeKey },
76+
#[error("could not find a node associated with the key `{key:?}`")]
77+
NodeNotFound { key: NodeKey },
78+
#[error(transparent)]
79+
FreezeError(#[from] FreezeError<E>),
80+
}
81+
6482
impl<'a, R, H> JellyfishMerkleTree<'a, R, H>
6583
where
6684
R: 'a + TreeReader,
@@ -97,8 +115,9 @@ where
97115
value_sets: Vec<Vec<(KeyHash, OwnedValue)>>,
98116
node_hashes: Option<Vec<&HashMap<NibblePath, [u8; 32]>>>,
99117
first_version: Version,
100-
) -> Result<(Vec<RootHash>, TreeUpdateBatch), anyhow::Error> {
101-
let mut tree_cache = TreeCache::new(self.reader, first_version)?;
118+
) -> Result<(Vec<RootHash>, TreeUpdateBatch), PutValueSetError<R::Error>> {
119+
let mut tree_cache =
120+
TreeCache::new(self.reader, first_version).map_err(PutValueSetError::Read)?;
102121
let hash_sets: Vec<_> = match node_hashes {
103122
Some(hashes) => hashes.into_iter().map(Some).collect(),
104123
None => (0..value_sets.len()).map(|_| None).collect(),
@@ -148,10 +167,15 @@ where
148167
depth: usize,
149168
hash_cache: &Option<&HashMap<NibblePath, [u8; 32]>>,
150169
tree_cache: &mut TreeCache<R>,
151-
) -> Result<(NodeKey, Node), anyhow::Error> {
170+
) -> Result<(NodeKey, Node), PutValueSetError<R::Error>> {
152171
assert!(!kvs.is_empty());
153172

154-
let node = tree_cache.get_node(&node_key)?;
173+
let node = tree_cache
174+
.get_node_option(&node_key)
175+
.map_err(PutValueSetError::Read)?
176+
.ok_or_else(|| PutValueSetError::NodeNotFound {
177+
key: node_key.clone(),
178+
})?;
155179
Ok(match node {
156180
Node::Internal(internal_node) => {
157181
// We always delete the existing internal node here because it will not be referenced anyway
@@ -224,10 +248,9 @@ where
224248
}
225249
Node::Null => {
226250
if !node_key.nibble_path().is_empty() {
227-
bail!(
228-
"Null node exists for non-root node with node_key {:?}",
229-
node_key
230-
);
251+
return Err(PutValueSetError::NonRootNullNodeExists {
252+
key: node_key.clone(),
253+
});
231254
}
232255

233256
if node_key.version() == version {
@@ -255,7 +278,7 @@ where
255278
depth: usize,
256279
hash_cache: &Option<&HashMap<NibblePath, [u8; 32]>>,
257280
tree_cache: &mut TreeCache<R>,
258-
) -> Result<(NodeKey, Node), anyhow::Error> {
281+
) -> Result<(NodeKey, Node), NodeAlreadyExists> {
259282
let existing_leaf_key = existing_leaf_node.key_hash();
260283

261284
if kvs.len() == 1 && kvs[0].0 == existing_leaf_key {
@@ -325,7 +348,7 @@ where
325348
depth: usize,
326349
hash_cache: &Option<&HashMap<NibblePath, [u8; 32]>>,
327350
tree_cache: &mut TreeCache<R>,
328-
) -> Result<(NodeKey, Node), anyhow::Error> {
351+
) -> Result<(NodeKey, Node), NodeAlreadyExists> {
329352
if kvs.len() == 1 {
330353
let new_leaf_node = Node::Leaf(LeafNode::new(kvs[0].0, kvs[0].1));
331354
tree_cache.put_node(node_key.clone(), new_leaf_node.clone())?;
@@ -366,7 +389,7 @@ where
366389
&self,
367390
value_set: impl IntoIterator<Item = (KeyHash, Option<OwnedValue>)>,
368391
version: Version,
369-
) -> Result<(RootHash, TreeUpdateBatch), anyhow::Error> {
392+
) -> Result<(RootHash, TreeUpdateBatch), PutValueSetError<R::Error>> {
370393
let (root_hashes, tree_update_batch) = self.put_value_sets(vec![value_set], version)?;
371394
assert_eq!(
372395
root_hashes.len(),
@@ -383,7 +406,7 @@ where
383406
&self,
384407
value_set: impl IntoIterator<Item = (KeyHash, Option<OwnedValue>)>,
385408
version: Version,
386-
) -> Result<(RootHash, UpdateMerkleProof<H>, TreeUpdateBatch), anyhow::Error> {
409+
) -> Result<(RootHash, UpdateMerkleProof<H>, TreeUpdateBatch), PutValueSetError<R::Error>> {
387410
let (mut hash_and_proof, batch_update) =
388411
self.put_value_sets_with_proof(vec![value_set], version)?;
389412
assert_eq!(
@@ -442,21 +465,15 @@ where
442465
&self,
443466
value_sets: impl IntoIterator<Item = impl IntoIterator<Item = (KeyHash, Option<OwnedValue>)>>,
444467
first_version: Version,
445-
) -> Result<(Vec<RootHash>, TreeUpdateBatch), anyhow::Error> {
446-
let mut tree_cache = TreeCache::new(self.reader, first_version)?;
468+
) -> Result<(Vec<RootHash>, TreeUpdateBatch), PutValueSetError<R::Error>> {
469+
let mut tree_cache =
470+
TreeCache::new(self.reader, first_version).map_err(PutValueSetError::Read)?;
447471
for (idx, value_set) in value_sets.into_iter().enumerate() {
448472
let version = first_version + idx as u64;
449-
for (i, (key, value)) in value_set.into_iter().enumerate() {
450-
let action = if value.is_some() { "insert" } else { "delete" };
473+
for (key, value) in value_set {
451474
let value_hash = value.as_ref().map(|v| ValueHash::with::<H>(v));
452475
tree_cache.put_value(version, key, value);
453-
self.put(key, value_hash, version, &mut tree_cache, false)
454-
.with_context(|| {
455-
format!(
456-
"failed to {} key {} for version {}, key = {:?}",
457-
action, i, version, key
458-
)
459-
})?;
476+
self.put(key, value_hash, version, &mut tree_cache, false)?;
460477
}
461478

462479
// Freezes the current cache to make all contents in the current cache immutable.
@@ -506,24 +523,19 @@ where
506523
&self,
507524
value_sets: impl IntoIterator<Item = impl IntoIterator<Item = (KeyHash, Option<OwnedValue>)>>,
508525
first_version: Version,
509-
) -> Result<(Vec<(RootHash, UpdateMerkleProof<H>)>, TreeUpdateBatch), anyhow::Error> {
510-
let mut tree_cache = TreeCache::new(self.reader, first_version)?;
526+
) -> Result<(Vec<(RootHash, UpdateMerkleProof<H>)>, TreeUpdateBatch), PutValueSetError<R::Error>>
527+
{
528+
let mut tree_cache =
529+
TreeCache::new(self.reader, first_version).map_err(PutValueSetError::Read)?;
511530
let mut batch_proofs = Vec::new();
512531
for (idx, value_set) in value_sets.into_iter().enumerate() {
513532
let version = first_version + idx as u64;
514533
let mut proofs = Vec::new();
515-
for (i, (key, value)) in value_set.into_iter().enumerate() {
516-
let action = if value.is_some() { "insert" } else { "delete" };
534+
for (key, value) in value_set {
517535
let value_hash = value.as_ref().map(|v| ValueHash::with::<H>(v));
518536
tree_cache.put_value(version, key, value.clone());
519537
let merkle_proof = self
520-
.put(key, value_hash, version, &mut tree_cache, true)
521-
.with_context(|| {
522-
format!(
523-
"failed to {} key {} for version {}, key = {:?}",
524-
action, i, version, key
525-
)
526-
})?
538+
.put(key, value_hash, version, &mut tree_cache, true)?
527539
.unwrap();
528540

529541
proofs.push(merkle_proof);
@@ -552,7 +564,7 @@ where
552564
version: Version,
553565
tree_cache: &mut TreeCache<R>,
554566
with_proof: bool,
555-
) -> Result<Option<SparseMerkleProof<H>>, anyhow::Error> {
567+
) -> Result<Option<SparseMerkleProof<H>>, PutValueSetError<R::Error>> {
556568
// tree_cache.ensure_initialized()?;
557569

558570
let nibble_path = NibblePath::new(key.0.to_vec());
@@ -602,35 +614,43 @@ where
602614
value: Option<ValueHash>,
603615
tree_cache: &mut TreeCache<R>,
604616
with_proof: bool,
605-
) -> Result<(PutResult<(NodeKey, Node)>, Option<SparseMerkleProof<H>>), anyhow::Error> {
617+
) -> Result<
618+
(PutResult<(NodeKey, Node)>, Option<SparseMerkleProof<H>>),
619+
PutValueSetError<R::Error>,
620+
> {
606621
// Because deletions could cause the root node not to exist, we try to get the root node,
607622
// and if it doesn't exist, we synthesize a `Null` node, noting that it hasn't yet been
608623
// committed anywhere (we need to track this because the tree cache will panic if we try to
609624
// delete a node that it doesn't know about).
610625
let (node, node_already_exists) = tree_cache
611-
.get_node_option(&root_node_key)?
626+
.get_node_option(&root_node_key)
627+
.map_err(PutValueSetError::Read)?
612628
.map(|node| (node, true))
613629
.unwrap_or((Node::Null, false));
614630

615631
match node {
616-
Node::Internal(internal_node) => self.insert_at_internal_node(
617-
root_node_key,
618-
internal_node,
619-
version,
620-
nibble_iter,
621-
value,
622-
tree_cache,
623-
with_proof,
624-
),
625-
Node::Leaf(leaf_node) => self.insert_at_leaf_node(
626-
root_node_key,
627-
leaf_node,
628-
version,
629-
nibble_iter,
630-
value,
631-
tree_cache,
632-
with_proof,
633-
),
632+
Node::Internal(internal_node) => self
633+
.insert_at_internal_node(
634+
root_node_key,
635+
internal_node,
636+
version,
637+
nibble_iter,
638+
value,
639+
tree_cache,
640+
with_proof,
641+
)
642+
.map_err(PutValueSetError::from),
643+
Node::Leaf(leaf_node) => self
644+
.insert_at_leaf_node(
645+
root_node_key,
646+
leaf_node,
647+
version,
648+
nibble_iter,
649+
value,
650+
tree_cache,
651+
with_proof,
652+
)
653+
.map_err(PutValueSetError::from),
634654
Node::Null => {
635655
let merkle_proof_null = if with_proof {
636656
Some(SparseMerkleProof::new(None, vec![]))
@@ -639,10 +659,7 @@ where
639659
};
640660

641661
if !root_node_key.nibble_path().is_empty() {
642-
bail!(
643-
"Null node exists for non-root node with node_key {:?}",
644-
root_node_key
645-
);
662+
return Err(PutValueSetError::NonRootNullNodeExists { key: root_node_key });
646663
}
647664
// Delete the old null node if the at the same version
648665
if root_node_key.version() == version && node_already_exists {
@@ -680,7 +697,10 @@ where
680697
value: Option<ValueHash>,
681698
tree_cache: &mut TreeCache<R>,
682699
with_proof: bool,
683-
) -> Result<(PutResult<(NodeKey, Node)>, Option<SparseMerkleProof<H>>), anyhow::Error> {
700+
) -> Result<
701+
(PutResult<(NodeKey, Node)>, Option<SparseMerkleProof<H>>),
702+
PutValueSetError<R::Error>,
703+
> {
684704
// Find the next node to visit following the next nibble as index.
685705
let child_index = nibble_iter.next().expect("Ran out of nibbles");
686706

@@ -813,7 +833,12 @@ where
813833
if it.next().is_none() && child.is_leaf() {
814834
// internal node has only one child left and it's leaf node, replace it with the leaf node
815835
let child_key = node_key.gen_child_node_key(child.version, child_nibble);
816-
let child_node = tree_cache.get_node(&child_key)?;
836+
let child_node = tree_cache
837+
.get_node_option(&child_key)
838+
.map_err(PutValueSetError::Read)?
839+
.ok_or_else(|| PutValueSetError::NodeNotFound {
840+
key: child_key.clone(),
841+
})?;
817842
tree_cache.delete_node(&child_key, true /* is_leaf */);
818843

819844
node_key.set_version(version);
@@ -853,7 +878,7 @@ where
853878
value_hash: Option<ValueHash>,
854879
tree_cache: &mut TreeCache<R>,
855880
with_proof: bool,
856-
) -> Result<(PutResult<(NodeKey, Node)>, Option<SparseMerkleProof<H>>), anyhow::Error> {
881+
) -> Result<(PutResult<(NodeKey, Node)>, Option<SparseMerkleProof<H>>), NodeAlreadyExists> {
857882
// We are inserting a new key that shares a common prefix with the existing leaf node.
858883
// This check is to make sure that the visited nibble path of the inserted key is a
859884
// subpath of the existing leaf node's nibble path.

0 commit comments

Comments
 (0)