From aa1d59f47aeea442cfae94ec640b3c3cffacc20d Mon Sep 17 00:00:00 2001 From: Kris Nuttycombe Date: Thu, 22 Dec 2022 11:44:30 -0700 Subject: [PATCH 1/6] `witness` now witnesses at a checkpoint, rather than a root. This change exposed an inconsistency in how `BridgeTree` and `CompleteTree` were tracking `Mark` operations, which has now also been fixed. The behavior of the "unmark" operation has been simplified, such that it now ensures that a marked node will be made eligible for garbage collection when the checkpoint for the tree state in which the mark was forgotten rolls off the back of the checkpoint queue. --- bridgetree/CHANGELOG.md | 5 + .../{bridgetree.txt => lib.txt} | 4 + bridgetree/src/lib.rs | 352 ++++++++---------- incrementalmerkletree/CHANGELOG.md | 4 +- incrementalmerkletree/src/lib.rs | 2 +- incrementalmerkletree/src/testing.rs | 208 ++++++----- .../src/testing/complete_tree.rs | 341 +++++++++-------- 7 files changed, 442 insertions(+), 474 deletions(-) rename bridgetree/proptest-regressions/{bridgetree.txt => lib.txt} (50%) diff --git a/bridgetree/CHANGELOG.md b/bridgetree/CHANGELOG.md index 8b12559..a6830ac 100644 --- a/bridgetree/CHANGELOG.md +++ b/bridgetree/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to Rust's notion of ## [Unreleased] +### Changed + +- `BridgeTree::witness` now takes a checkpoint depth rather than a root hash to + identify the tree state with respect to which the witness should be constructed. + ### Removed - The `testing` module has been removed in favor of depending on diff --git a/bridgetree/proptest-regressions/bridgetree.txt b/bridgetree/proptest-regressions/lib.txt similarity index 50% rename from bridgetree/proptest-regressions/bridgetree.txt rename to bridgetree/proptest-regressions/lib.txt index afa608a..72c9155 100644 --- a/bridgetree/proptest-regressions/bridgetree.txt +++ b/bridgetree/proptest-regressions/lib.txt @@ -6,3 +6,7 @@ # everyone who runs the test benefits from these saved cases. cc a64e0d40c9287f2975a27447d266e4dd7228e44de99d000df39c4610598f486a # shrinks to tree = BridgeTree { depth: 8, prior_bridges: [MerkleBridge { prior_position: None, auth_fragments: {}, frontier: NonEmptyFrontier { position: Position(0), leaf: Left("a"), ommers: [] } }, MerkleBridge { prior_position: Some(Position(0)), auth_fragments: {}, frontier: NonEmptyFrontier { position: Position(1), leaf: Right("a", "a"), ommers: [] } }], current_bridge: Some(MerkleBridge { prior_position: Some(Position(1)), auth_fragments: {Position(1): AuthFragment { position: Position(1), altitudes_observed: 0, values: [] }}, frontier: NonEmptyFrontier { position: Position(1), leaf: Right("a", "a"), ommers: [] } }), saved: {Position(1): 1}, checkpoints: [Checkpoint { bridges_len: 1, is_witnessed: false, forgotten: {} }], max_checkpoints: 10 } cc 736aee7c92f418b3b7391b0ae253ca4dc18f9b6cc625c0c34f0e568d26421e92 # shrinks to tree = BridgeTree { depth: 8, prior_bridges: [], current_bridge: None, saved: {}, checkpoints: [], max_checkpoints: 10 } +cc f1a0c0e8114919f9f675e0b31cdecf37af31579e365119312a7ebefcabb639f1 # shrinks to ops = [Append(SipHashable(0)), Checkpoint, Checkpoint, Mark, Checkpoint, Checkpoint, Checkpoint, Checkpoint, Checkpoint, Authpath(Position(0), 6)] +cc ac8dad1a92cb9563a291802cbd3dfca2a89da35fe4ed377701f5ad85b43700f4 # shrinks to ops = [Append("a"), Append("a"), Checkpoint, Checkpoint, Checkpoint, Mark, Checkpoint, Authpath(Position(1), 2)] +cc cebb95fe896dc1d1e3c9a65efd50e786e6a4a3503c86fb2a6817bb05d25e751e # shrinks to tree = BridgeTree { depth: 8, prior_bridges: [MerkleBridge { prior_position: None, tracking: {Address { level: Level(2), index: 0 }}, ommers: {}, frontier: NonEmptyFrontier { position: Position(3), leaf: "a", ommers: ["a", "aa"] } }], current_bridge: Some(MerkleBridge { prior_position: Some(Position(3)), tracking: {Address { level: Level(2), index: 0 }}, ommers: {}, frontier: NonEmptyFrontier { position: Position(3), leaf: "a", ommers: ["a", "aa"] } }), saved: {Position(3): 0}, checkpoints: [Checkpoint { bridges_len: 0, marked: {Position(3)}, forgotten: {Position(3): 0} }], max_checkpoints: 10 } +cc cdab33688ef9270768481b72d1615ff1d209fdb3d8d45be2ad564a8d5e0addc1 # shrinks to ops = [Append(SipHashable(0)), Mark, Mark, Checkpoint, Mark, Rewind, Append(SipHashable(0)), Mark] diff --git a/bridgetree/src/lib.rs b/bridgetree/src/lib.rs index daf39f4..75961d0 100644 --- a/bridgetree/src/lib.rs +++ b/bridgetree/src/lib.rs @@ -31,7 +31,7 @@ //! //! In this module, the term "ommer" is used as for the sibling of a parent node in a binary tree. use serde::{Deserialize, Serialize}; -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::convert::TryFrom; use std::fmt::Debug; use std::mem::size_of; @@ -482,7 +482,7 @@ impl<'a, H: Hashable + Ord + Clone + 'a> MerkleBridge { /// will track the information necessary to create a witness for the leaf most /// recently appended to this bridge's frontier. #[must_use] - pub fn successor(&self, mark_current_leaf: bool) -> Self { + pub fn successor(&self, track_current_leaf: bool) -> Self { let mut result = Self { prior_position: Some(self.frontier.position()), tracking: self.tracking.clone(), @@ -490,7 +490,7 @@ impl<'a, H: Hashable + Ord + Clone + 'a> MerkleBridge { frontier: self.frontier.clone(), }; - if mark_current_leaf { + if track_current_leaf { result.track_current_leaf(); } @@ -616,43 +616,35 @@ impl<'a, H: Hashable + Ord + Clone + 'a> MerkleBridge { pub struct Checkpoint { /// The number of bridges that will be retained in a rewind. bridges_len: usize, - /// A flag indicating whether or not the current state of the tree - /// had been marked at the time the checkpoint was created. - is_marked: bool, /// A set of the positions that have been marked during the period that this /// checkpoint is the current checkpoint. marked: BTreeSet, - /// When a mark is forgotten, if the index of the forgotten mark is <= bridge_idx we - /// record it in the current checkpoint so that on rollback, we restore the forgotten - /// marks to the BridgeTree's "saved" list. If the mark was newly created since the - /// checkpoint, we don't need to remember when we forget it because both the mark - /// creation and removal will be reverted in the rollback. - forgotten: BTreeMap, + /// When a mark is forgotten, we add it to the checkpoint's forgotten set but + /// don't immediately remove it from the `saved` map; that removal occurs when + /// the checkpoint is eventually dropped. + forgotten: BTreeSet, } impl Checkpoint { /// Creates a new checkpoint from its constituent parts. pub fn from_parts( bridges_len: usize, - is_marked: bool, marked: BTreeSet, - forgotten: BTreeMap, + forgotten: BTreeSet, ) -> Self { Self { bridges_len, - is_marked, marked, forgotten, } } /// Creates a new empty checkpoint for the specified [`BridgeTree`] state. - pub fn at_length(bridges_len: usize, is_marked: bool) -> Self { + pub fn at_length(bridges_len: usize) -> Self { Checkpoint { bridges_len, - is_marked, marked: BTreeSet::new(), - forgotten: BTreeMap::new(), + forgotten: BTreeSet::new(), } } @@ -665,15 +657,6 @@ impl Checkpoint { self.bridges_len } - /// Returns whether the current state of the tree had been marked at the point that - /// this checkpoint was made. - /// - /// In the event of a rewind, the rewind logic will ensure that mark information is - /// properly reconstituted for the checkpointed tree state. - pub fn is_marked(&self) -> bool { - self.is_marked - } - /// Returns a set of the positions that have been marked during the period that this /// checkpoint is the current checkpoint. pub fn marked(&self) -> &BTreeSet { @@ -682,7 +665,7 @@ impl Checkpoint { /// Returns the set of previously-marked positions that have had their marks removed /// during the period that this checkpoint is the current checkpoint. - pub fn forgotten(&self) -> &BTreeMap { + pub fn forgotten(&self) -> &BTreeSet { &self.forgotten } @@ -714,9 +697,6 @@ impl Checkpoint { // using the specified rewrite function. Used during garbage collection. fn rewrite_indices usize>(&mut self, f: F) { self.bridges_len = f(self.bridges_len); - for v in self.forgotten.values_mut() { - *v = f(*v) - } } } @@ -732,8 +712,11 @@ pub struct BridgeTree { /// A map from positions for which we wish to be able to compute a /// witness to index in the bridges vector. saved: BTreeMap, - /// A stack of bridge indices to which it's possible to rewind directly. - checkpoints: Vec, + /// A deque of bridge indices to which it's possible to rewind directly. + /// This deque must be maintained to have a minimum size of 1 and a maximum + /// size of `max_checkpoints` in order to correctly maintain mark & rewind + /// semantics. + checkpoints: VecDeque, /// The maximum number of checkpoints to retain. If this number is /// exceeded, the oldest checkpoint will be dropped when creating /// a new checkpoint. @@ -764,24 +747,34 @@ pub enum BridgeTreeError { impl BridgeTree { /// Construct an empty BridgeTree value with the specified maximum number of checkpoints. + /// + /// Panics if `max_checkpoints < 1` because mark/rewind logic depends upon the presence + /// of checkpoints to function. pub fn new(max_checkpoints: usize) -> Self { + assert!(max_checkpoints >= 1); Self { prior_bridges: vec![], current_bridge: None, saved: BTreeMap::new(), - checkpoints: vec![], + checkpoints: VecDeque::from(vec![Checkpoint::at_length(0)]), max_checkpoints, } } - /// Removes the oldest checkpoint. Returns true if successful and false if - /// there are no checkpoints. + /// Removes the oldest checkpoin if there are more than `max_checkpoints`. Returns true if + /// successful and false if there are not enough checkpoints. fn drop_oldest_checkpoint(&mut self) -> bool { - if self.checkpoints.is_empty() { - false - } else { - self.checkpoints.remove(0); + if self.checkpoints.len() > self.max_checkpoints { + let c = self + .checkpoints + .pop_front() + .expect("Checkpoints deque is known to be non-empty."); + for pos in c.forgotten.iter() { + self.saved.remove(pos); + } true + } else { + false } } @@ -802,7 +795,7 @@ impl BridgeTree { } /// Returns the checkpoints to which this tree may be rewound. - pub fn checkpoints(&self) -> &[Checkpoint] { + pub fn checkpoints(&self) -> &VecDeque { &self.checkpoints } @@ -832,7 +825,7 @@ impl BridgeTree { frontier, )), saved: BTreeMap::new(), - checkpoints: vec![], + checkpoints: VecDeque::new(), max_checkpoints, } } @@ -843,7 +836,7 @@ impl BridgeTree { prior_bridges: Vec>, current_bridge: Option>, saved: BTreeMap, - checkpoints: Vec, + checkpoints: VecDeque, max_checkpoints: usize, ) -> Result { Self::check_consistency_internal( @@ -876,7 +869,7 @@ impl BridgeTree { prior_bridges: &[MerkleBridge], current_bridge: &Option>, saved: &BTreeMap, - checkpoints: &[Checkpoint], + checkpoints: &VecDeque, max_checkpoints: usize, ) -> Result<(), BridgeTreeError> { // check that saved values correspond to bridges @@ -969,10 +962,10 @@ impl BridgeTree { self.current_bridge.as_ref().map(|b| b.current_leaf()) } - /// Marks the current leaf as one for which we're interested in producing - /// a witness. Returns an optional value containing the - /// current position if successful or if the current value was already - /// marked, or None if the tree is empty. + /// Marks the current leaf as one for which we're interested in producing a witness. + /// + /// Returns an optional value containing the current position if successful or if the current + /// value was already marked, or None if the tree is empty. pub fn mark(&mut self) -> Option { match self.current_bridge.take() { Some(mut cur_b) => { @@ -995,15 +988,14 @@ impl BridgeTree { self.current_bridge = Some(successor); } - self.saved - .entry(pos) - .or_insert(self.prior_bridges.len() - 1); - // mark the position as having been marked in the current checkpoint - if let Some(c) = self.checkpoints.last_mut() { - if !c.is_marked { - c.marked.insert(pos); - } + if let std::collections::btree_map::Entry::Vacant(e) = self.saved.entry(pos) { + let c = self + .checkpoints + .back_mut() + .expect("Checkpoints deque must never be empty"); + c.marked.insert(pos); + e.insert(self.prior_bridges.len() - 1); } Some(pos) @@ -1029,15 +1021,12 @@ impl BridgeTree { /// interested in maintaining a mark for. Returns true if successful and /// false if we were already not maintaining a mark at this position. pub fn remove_mark(&mut self, position: Position) -> bool { - if let Some(idx) = self.saved.remove(&position) { - // If the position is one that has *not* just been marked since the last checkpoint, - // then add it to the set of those forgotten during the current checkpoint span so that - // it can be restored on rollback. - if let Some(c) = self.checkpoints.last_mut() { - if !c.marked.contains(&position) { - c.forgotten.insert(position, idx); - } - } + if self.saved.contains_key(&position) { + let c = self + .checkpoints + .back_mut() + .expect("Checkpoints deque must never be empty."); + c.forgotten.insert(position); true } else { false @@ -1050,8 +1039,6 @@ impl BridgeTree { pub fn checkpoint(&mut self) { match self.current_bridge.take() { Some(cur_b) => { - let is_marked = self.get_marked_leaf(cur_b.position()).is_some(); - // Do not create a duplicate bridge if self .prior_bridges @@ -1065,10 +1052,10 @@ impl BridgeTree { } self.checkpoints - .push(Checkpoint::at_length(self.prior_bridges.len(), is_marked)); + .push_back(Checkpoint::at_length(self.prior_bridges.len())); } None => { - self.checkpoints.push(Checkpoint::at_length(0, false)); + self.checkpoints.push_back(Checkpoint::at_length(0)); } } @@ -1083,20 +1070,25 @@ impl BridgeTree { /// at that tree state have been removed using `rewind`. This function /// return false and leave the tree unmodified if no checkpoints exist. pub fn rewind(&mut self) -> bool { - match self.checkpoints.pop() { - Some(mut c) => { - // drop marked values at and above the checkpoint height; - // we will re-mark if necessary. - self.saved.append(&mut c.forgotten); - self.saved.retain(|_, i| *i + 1 < c.bridges_len); - self.prior_bridges.truncate(c.bridges_len); - self.current_bridge = self.prior_bridges.last().map(|b| b.successor(c.is_marked)); - if c.is_marked { - self.mark(); - } - true + if self.checkpoints.len() > 1 { + let c = self + .checkpoints + .pop_back() + .expect("Checkpoints deque is known to be non-empty."); + + // Remove marks for positions that were marked during the lifetime of this checkpoint. + for pos in c.marked { + self.saved.remove(&pos); } - None => false, + + self.prior_bridges.truncate(c.bridges_len); + self.current_bridge = self + .prior_bridges + .last() + .map(|b| b.successor(self.saved.contains_key(&b.position()))); + true + } else { + false } } @@ -1105,58 +1097,47 @@ impl BridgeTree { /// Returns `None` if there is no available witness to that /// position or if the root does not correspond to a checkpointed /// root of the tree. - pub fn witness(&self, position: Position, as_of_root: &H) -> Option> { - self.witness_inner(position, as_of_root).ok() - } - - fn witness_inner(&self, position: Position, as_of_root: &H) -> Result, WitnessingError> { + pub fn witness( + &self, + position: Position, + checkpoint_depth: usize, + ) -> Result, WitnessingError> { #[derive(Debug)] enum AuthBase<'a> { Current, Checkpoint(usize, &'a Checkpoint), - NotFound, } - let max_alt = Level::from(DEPTH); - // Find the earliest checkpoint having a matching root, or the current // root if it matches and there is no earlier matching checkpoint. - let auth_base = self - .checkpoints - .iter() - .enumerate() - .rev() - .take_while(|(_, c)| c.position(&self.prior_bridges) >= Some(position)) - .filter(|(_, c)| &c.root(&self.prior_bridges, max_alt) == as_of_root) - .last() - .map(|(i, c)| AuthBase::Checkpoint(i, c)) - .unwrap_or_else(|| { - if self.root(0).as_ref() == Some(as_of_root) { - AuthBase::Current - } else { - AuthBase::NotFound - } - }); + let auth_base = if checkpoint_depth == 0 { + Ok(AuthBase::Current) + } else if self.checkpoints.len() >= checkpoint_depth { + let c_idx = self.checkpoints.len() - checkpoint_depth; + if self + .checkpoints + .iter() + .skip(c_idx) + .take_while(|c| { + c.position(&self.prior_bridges) + .iter() + .any(|p| p <= &position) + }) + .any(|c| c.marked.contains(&position)) + { + // The mark had not yet been established at the point the checkpoint was + // created, so we can't treat it as marked. + Err(WitnessingError::PositionNotMarked(position)) + } else { + Ok(AuthBase::Checkpoint(c_idx, &self.checkpoints[c_idx])) + } + } else { + Err(WitnessingError::CheckpointInvalid) + }?; let saved_idx = self .saved .get(&position) - .or_else(|| { - if let AuthBase::Checkpoint(i, _) = auth_base { - // The saved position might have been forgotten since the checkpoint, - // so look for it in each of the subsequent checkpoints' forgotten - // items. - self.checkpoints[i..].iter().find_map(|c| { - // restore the forgotten position, if that position was not also marked - // in the same checkpoint - c.forgotten - .get(&position) - .filter(|_| !c.marked.contains(&position)) - }) - } else { - None - } - }) .ok_or(WitnessingError::PositionNotMarked(position))?; let prior_frontier = &self.prior_bridges[*saved_idx].frontier; @@ -1198,10 +1179,6 @@ impl BridgeTree { fuse_from - checkpoint.bridges_len, )) } - AuthBase::NotFound => { - // we didn't find any suitable auth base - Err(WitnessingError::AuthBaseNotFound) - } }?; successor.witness(DEPTH, prior_frontier) @@ -1216,17 +1193,7 @@ impl BridgeTree { // checkpoints; we cannot remove information that we might need to restore in // a rewind. if self.checkpoints.len() == self.max_checkpoints { - let gc_len = self.checkpoints.first().unwrap().bridges_len; - // Get a list of the leaf positions that we need to retain. This consists of - // all the saved leaves, plus all the leaves that have been forgotten since - // the most distant checkpoint to which we could rewind. - let remember: BTreeSet = self - .saved - .keys() - .chain(self.checkpoints.iter().flat_map(|c| c.forgotten.keys())) - .cloned() - .collect(); - + let gc_len = self.checkpoints.front().unwrap().bridges_len; let mut cur: Option> = None; let mut merged = 0; let mut ommer_addrs: BTreeSet
= BTreeSet::new(); @@ -1236,7 +1203,7 @@ impl BridgeTree { { if let Some(cur_bridge) = cur { let pos = cur_bridge.position(); - let mut new_cur = if remember.contains(&pos) || i > gc_len { + let mut new_cur = if self.saved.contains_key(&pos) || i > gc_len { // We need to remember cur_bridge; update its save index & put next_bridge // on the chopping block if let Some(idx) = self.saved.get_mut(&pos) { @@ -1300,9 +1267,8 @@ mod tests { use incrementalmerkletree::{ testing::{ apply_operation, arb_operation, check_checkpoint_rewind, check_operations, - check_rewind_remove_mark, check_rewind_remove_mark_consistency, check_root_hashes, - check_witnesses, complete_tree::CompleteTree, CombinedTree, Frontier, SipHashable, - Tree, + check_remove_mark, check_rewind_remove_mark, check_root_hashes, check_witnesses, + complete_tree::CompleteTree, CombinedTree, Frontier, SipHashable, Tree, }, Hashable, }; @@ -1322,6 +1288,10 @@ mod tests { BridgeTree::append(self, value) } + fn depth(&self) -> u8 { + DEPTH + } + fn current_position(&self) -> Option { BridgeTree::current_position(self) } @@ -1346,8 +1316,8 @@ mod tests { BridgeTree::root(self, checkpoint_depth) } - fn witness(&self, position: Position, as_of_root: &H) -> Option> { - BridgeTree::witness(self, position, as_of_root) + fn witness(&self, position: Position, checkpoint_depth: usize) -> Option> { + BridgeTree::witness(self, position, checkpoint_depth).ok() } fn remove_mark(&mut self, position: Position) -> bool { @@ -1467,6 +1437,22 @@ mod tests { assert!(!tree.append('i'.to_string())); } + fn check_garbage_collect( + mut tree: BridgeTree, + ) { + // Add checkpoints until we're sure everything that can be gc'ed will be gc'ed + for _ in 0..tree.max_checkpoints { + tree.checkpoint(); + } + + let mut tree_mut = tree.clone(); + tree_mut.garbage_collect(); + + for pos in tree.saved.keys() { + assert_eq!(tree.witness(*pos, 0), tree_mut.witness(*pos, 0)); + } + } + fn arb_bridgetree( item_gen: G, max_count: usize, @@ -1506,40 +1492,10 @@ mod tests { fn prop_garbage_collect( tree in arb_bridgetree((97u8..123).prop_map(|c| char::from(c).to_string()), 100) ) { - let mut tree_mut = tree.clone(); - // ensure we have enough checkpoints to not rewind past the state `tree` is in - for _ in 0..10 { - tree_mut.checkpoint(); - } - - tree_mut.garbage_collect(); - - tree_mut.rewind(); - - for pos in tree.saved.keys() { - assert_eq!( - tree.witness(*pos, &tree.root(0).unwrap()), - tree_mut.witness(*pos, &tree.root(0).unwrap()) - ); - } + check_garbage_collect(tree); } } - #[test] - fn drop_oldest_checkpoint() { - let mut t = BridgeTree::::new(100); - t.checkpoint(); - t.append("a".to_string()); - t.mark(); - t.append("b".to_string()); - t.append("c".to_string()); - assert!( - t.drop_oldest_checkpoint(), - "Checkpoint drop is expected to succeed" - ); - assert!(!t.rewind(), "Rewind is expected to fail."); - } - #[test] fn root_hashes() { check_root_hashes(BridgeTree::::new); @@ -1562,6 +1518,17 @@ mod tests { #[test] fn garbage_collect() { + let mut tree: BridgeTree = BridgeTree::new(100); + let empty_root = tree.root(0); + tree.append("a".to_string()); + for _ in 0..100 { + tree.checkpoint(); + } + tree.garbage_collect(); + assert!(tree.root(0) != empty_root); + tree.rewind(); + assert!(tree.root(0) != empty_root); + let mut t = BridgeTree::::new(10); let mut to_unmark = vec![]; let mut has_witness = vec![]; @@ -1588,7 +1555,7 @@ mod tests { assert_eq!(t.prior_bridges().len(), 20 + 14 - 2); let witnesss = has_witness .iter() - .map(|pos| match t.witness_inner(*pos, &t.root(0).unwrap()) { + .map(|pos| match t.witness(*pos, 0) { Ok(path) => path, Err(e) => panic!("Failed to get auth path: {:?}", e), }) @@ -1598,46 +1565,29 @@ mod tests { assert_eq!(t.prior_bridges().len(), 32 - 10 + 1 - 3); let retained_witnesss = has_witness .iter() - .map(|pos| { - t.witness(*pos, &t.root(0).unwrap()) - .expect("Must be able to get auth path") - }) + .map(|pos| t.witness(*pos, 0).expect("Must be able to get auth path")) .collect::>(); assert_eq!(witnesss, retained_witnesss); } - #[test] - fn garbage_collect_idx() { - let mut tree: BridgeTree = BridgeTree::new(100); - let empty_root = tree.root(0); - tree.append("a".to_string()); - for _ in 0..100 { - tree.checkpoint(); - } - tree.garbage_collect(); - assert!(tree.root(0) != empty_root); - tree.rewind(); - assert!(tree.root(0) != empty_root); - } - // Combined tree tests fn new_combined_tree( max_checkpoints: usize, - ) -> CombinedTree, BridgeTree> { + ) -> CombinedTree, BridgeTree> { CombinedTree::new( - CompleteTree::new(4, max_checkpoints), + CompleteTree::::new(max_checkpoints), BridgeTree::::new(max_checkpoints), ) } #[test] - fn test_rewind_remove_mark() { - check_rewind_remove_mark(new_combined_tree); + fn combined_remove_mark() { + check_remove_mark(new_combined_tree); } #[test] - fn test_rewind_remove_mark_consistency() { - check_rewind_remove_mark_consistency(new_combined_tree); + fn combined_rewind_remove_mark() { + check_rewind_remove_mark(new_combined_tree); } proptest! { @@ -1651,7 +1601,7 @@ mod tests { ) ) { let tree = new_combined_tree(100); - check_operations(tree, 4, &ops)?; + check_operations(tree, &ops)?; } #[test] @@ -1662,7 +1612,7 @@ mod tests { ) ) { let tree = new_combined_tree(100); - check_operations(tree, 4, &ops)?; + check_operations(tree, &ops)?; } } } diff --git a/incrementalmerkletree/CHANGELOG.md b/incrementalmerkletree/CHANGELOG.md index abfc85f..959e55f 100644 --- a/incrementalmerkletree/CHANGELOG.md +++ b/incrementalmerkletree/CHANGELOG.md @@ -33,7 +33,9 @@ is not another good use case for polymorphism over tree implementations. - `Tree::witnessed_positions` has been renamed to `Tree::marked_positions` - `Tree::get_witnessed_leaf` has been renamed to `Tree::get_marked_leaf` - `Tree::remove_witness` has been renamed to `Tree::remove_mark` - - `Tree::authentication_path` has been renamed to `Tree::witness` + - `Tree::authentication_path` has been renamed to `Tree::witness`. Also, this method + now takes a checkpoint depth as its second argument rather than a Merkle root, + to better support future changes. - `Tree::append` now takes ownership of the value being appended instead of a value passed by reference. diff --git a/incrementalmerkletree/src/lib.rs b/incrementalmerkletree/src/lib.rs index e2ca138..9f667ce 100644 --- a/incrementalmerkletree/src/lib.rs +++ b/incrementalmerkletree/src/lib.rs @@ -373,7 +373,7 @@ impl<'a> From<&'a Address> for Option { /// A trait describing the operations that make a type suitable for use as /// a leaf or node value in a merkle tree. -pub trait Hashable: Sized { +pub trait Hashable: Sized + core::fmt::Debug { fn empty_leaf() -> Self; fn combine(level: Level, a: &Self, b: &Self) -> Self; diff --git a/incrementalmerkletree/src/testing.rs b/incrementalmerkletree/src/testing.rs index 81341f0..94a243c 100644 --- a/incrementalmerkletree/src/testing.rs +++ b/incrementalmerkletree/src/testing.rs @@ -29,6 +29,9 @@ pub trait Frontier { /// A Merkle tree that supports incremental appends, marking of /// leaf nodes for construction of witnesses, checkpoints and rollbacks. pub trait Tree { + /// Returns the depth of the tree. + fn depth(&self) -> u8; + /// Appends a new value to the tree at the next available slot. /// Returns true if successful and false if the tree would exceed /// the maximum allowed depth. @@ -64,7 +67,7 @@ pub trait Tree { /// Returns `None` if there is no available witness to that /// position or if the root does not correspond to a checkpointed /// root of the tree. - fn witness(&self, position: Position, as_of_root: &H) -> Option>; + fn witness(&self, position: Position, checkpoint_depth: usize) -> Option>; /// Marks the value at the specified position as a value we're no longer /// interested in maintaining a mark for. Returns true if successful and @@ -117,6 +120,19 @@ impl Hashable for String { } } +impl Hashable for Option { + fn empty_leaf() -> Self { + Some(H::empty_leaf()) + } + + fn combine(l: Level, a: &Self, b: &Self) -> Self { + match (a, b) { + (Some(a), Some(b)) => Some(H::combine(l, a, b)), + _ => None, + } + } +} + // // Operations // @@ -177,10 +193,7 @@ impl Operation { assert!(tree.rewind(), "rewind failed"); None } - Authpath(p, d) => tree - .root(*d) - .and_then(|root| tree.witness(*p, &root)) - .map(|xs| (*p, xs)), + Authpath(p, d) => tree.witness(*p, *d).map(|xs| (*p, xs)), GarbageCollect => None, } } @@ -252,7 +265,6 @@ pub fn apply_operation>(tree: &mut T, op: Operation) { pub fn check_operations>( mut tree: T, - tree_depth: u8, ops: &[Operation], ) -> Result<(), TestCaseError> { let mut tree_size = 0; @@ -265,11 +277,11 @@ pub fn check_operations>( match op { Append(value) => { if tree.append(value.clone()) { - prop_assert!(tree_size < (1 << tree_depth)); + prop_assert!(tree_size < (1 << tree.depth())); tree_size += 1; tree_values.push(value.clone()); } else { - prop_assert_eq!(tree_size, 1 << tree_depth); + prop_assert_eq!(tree_size, 1 << tree.depth()); } } CurrentPosition => { @@ -310,7 +322,7 @@ pub fn check_operations>( } } Authpath(position, depth) => { - if let Some(path) = tree.root(*depth).and_then(|r| tree.witness(*position, &r)) { + if let Some(path) = tree.witness(*position, *depth) { let value: H = tree_values[::from(*position)].clone(); let tree_root = tree.root(*depth); @@ -324,11 +336,10 @@ pub fn check_operations>( extended_tree_values.truncate(*checkpointed_tree_size); } } - // extend the tree with empty leaves until it is full - extended_tree_values.resize(1 << tree_depth, H::empty_leaf()); // compute the root - let expected_root = complete_tree::root::(extended_tree_values); + let expected_root = + complete_tree::root::(&extended_tree_values, tree.depth()); prop_assert_eq!(&tree_root.unwrap(), &expected_root); prop_assert_eq!( @@ -376,6 +387,7 @@ pub struct CombinedTree, E: Tree> { impl, E: Tree> CombinedTree { pub fn new(inefficient: I, efficient: E) -> Self { + assert_eq!(inefficient.depth(), efficient.depth()); CombinedTree { inefficient, efficient, @@ -385,6 +397,10 @@ impl, E: Tree> CombinedTree, E: Tree> Tree for CombinedTree { + fn depth(&self) -> u8 { + self.inefficient.depth() + } + fn append(&mut self, value: H) -> bool { let a = self.inefficient.append(value.clone()); let b = self.efficient.append(value); @@ -437,9 +453,9 @@ impl, E: Tree> Tree for Comb a } - fn witness(&self, position: Position, as_of_root: &H) -> Option> { - let a = self.inefficient.witness(position, as_of_root); - let b = self.efficient.witness(position, as_of_root); + fn witness(&self, position: Position, checkpoint_depth: usize) -> Option> { + let a = self.inefficient.witness(position, checkpoint_depth); + let b = self.efficient.witness(position, checkpoint_depth); assert_eq!(a, b); a } @@ -496,7 +512,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.append("a".to_string()); tree.mark(); assert_eq!( - tree.witness(Position::from(0), &tree.root(0).unwrap()), + tree.witness(Position::from(0), 0), Some(vec![ "_".to_string(), "__".to_string(), @@ -507,7 +523,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.append("b".to_string()); assert_eq!( - tree.witness(0.into(), &tree.root(0).unwrap()), + tree.witness(0.into(), 0), Some(vec![ "b".to_string(), "__".to_string(), @@ -519,7 +535,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.append("c".to_string()); tree.mark(); assert_eq!( - tree.witness(Position::from(2), &tree.root(0).unwrap()), + tree.witness(Position::from(2), 0), Some(vec![ "_".to_string(), "ab".to_string(), @@ -530,7 +546,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.append("d".to_string()); assert_eq!( - tree.witness(Position::from(2), &tree.root(0).unwrap()), + tree.witness(Position::from(2), 0), Some(vec![ "d".to_string(), "ab".to_string(), @@ -541,7 +557,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.append("e".to_string()); assert_eq!( - tree.witness(Position::from(2), &tree.root(0).unwrap()), + tree.witness(Position::from(2), 0), Some(vec![ "d".to_string(), "ab".to_string(), @@ -560,7 +576,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.append("h".to_string()); assert_eq!( - tree.witness(0.into(), &tree.root(0).unwrap()), + tree.witness(0.into(), 0), Some(vec![ "b".to_string(), "cd".to_string(), @@ -583,7 +599,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.append("g".to_string()); assert_eq!( - tree.witness(Position::from(5), &tree.root(0).unwrap()), + tree.witness(Position::from(5), 0), Some(vec![ "e".to_string(), "g_".to_string(), @@ -600,7 +616,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.append('l'.to_string()); assert_eq!( - tree.witness(Position::from(10), &tree.root(0).unwrap()), + tree.witness(Position::from(10), 0), Some(vec![ "l".to_string(), "ij".to_string(), @@ -621,9 +637,8 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new for c in 'f'..'i' { tree.append(c.to_string()); } - assert_eq!( - tree.witness(0.into(), &tree.root(0).unwrap()), + tree.witness(0.into(), 0), Some(vec![ "b".to_string(), "cd".to_string(), @@ -645,9 +660,8 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.checkpoint(); tree.append('h'.to_string()); assert!(tree.rewind()); - assert_eq!( - tree.witness(Position::from(2), &tree.root(0).unwrap()), + tree.witness(Position::from(2), 0), Some(vec![ "d".to_string(), "ab".to_string(), @@ -660,10 +674,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.append('a'.to_string()); tree.append('b'.to_string()); tree.mark(); - assert_eq!( - tree.witness(Position::from(0), &tree.root(0).unwrap()), - None - ); + assert_eq!(tree.witness(Position::from(0), 0), None); let mut tree = new_tree(100); for c in 'a'..'n' { @@ -676,7 +687,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new tree.append('p'.to_string()); assert_eq!( - tree.witness(Position::from(12), &tree.root(0).unwrap()), + tree.witness(Position::from(12), 0), Some(vec![ "n".to_string(), "op".to_string(), @@ -750,7 +761,50 @@ pub fn check_checkpoint_rewind, F: Fn(usize) -> T>(new_tree: F) assert_eq!(t.root(0).unwrap(), "ab______________"); } +pub fn check_remove_mark, F: Fn(usize) -> T>(new_tree: F) { + let samples = vec![ + vec![ + append_str("a"), + append_str("a"), + Checkpoint, + Mark, + witness(1, 1), + ], + vec![ + append_str("a"), + append_str("a"), + append_str("a"), + append_str("a"), + Mark, + Checkpoint, + unmark(3), + witness(3, 0), + ], + vec![ + append_str("a"), + append_str("a"), + Checkpoint, + Checkpoint, + Checkpoint, + Mark, + Checkpoint, + witness(1, 3), + ], + ]; + + for (i, sample) in samples.iter().enumerate() { + let result = check_operations(new_tree(100), sample); + assert!( + matches!(result, Ok(())), + "Reference/Test mismatch at index {}: {:?}", + i, + result + ); + } +} + pub fn check_rewind_remove_mark, F: Fn(usize) -> T>(new_tree: F) { + // rewinding doesn't remove a mark let mut tree = new_tree(100); tree.append("e".to_string()); tree.mark(); @@ -758,6 +812,7 @@ pub fn check_rewind_remove_mark, F: Fn(usize) -> T>(new_tree: F) assert!(tree.rewind()); assert!(tree.remove_mark(0usize.into())); + // the order of checkpoint & mark does not matter let mut tree = new_tree(100); tree.append("e".to_string()); tree.checkpoint(); @@ -765,35 +820,21 @@ pub fn check_rewind_remove_mark, F: Fn(usize) -> T>(new_tree: F) assert!(tree.rewind()); assert!(!tree.remove_mark(0usize.into())); - let mut tree = new_tree(100); + // use a maximum number of checkpoints of 1 + let mut tree = new_tree(1); tree.append("e".to_string()); tree.mark(); tree.checkpoint(); + assert!(tree.marked_positions().contains(&0usize.into())); + tree.append("f".to_string()); + // simulate a spend of `e` at `f` assert!(tree.remove_mark(0usize.into())); - assert!(tree.rewind()); - assert!(tree.remove_mark(0usize.into())); - - let mut tree = new_tree(100); - tree.append("e".to_string()); - tree.mark(); - assert!(tree.remove_mark(0usize.into())); + // even though the mark has been staged for removal, it's not gone yet + assert!(tree.marked_positions().contains(&0usize.into())); tree.checkpoint(); - assert!(tree.rewind()); - assert!(!tree.remove_mark(0usize.into())); - - let mut tree = new_tree(100); - tree.append("a".to_string()); - assert!(!tree.remove_mark(0usize.into())); - tree.checkpoint(); - assert!(tree.mark().is_some()); - assert!(tree.rewind()); - - let mut tree = new_tree(100); - tree.append("a".to_string()); - tree.checkpoint(); - assert!(tree.mark().is_some()); - assert!(tree.remove_mark(0usize.into())); - assert!(tree.rewind()); + // the newest checkpoint will have caused the oldest to roll off, and + // so the forgotten node will be unmarked + assert!(!tree.marked_positions().contains(&0usize.into())); assert!(!tree.remove_mark(0usize.into())); // The following check_operations tests cover errors where the @@ -829,10 +870,19 @@ pub fn check_rewind_remove_mark, F: Fn(usize) -> T>(new_tree: F) unmark(0), unmark(0), ], + vec![ + append_str("a"), + Mark, + Checkpoint, + Mark, + Rewind, + append_str("a"), + Mark, + ], ]; for (i, sample) in samples.iter().enumerate() { - let result = check_operations(new_tree(100), 4, sample); + let result = check_operations(new_tree(100), sample); assert!( matches!(result, Ok(())), "Reference/Test mismatch at index {}: {:?}", @@ -972,49 +1022,7 @@ pub fn check_witness_consistency, F: Fn(usize) -> T>(new_tree: F ]; for (i, sample) in samples.iter().enumerate() { - let result = check_operations(new_tree(100), 4, sample); - assert!( - matches!(result, Ok(())), - "Reference/Test mismatch at index {}: {:?}", - i, - result - ); - } -} - -pub fn check_rewind_remove_mark_consistency, F: Fn(usize) -> T>(new_tree: F) { - let samples = vec![ - vec![append_str("x"), Checkpoint, Mark, Rewind, unmark(0)], - vec![ - append_str("d"), - Checkpoint, - Mark, - unmark(0), - Rewind, - unmark(0), - ], - vec![ - append_str("o"), - Checkpoint, - Mark, - Checkpoint, - unmark(0), - Rewind, - Rewind, - ], - vec![ - append_str("s"), - Mark, - append_str("m"), - Checkpoint, - unmark(0), - Rewind, - unmark(0), - unmark(0), - ], - ]; - for (i, sample) in samples.iter().enumerate() { - let result = check_operations(new_tree(100), 4, sample); + let result = check_operations(new_tree(100), sample); assert!( matches!(result, Ok(())), "Reference/Test mismatch at index {}: {:?}", diff --git a/incrementalmerkletree/src/testing/complete_tree.rs b/incrementalmerkletree/src/testing/complete_tree.rs index 080f645..6c41d05 100644 --- a/incrementalmerkletree/src/testing/complete_tree.rs +++ b/incrementalmerkletree/src/testing/complete_tree.rs @@ -1,13 +1,17 @@ //! Sample implementation of the Tree interface. -use std::collections::BTreeSet; +use std::cmp::min; +use std::collections::{BTreeSet, VecDeque}; -use crate::{ - testing::{Frontier, Tree}, - Hashable, Level, Position, -}; +use crate::{testing::Tree, Hashable, Level, Position}; -pub(crate) fn root(mut leaves: Vec) -> H { - leaves.resize(leaves.len().next_power_of_two(), H::empty_leaf()); +pub(crate) fn root(leaves: &[H], depth: u8) -> H { + let empty_leaf = H::empty_leaf(); + let mut leaves = leaves + .iter() + .chain(std::iter::repeat(&empty_leaf)) + .take(1 << depth) + .cloned() + .collect::>(); //leaves are always at level zero, so we start there. let mut level = Level::from(0); @@ -32,214 +36,209 @@ pub(crate) fn root(mut leaves: Vec) -> H { leaves[0].clone() } +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Checkpoint { + /// The number of leaves that will be retained + leaves_len: usize, + /// A set of the positions that have been marked during the period that this + /// checkpoint is the current checkpoint. + marked: BTreeSet, + /// When a mark is forgotten, we add it to the checkpoint's forgotten set but + /// don't immediately remove it from the `marked` set; that removal occurs when + /// the checkpoint is eventually dropped. + forgotten: BTreeSet, +} + +impl Checkpoint { + pub fn at_length(leaves_len: usize) -> Self { + Checkpoint { + leaves_len, + marked: BTreeSet::new(), + forgotten: BTreeSet::new(), + } + } +} + #[derive(Clone, Debug)] -pub struct TreeState { - leaves: Vec, - current_offset: usize, +pub struct CompleteTree { + leaves: Vec>, marks: BTreeSet, - depth: usize, -} - -impl TreeState { - /// Creates a new, empty binary tree of specified depth. - pub fn new(depth: usize) -> Self { - Self { - leaves: vec![H::empty_leaf(); 1 << depth], - current_offset: 0, - marks: BTreeSet::new(), - depth, - } - } -} - -impl Frontier for TreeState { - fn append(&mut self, value: H) -> bool { - if self.current_offset == (1 << self.depth) { - false - } else { - self.leaves[self.current_offset] = value; - self.current_offset += 1; - true - } - } - - /// Obtains the current root of this Merkle tree. - fn root(&self) -> H { - root(self.leaves.clone()) - } -} - -impl TreeState { - fn current_position(&self) -> Option { - if self.current_offset == 0 { - None - } else { - Some((self.current_offset - 1).into()) - } - } - - /// Returns the leaf most recently appended to the tree - fn current_leaf(&self) -> Option<&H> { - self.current_position() - .map(|p| &self.leaves[::from(p)]) - } - - /// Returns the leaf at the specified position if the tree can produce - /// a witness for it. - fn get_marked_leaf(&self, position: Position) -> Option<&H> { - if self.marks.contains(&position) { - self.leaves.get(::from(position)) - } else { - None - } - } - - /// Marks the current tree state leaf as a value that we're interested in - /// marking. Returns the current position if the tree is non-empty. - fn mark(&mut self) -> Option { - self.current_position().map(|pos| { - self.marks.insert(pos); - pos - }) - } - - /// Obtains a witness to the value at the specified position. - /// Returns `None` if there is no available witness to that - /// value. - fn witness(&self, position: Position) -> Option> { - if Some(position) <= self.current_position() { - let mut path = vec![]; - - let mut leaf_idx: usize = position.into(); - for bit in 0..self.depth { - leaf_idx ^= 1 << bit; - path.push(root::(self.leaves[leaf_idx..][0..(1 << bit)].to_vec())); - leaf_idx &= usize::MAX << (bit + 1); - } - - Some(path) - } else { - None - } - } - - /// Marks the value at the specified position as a value we're no longer - /// interested in maintaining a mark for. Returns true if successful and - /// false if we were already not maintaining a mark at this position. - fn remove_mark(&mut self, position: Position) -> bool { - self.marks.remove(&position) - } -} - -#[derive(Clone, Debug)] -pub struct CompleteTree { - tree_state: TreeState, - checkpoints: Vec>, + checkpoints: VecDeque, max_checkpoints: usize, } -impl CompleteTree { - /// Creates a new, empty binary tree of specified depth. - pub fn new(depth: usize, max_checkpoints: usize) -> Self { +impl CompleteTree { + /// Creates a new, empty binary tree + pub fn new(max_checkpoints: usize) -> Self { Self { - tree_state: TreeState::new(depth), - checkpoints: vec![], + leaves: vec![], + marks: BTreeSet::new(), + checkpoints: VecDeque::from(vec![Checkpoint::at_length(0)]), max_checkpoints, } } -} -impl CompleteTree { - /// Removes the oldest checkpoint. Returns true if successful and false if - /// there are no checkpoints. - fn drop_oldest_checkpoint(&mut self) -> bool { - if self.checkpoints.is_empty() { + fn append(&mut self, value: H) -> bool { + if self.leaves.len() == (1 << DEPTH) { false } else { - self.checkpoints.remove(0); + self.leaves.push(Some(value)); true } } - /// Retrieve the tree state at the specified checkpoint depth. This - /// is the current tree state if the depth is 0, and this will return - /// None if not enough checkpoints exist to obtain the requested depth. - fn tree_state_at_checkpoint_depth(&self, checkpoint_depth: usize) -> Option<&TreeState> { - if self.checkpoints.len() < checkpoint_depth { - None - } else if checkpoint_depth == 0 { - Some(&self.tree_state) - } else { + fn leaves_at_checkpoint_depth(&self, checkpoint_depth: usize) -> Option { + if checkpoint_depth == 0 { + Some(self.leaves.len()) + } else if checkpoint_depth <= self.checkpoints.len() { self.checkpoints .get(self.checkpoints.len() - checkpoint_depth) + .map(|c| c.leaves_len) + } else { + None } } } -impl Tree for CompleteTree { - /// Appends a new value to the tree at the next available slot. Returns true - /// if successful and false if the tree is full. - fn append(&mut self, value: H) -> bool { - self.tree_state.append(value) +impl CompleteTree { + /// Removes the oldest checkpoint. Returns true if successful and false if + /// there are fewer than `self.max_checkpoints` checkpoints. + fn drop_oldest_checkpoint(&mut self) -> bool { + if self.checkpoints.len() > self.max_checkpoints { + let c = self.checkpoints.pop_front().unwrap(); + for pos in c.forgotten.iter() { + self.marks.remove(pos); + } + true + } else { + false + } + } +} + +impl Tree + for CompleteTree +{ + fn depth(&self) -> u8 { + DEPTH + } + + fn append(&mut self, value: H) -> bool { + Self::append(self, value) } - /// Returns the most recently appended leaf value. fn current_position(&self) -> Option { - self.tree_state.current_position() + if self.leaves.is_empty() { + None + } else { + Some((self.leaves.len() - 1).into()) + } } fn current_leaf(&self) -> Option<&H> { - self.tree_state.current_leaf() + self.leaves.last().and_then(|opt: &Option| opt.as_ref()) } fn get_marked_leaf(&self, position: Position) -> Option<&H> { - self.tree_state.get_marked_leaf(position) + if self.marks.contains(&position) { + self.leaves + .get(usize::from(position)) + .and_then(|opt: &Option| opt.as_ref()) + } else { + None + } } fn mark(&mut self) -> Option { - self.tree_state.mark() + match self.current_position() { + Some(pos) => { + if !self.marks.contains(&pos) { + self.marks.insert(pos); + self.checkpoints.back_mut().unwrap().marked.insert(pos); + } + Some(pos) + } + None => None, + } } fn marked_positions(&self) -> BTreeSet { - self.tree_state.marks.clone() + self.marks.clone() } fn root(&self, checkpoint_depth: usize) -> Option { - self.tree_state_at_checkpoint_depth(checkpoint_depth) - .map(|s| s.root()) + self.leaves_at_checkpoint_depth(checkpoint_depth) + .and_then(|len| root(&self.leaves[0..len], DEPTH)) } - fn witness(&self, position: Position, root: &H) -> Option> { - // Search for the checkpointed state corresponding to the provided root, and if one is - // found, compute the witness as of that root. - self.checkpoints - .iter() - .chain(Some(&self.tree_state)) - .rev() - .skip_while(|c| !c.marks.contains(&position)) - .find_map(|c| { - if &c.root() == root { - c.witness(position) - } else { - None + fn witness(&self, position: Position, checkpoint_depth: usize) -> Option> { + if self.marks.contains(&position) && checkpoint_depth <= self.checkpoints.len() { + let checkpoint_idx = self.checkpoints.len() - checkpoint_depth; + let len = if checkpoint_depth == 0 { + self.leaves.len() + } else { + self.checkpoints[checkpoint_idx].leaves_len + }; + + if self + .checkpoints + .iter() + .skip(checkpoint_idx) + .any(|c| c.marked.contains(&position)) + { + // The requested position was marked after the checkpoint was created, so we + // cannot create a witness. + None + } else { + let mut path = vec![]; + + let mut leaf_idx: usize = position.into(); + for bit in 0..DEPTH { + leaf_idx ^= 1 << bit; + path.push(if leaf_idx < len { + let subtree_end = min(leaf_idx + (1 << bit), len); + root(&self.leaves[leaf_idx..subtree_end], bit)? + } else { + H::empty_root(Level::from(bit)) + }); + leaf_idx &= usize::MAX << (bit + 1); } - }) + + Some(path) + } + } else { + None + } } fn remove_mark(&mut self, position: Position) -> bool { - self.tree_state.remove_mark(position) + if self.marks.contains(&position) { + self.checkpoints + .back_mut() + .unwrap() + .forgotten + .insert(position); + true + } else { + false + } } fn checkpoint(&mut self) { - self.checkpoints.push(self.tree_state.clone()); + self.checkpoints + .push_back(Checkpoint::at_length(self.leaves.len())); if self.checkpoints.len() > self.max_checkpoints { self.drop_oldest_checkpoint(); } } fn rewind(&mut self) -> bool { - if let Some(checkpointed_state) = self.checkpoints.pop() { - self.tree_state = checkpointed_state; + if self.checkpoints.len() > 1 { + let c = self.checkpoints.pop_back().unwrap(); + self.leaves.truncate(c.leaves_len); + for pos in c.marked.iter() { + self.marks.remove(pos); + } true } else { false @@ -268,16 +267,16 @@ mod tests { expected = SipHashable::combine(lvl.into(), &expected, &expected); } - let tree = CompleteTree::::new(DEPTH as usize, 100); + let tree = CompleteTree::::new(100); assert_eq!(tree.root(0).unwrap(), expected); } #[test] fn correct_root() { - const DEPTH: usize = 3; + const DEPTH: u8 = 3; let values = (0..(1 << DEPTH)).into_iter().map(SipHashable); - let mut tree = CompleteTree::::new(DEPTH, 100); + let mut tree = CompleteTree::::new(100); for value in values { assert!(tree.append(value)); } @@ -302,20 +301,20 @@ mod tests { #[test] fn root_hashes() { - check_root_hashes(|max_c| CompleteTree::::new(4, max_c)); + check_root_hashes(CompleteTree::::new); } #[test] - fn witnesss() { - check_witnesses(|max_c| CompleteTree::::new(4, max_c)); + fn witness() { + check_witnesses(CompleteTree::::new); } #[test] fn correct_witness() { - const DEPTH: usize = 3; + const DEPTH: u8 = 3; let values = (0..(1 << DEPTH)).into_iter().map(SipHashable); - let mut tree = CompleteTree::::new(DEPTH, 100); + let mut tree = CompleteTree::::new(100); for value in values { assert!(tree.append(value)); tree.mark(); @@ -340,7 +339,7 @@ mod tests { for i in 0u64..(1 << DEPTH) { let position = Position::try_from(i).unwrap(); - let path = tree.witness(position, &tree.root(0).unwrap()).unwrap(); + let path = tree.witness(position, 0).unwrap(); assert_eq!( compute_root_from_witness(SipHashable(i), position, &path), expected @@ -350,11 +349,11 @@ mod tests { #[test] fn checkpoint_rewind() { - check_checkpoint_rewind(|max_c| CompleteTree::::new(4, max_c)); + check_checkpoint_rewind(CompleteTree::::new); } #[test] fn rewind_remove_mark() { - check_rewind_remove_mark(|max_c| CompleteTree::::new(4, max_c)); + check_rewind_remove_mark(CompleteTree::::new); } } From 07a88c34ac6db20f428a1efbc4b5390788ad3c19 Mon Sep 17 00:00:00 2001 From: Kris Nuttycombe Date: Thu, 15 Dec 2022 12:19:59 -0700 Subject: [PATCH 2/6] Adds a `Retention` enum to the `incrementalmerkletree` crate. This will be necessary in order to make the testing infrastructure reusable for shardtree. --- incrementalmerkletree/src/lib.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/incrementalmerkletree/src/lib.rs b/incrementalmerkletree/src/lib.rs index 9f667ce..a59668d 100644 --- a/incrementalmerkletree/src/lib.rs +++ b/incrementalmerkletree/src/lib.rs @@ -10,6 +10,28 @@ use std::ops::{Add, AddAssign, Range, Sub}; #[cfg(feature = "test-dependencies")] pub mod testing; +/// A type for metadata that is used to determine when and how a leaf can be pruned from a tree. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub enum Retention { + Ephemeral, + Checkpoint { id: C, is_marked: bool }, + Marked, +} + +impl Retention { + pub fn is_checkpoint(&self) -> bool { + matches!(self, Retention::Checkpoint { .. }) + } + + pub fn is_marked(&self) -> bool { + match self { + Retention::Ephemeral => false, + Retention::Checkpoint { is_marked, .. } => *is_marked, + Retention::Marked => true, + } + } +} + /// A type representing the position of a leaf in a Merkle tree. #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] #[repr(transparent)] From 14bb9c6b1bd2699496bd2fbafd51bbad326ddf34 Mon Sep 17 00:00:00 2001 From: Kris Nuttycombe Date: Thu, 15 Dec 2022 15:02:44 -0700 Subject: [PATCH 3/6] Modify test infrastructure to allow shardtree testing. This removes the `mark` function from the `tree` interface, in favor of always appending nodes as marked, as that's what's needed in practice, rather than the ability to mutably mark the latest position at arbitrary states of the tree. This removal does mean that a few patterns of interleaved mark and checkpoint operations are no longer being tested in the shared test suite; however, such interleaving of operations is not something that we should need to support anyway. --- bridgetree/proptest-regressions/lib.txt | 2 + bridgetree/src/lib.rs | 152 +++-- incrementalmerkletree/src/lib.rs | 11 + incrementalmerkletree/src/testing.rs | 567 +++++++++--------- .../src/testing/complete_tree.rs | 271 ++++++--- 5 files changed, 550 insertions(+), 453 deletions(-) diff --git a/bridgetree/proptest-regressions/lib.txt b/bridgetree/proptest-regressions/lib.txt index 72c9155..be5eefa 100644 --- a/bridgetree/proptest-regressions/lib.txt +++ b/bridgetree/proptest-regressions/lib.txt @@ -10,3 +10,5 @@ cc f1a0c0e8114919f9f675e0b31cdecf37af31579e365119312a7ebefcabb639f1 # shrinks to cc ac8dad1a92cb9563a291802cbd3dfca2a89da35fe4ed377701f5ad85b43700f4 # shrinks to ops = [Append("a"), Append("a"), Checkpoint, Checkpoint, Checkpoint, Mark, Checkpoint, Authpath(Position(1), 2)] cc cebb95fe896dc1d1e3c9a65efd50e786e6a4a3503c86fb2a6817bb05d25e751e # shrinks to tree = BridgeTree { depth: 8, prior_bridges: [MerkleBridge { prior_position: None, tracking: {Address { level: Level(2), index: 0 }}, ommers: {}, frontier: NonEmptyFrontier { position: Position(3), leaf: "a", ommers: ["a", "aa"] } }], current_bridge: Some(MerkleBridge { prior_position: Some(Position(3)), tracking: {Address { level: Level(2), index: 0 }}, ommers: {}, frontier: NonEmptyFrontier { position: Position(3), leaf: "a", ommers: ["a", "aa"] } }), saved: {Position(3): 0}, checkpoints: [Checkpoint { bridges_len: 0, marked: {Position(3)}, forgotten: {Position(3): 0} }], max_checkpoints: 10 } cc cdab33688ef9270768481b72d1615ff1d209fdb3d8d45be2ad564a8d5e0addc1 # shrinks to ops = [Append(SipHashable(0)), Mark, Mark, Checkpoint, Mark, Rewind, Append(SipHashable(0)), Mark] +cc 4058fadeb645f982bb83d666ea20c9541c9b920438877c145ee750cb09e22337 # shrinks to ops = [Append(SipHashable(0), Ephemeral), Append(SipHashable(0), Ephemeral), CurrentLeaf, CurrentLeaf, Append(SipHashable(0), Ephemeral), Append(SipHashable(0), Ephemeral), Append(SipHashable(0), Checkpoint { id: (), is_marked: true }), Unmark(Position(4))] +cc befc65ab8360534feebaa4c0e3da24518d134b03326235f1ccede15ad998d066 # shrinks to ops = [Append("a", Ephemeral), Append("a", Checkpoint { id: (), is_marked: true }), Append("a", Ephemeral), Append("a", Checkpoint { id: (), is_marked: false }), Authpath(Position(1), 0)] diff --git a/bridgetree/src/lib.rs b/bridgetree/src/lib.rs index 75961d0..883c7e5 100644 --- a/bridgetree/src/lib.rs +++ b/bridgetree/src/lib.rs @@ -37,7 +37,7 @@ use std::fmt::Debug; use std::mem::size_of; use std::ops::Range; -pub use incrementalmerkletree::{Address, Hashable, Level, Position}; +pub use incrementalmerkletree::{Address, Hashable, Level, Position, Retention}; /// Validation errors that can occur during reconstruction of a Merkle frontier from /// its constituent parts. @@ -614,6 +614,8 @@ impl<'a, H: Hashable + Ord + Clone + 'a> MerkleBridge { /// previous state. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct Checkpoint { + /// The unique identifier for this checkpoint. + id: usize, /// The number of bridges that will be retained in a rewind. bridges_len: usize, /// A set of the positions that have been marked during the period that this @@ -628,11 +630,13 @@ pub struct Checkpoint { impl Checkpoint { /// Creates a new checkpoint from its constituent parts. pub fn from_parts( + id: usize, bridges_len: usize, marked: BTreeSet, forgotten: BTreeSet, ) -> Self { Self { + id, bridges_len, marked, forgotten, @@ -640,14 +644,21 @@ impl Checkpoint { } /// Creates a new empty checkpoint for the specified [`BridgeTree`] state. - pub fn at_length(bridges_len: usize) -> Self { + pub fn at_length(bridges_len: usize, id: usize) -> Self { Checkpoint { + id, bridges_len, marked: BTreeSet::new(), forgotten: BTreeSet::new(), } } + /// The unique identifier for the checkpoint, which is simply an automatically incrementing + /// index over all checkpoints that have ever been created in the history of the tree. + pub fn id(&self) -> usize { + self.id + } + /// Returns the length of the [`BridgeTree::prior_bridges`] vector of the [`BridgeTree`] to /// which this checkpoint refers. /// @@ -750,13 +761,13 @@ impl BridgeTree { /// /// Panics if `max_checkpoints < 1` because mark/rewind logic depends upon the presence /// of checkpoints to function. - pub fn new(max_checkpoints: usize) -> Self { + pub fn new(max_checkpoints: usize, initial_checkpoint_id: usize) -> Self { assert!(max_checkpoints >= 1); Self { prior_bridges: vec![], current_bridge: None, saved: BTreeMap::new(), - checkpoints: VecDeque::from(vec![Checkpoint::at_length(0)]), + checkpoints: VecDeque::from(vec![Checkpoint::at_length(0, initial_checkpoint_id)]), max_checkpoints, } } @@ -1033,34 +1044,42 @@ impl BridgeTree { } } - /// Creates a new checkpoint for the current tree state. It is valid to - /// have multiple checkpoints for the same tree state, and each `rewind` - /// call will remove a single checkpoint. - pub fn checkpoint(&mut self) { - match self.current_bridge.take() { - Some(cur_b) => { - // Do not create a duplicate bridge - if self - .prior_bridges - .last() - .map_or(false, |pb| pb.position() == cur_b.position()) - { - self.current_bridge = Some(cur_b); - } else { - self.current_bridge = Some(cur_b.successor(false)); - self.prior_bridges.push(cur_b); + /// Creates a new checkpoint for the current tree state, with the given identifier. + /// + /// It is valid to have multiple checkpoints for the same tree state, and each `rewind` call + /// will remove a single checkpoint. Successive checkpoint identifiers must always be provided + /// in increasing order. + pub fn checkpoint(&mut self, id: usize) -> bool { + if Some(id) > self.checkpoints.back().map(|c| c.id) { + match self.current_bridge.take() { + Some(cur_b) => { + // Do not create a duplicate bridge + if self + .prior_bridges + .last() + .map_or(false, |pb| pb.position() == cur_b.position()) + { + self.current_bridge = Some(cur_b); + } else { + self.current_bridge = Some(cur_b.successor(false)); + self.prior_bridges.push(cur_b); + } + + self.checkpoints + .push_back(Checkpoint::at_length(self.prior_bridges.len(), id)); + } + None => { + self.checkpoints.push_back(Checkpoint::at_length(0, id)); } - - self.checkpoints - .push_back(Checkpoint::at_length(self.prior_bridges.len())); } - None => { - self.checkpoints.push_back(Checkpoint::at_length(0)); - } - } - if self.checkpoints.len() > self.max_checkpoints { - self.drop_oldest_checkpoint(); + if self.checkpoints.len() > self.max_checkpoints { + self.drop_oldest_checkpoint(); + } + + true + } else { + false } } @@ -1283,9 +1302,18 @@ mod tests { } } - impl Tree for BridgeTree { - fn append(&mut self, value: H) -> bool { - BridgeTree::append(self, value) + impl Tree for BridgeTree { + fn append(&mut self, value: H, retention: Retention) -> bool { + let appended = BridgeTree::append(self, value); + if appended { + if retention.is_marked() { + BridgeTree::mark(self); + } + if let Retention::Checkpoint { id, .. } = retention { + BridgeTree::checkpoint(self, id); + } + } + appended } fn depth(&self) -> u8 { @@ -1304,10 +1332,6 @@ mod tests { BridgeTree::get_marked_leaf(self, position) } - fn mark(&mut self) -> Option { - BridgeTree::mark(self) - } - fn marked_positions(&self) -> BTreeSet { BridgeTree::marked_positions(self) } @@ -1324,8 +1348,8 @@ mod tests { BridgeTree::remove_mark(self, position) } - fn checkpoint(&mut self) { - BridgeTree::checkpoint(self) + fn checkpoint(&mut self, id: usize) -> bool { + BridgeTree::checkpoint(self, id) } fn rewind(&mut self) -> bool { @@ -1430,7 +1454,7 @@ mod tests { #[test] fn tree_depth() { - let mut tree = BridgeTree::::new(100); + let mut tree = BridgeTree::::new(100, 0); for c in 'a'..'i' { assert!(tree.append(c.to_string())) } @@ -1441,8 +1465,8 @@ mod tests { mut tree: BridgeTree, ) { // Add checkpoints until we're sure everything that can be gc'ed will be gc'ed - for _ in 0..tree.max_checkpoints { - tree.checkpoint(); + for i in 0..tree.max_checkpoints { + tree.checkpoint(i + 1); } let mut tree_mut = tree.clone(); @@ -1462,9 +1486,9 @@ mod tests { { proptest::collection::vec(arb_operation(item_gen, 0..max_count), 0..max_count).prop_map( |ops| { - let mut tree: BridgeTree = BridgeTree::new(10); - for op in ops { - apply_operation(&mut tree, op); + let mut tree: BridgeTree = BridgeTree::new(10, 0); + for (i, op) in ops.into_iter().enumerate() { + apply_operation(&mut tree, op.map_checkpoint_id(|_| i)); } tree }, @@ -1498,45 +1522,47 @@ mod tests { #[test] fn root_hashes() { - check_root_hashes(BridgeTree::::new); + check_root_hashes(|max_checkpoints| BridgeTree::::new(max_checkpoints, 0)); } #[test] - fn witnesss() { - check_witnesses(BridgeTree::::new); + fn witness() { + check_witnesses(|max_checkpoints| BridgeTree::::new(max_checkpoints, 0)); } #[test] fn checkpoint_rewind() { - check_checkpoint_rewind(BridgeTree::::new); + check_checkpoint_rewind(|max_checkpoints| BridgeTree::::new(max_checkpoints, 0)); } #[test] fn rewind_remove_mark() { - check_rewind_remove_mark(BridgeTree::::new); + check_rewind_remove_mark(|max_checkpoints| { + BridgeTree::::new(max_checkpoints, 0) + }); } #[test] fn garbage_collect() { - let mut tree: BridgeTree = BridgeTree::new(100); + let mut tree: BridgeTree = BridgeTree::new(1000, 0); let empty_root = tree.root(0); tree.append("a".to_string()); - for _ in 0..100 { - tree.checkpoint(); + for i in 0..100 { + tree.checkpoint(i + 1); } tree.garbage_collect(); assert!(tree.root(0) != empty_root); tree.rewind(); assert!(tree.root(0) != empty_root); - let mut t = BridgeTree::::new(10); + let mut t = BridgeTree::::new(10, 0); let mut to_unmark = vec![]; let mut has_witness = vec![]; for i in 0usize..100 { let elem: String = format!("{},", i); assert!(t.append(elem), "Append should succeed."); if i % 5 == 0 { - t.checkpoint(); + t.checkpoint(i + 1); } if i % 7 == 0 { t.mark(); @@ -1553,7 +1579,7 @@ mod tests { } // 32 = 20 (checkpointed) + 14 (marked) - 2 (marked & checkpointed) assert_eq!(t.prior_bridges().len(), 20 + 14 - 2); - let witnesss = has_witness + let witness = has_witness .iter() .map(|pos| match t.witness(*pos, 0) { Ok(path) => path, @@ -1563,20 +1589,20 @@ mod tests { t.garbage_collect(); // 20 = 32 - 10 (removed checkpoints) + 1 (not removed due to mark) - 3 (removed marks) assert_eq!(t.prior_bridges().len(), 32 - 10 + 1 - 3); - let retained_witnesss = has_witness + let retained_witness = has_witness .iter() .map(|pos| t.witness(*pos, 0).expect("Must be able to get auth path")) .collect::>(); - assert_eq!(witnesss, retained_witnesss); + assert_eq!(witness, retained_witness); } // Combined tree tests fn new_combined_tree( max_checkpoints: usize, - ) -> CombinedTree, BridgeTree> { + ) -> CombinedTree, BridgeTree> { CombinedTree::new( - CompleteTree::::new(max_checkpoints), - BridgeTree::::new(max_checkpoints), + CompleteTree::::new(max_checkpoints, 0), + BridgeTree::::new(max_checkpoints, 0), ) } @@ -1601,7 +1627,8 @@ mod tests { ) ) { let tree = new_combined_tree(100); - check_operations(tree, &ops)?; + let indexed_ops = ops.iter().enumerate().map(|(i, op)| op.map_checkpoint_id(|_| i + 1)).collect::>(); + check_operations(tree, &indexed_ops)?; } #[test] @@ -1612,7 +1639,8 @@ mod tests { ) ) { let tree = new_combined_tree(100); - check_operations(tree, &ops)?; + let indexed_ops = ops.iter().enumerate().map(|(i, op)| op.map_checkpoint_id(|_| i + 1)).collect::>(); + check_operations(tree, &indexed_ops)?; } } } diff --git a/incrementalmerkletree/src/lib.rs b/incrementalmerkletree/src/lib.rs index a59668d..27771b3 100644 --- a/incrementalmerkletree/src/lib.rs +++ b/incrementalmerkletree/src/lib.rs @@ -30,6 +30,17 @@ impl Retention { Retention::Marked => true, } } + + pub fn map<'a, D, F: Fn(&'a C) -> D>(&'a self, f: F) -> Retention { + match self { + Retention::Ephemeral => Retention::Ephemeral, + Retention::Checkpoint { id, is_marked } => Retention::Checkpoint { + id: f(id), + is_marked: *is_marked, + }, + Retention::Marked => Retention::Marked, + } + } } /// A type representing the position of a leaf in a Merkle tree. diff --git a/incrementalmerkletree/src/testing.rs b/incrementalmerkletree/src/testing.rs index 94a243c..fad30cf 100644 --- a/incrementalmerkletree/src/testing.rs +++ b/incrementalmerkletree/src/testing.rs @@ -5,7 +5,7 @@ use core::marker::PhantomData; use proptest::prelude::*; use std::collections::BTreeSet; -use crate::{Hashable, Level, Position}; +use crate::{Hashable, Level, Position, Retention}; pub mod complete_tree; @@ -28,14 +28,14 @@ pub trait Frontier { /// A Merkle tree that supports incremental appends, marking of /// leaf nodes for construction of witnesses, checkpoints and rollbacks. -pub trait Tree { +pub trait Tree { /// Returns the depth of the tree. fn depth(&self) -> u8; /// Appends a new value to the tree at the next available slot. /// Returns true if successful and false if the tree would exceed /// the maximum allowed depth. - fn append(&mut self, value: H) -> bool; + fn append(&mut self, value: H, retention: Retention) -> bool; /// Returns the most recently appended leaf value. fn current_position(&self) -> Option; @@ -47,12 +47,6 @@ pub trait Tree { /// a witness for it. fn get_marked_leaf(&self, position: Position) -> Option<&H>; - /// Marks the current leaf as one for which we're interested in producing - /// a witness. Returns an optional value containing the - /// current position if successful or if the current value was already - /// marked, or None if the tree is empty. - fn mark(&mut self) -> Option; - /// Return a set of all the positions for which we have marked. fn marked_positions(&self) -> BTreeSet; @@ -74,10 +68,13 @@ pub trait Tree { /// false if we were already not maintaining a mark at this position. fn remove_mark(&mut self, position: Position) -> bool; - /// Creates a new checkpoint for the current tree state. It is valid to - /// have multiple checkpoints for the same tree state, and each `rewind` - /// call will remove a single checkpoint. - fn checkpoint(&mut self); + /// Creates a new checkpoint for the current tree state. + /// + /// It is valid to have multiple checkpoints for the same tree state, and + /// each `rewind` call will remove a single checkpoint. Returns `false` + /// if the checkpoint identifier provided is less than or equal to the + /// maximum checkpoint identifier observed. + fn checkpoint(&mut self, id: C) -> bool; /// Rewinds the tree state to the previous checkpoint, and then removes /// that checkpoint record. If there are multiple checkpoints at a given @@ -138,15 +135,14 @@ impl Hashable for Option { // #[derive(Clone, Debug)] -pub enum Operation { - Append(A), +pub enum Operation { + Append(A, Retention), CurrentPosition, CurrentLeaf, - Mark, MarkedLeaf(Position), MarkedPositions, Unmark(Position), - Checkpoint, + Checkpoint(C), Rewind, Authpath(Position, usize), GarbageCollect, @@ -154,39 +150,35 @@ pub enum Operation { use Operation::*; -pub fn append_str(x: &str) -> Operation { - Operation::Append(x.to_string()) +pub fn append_str(x: &str, retention: Retention) -> Operation { + Operation::Append(x.to_string(), retention) } -pub fn unmark(pos: usize) -> Operation { +pub fn unmark(pos: usize) -> Operation { Operation::Unmark(Position::from(pos)) } -pub fn witness(pos: usize, depth: usize) -> Operation { +pub fn witness(pos: usize, depth: usize) -> Operation { Operation::Authpath(Position::from(pos), depth) } -impl Operation { - pub fn apply>(&self, tree: &mut T) -> Option<(Position, Vec)> { +impl Operation { + pub fn apply>(&self, tree: &mut T) -> Option<(Position, Vec)> { match self { - Append(a) => { - assert!(tree.append(a.clone()), "append failed"); + Append(a, r) => { + assert!(tree.append(a.clone(), r.clone()), "append failed"); None } CurrentPosition => None, CurrentLeaf => None, - Mark => { - assert!(tree.mark().is_some(), "mark failed"); - None - } MarkedLeaf(_) => None, MarkedPositions => None, Unmark(p) => { assert!(tree.remove_mark(*p), "remove mark failed"); None } - Checkpoint => { - tree.checkpoint(); + Checkpoint(id) => { + tree.checkpoint(id.clone()); None } Rewind => { @@ -198,25 +190,50 @@ impl Operation { } } - pub fn apply_all>(ops: &[Operation], tree: &mut T) -> Option<(Position, Vec)> { + pub fn apply_all>( + ops: &[Operation], + tree: &mut T, + ) -> Option<(Position, Vec)> { let mut result = None; for op in ops { result = op.apply(tree); } result } + + pub fn map_checkpoint_id D>(&self, f: F) -> Operation { + match self { + Append(a, r) => Append(a.clone(), r.map(f)), + CurrentPosition => CurrentPosition, + CurrentLeaf => CurrentLeaf, + MarkedLeaf(l) => MarkedLeaf(*l), + MarkedPositions => MarkedPositions, + Unmark(p) => Unmark(*p), + Checkpoint(id) => Checkpoint(f(id)), + Rewind => Rewind, + Authpath(p, d) => Authpath(*p, *d), + GarbageCollect => GarbageCollect, + } + } +} + +pub fn arb_retention() -> impl Strategy> { + prop_oneof![ + Just(Retention::Ephemeral), + any::().prop_map(|is_marked| Retention::Checkpoint { id: (), is_marked }), + Just(Retention::Marked), + ] } pub fn arb_operation( item_gen: G, pos_gen: impl Strategy + Clone, -) -> impl Strategy> +) -> impl Strategy> where G::Value: Clone + 'static, { prop_oneof![ - item_gen.prop_map(Operation::Append), - Just(Operation::Mark), + (item_gen, arb_retention()).prop_map(|(i, r)| Operation::Append(i, r)), prop_oneof![ Just(Operation::CurrentLeaf), Just(Operation::CurrentPosition), @@ -229,7 +246,7 @@ where pos_gen .clone() .prop_map(|i| Operation::Unmark(Position::from(i))), - Just(Operation::Checkpoint), + Just(Operation::Checkpoint(())), Just(Operation::Rewind), pos_gen .prop_flat_map(|i| (0usize..10) @@ -237,19 +254,16 @@ where ] } -pub fn apply_operation>(tree: &mut T, op: Operation) { +pub fn apply_operation>(tree: &mut T, op: Operation) { match op { - Append(value) => { - tree.append(value); - } - Mark => { - tree.mark(); + Append(value, r) => { + tree.append(value, r); } Unmark(position) => { tree.remove_mark(position); } - Checkpoint => { - tree.checkpoint(); + Checkpoint(id) => { + tree.checkpoint(id); } Rewind => { tree.rewind(); @@ -263,9 +277,9 @@ pub fn apply_operation>(tree: &mut T, op: Operation) { } } -pub fn check_operations>( +pub fn check_operations>( mut tree: T, - ops: &[Operation], + ops: &[Operation], ) -> Result<(), TestCaseError> { let mut tree_size = 0; let mut tree_values: Vec = vec![]; @@ -275,11 +289,14 @@ pub fn check_operations>( for op in ops { prop_assert_eq!(tree_size, tree_values.len()); match op { - Append(value) => { - if tree.append(value.clone()) { + Append(value, r) => { + if tree.append(value.clone(), r.clone()) { prop_assert!(tree_size < (1 << tree.depth())); tree_size += 1; tree_values.push(value.clone()); + if r.is_checkpoint() { + tree_checkpoints.push(tree_size); + } } else { prop_assert_eq!(tree_size, 1 << tree.depth()); } @@ -293,13 +310,6 @@ pub fn check_operations>( CurrentLeaf => { prop_assert_eq!(tree_values.last(), tree.current_leaf()); } - Mark => { - if tree.mark().is_some() { - prop_assert!(tree_size != 0); - } else { - prop_assert_eq!(tree_size, 0); - } - } MarkedLeaf(position) => { if tree.get_marked_leaf(*position).is_some() { prop_assert!(::from(*position) < tree_size); @@ -309,9 +319,9 @@ pub fn check_operations>( tree.remove_mark(*position); } MarkedPositions => {} - Checkpoint => { + Checkpoint(id) => { tree_checkpoints.push(tree_size); - tree.checkpoint(); + tree.checkpoint(id.clone()); } Rewind => { if tree.rewind() { @@ -379,31 +389,35 @@ pub fn compute_root_from_witness(value: H, position: Position, path // #[derive(Clone)] -pub struct CombinedTree, E: Tree> { +pub struct CombinedTree, E: Tree> { inefficient: I, efficient: E, - _phantom: PhantomData, + _phantom_h: PhantomData, + _phantom_c: PhantomData, } -impl, E: Tree> CombinedTree { +impl, E: Tree> CombinedTree { pub fn new(inefficient: I, efficient: E) -> Self { assert_eq!(inefficient.depth(), efficient.depth()); CombinedTree { inefficient, efficient, - _phantom: PhantomData, + _phantom_h: PhantomData, + _phantom_c: PhantomData, } } } -impl, E: Tree> Tree for CombinedTree { +impl, E: Tree> Tree + for CombinedTree +{ fn depth(&self) -> u8 { self.inefficient.depth() } - fn append(&mut self, value: H) -> bool { - let a = self.inefficient.append(value.clone()); - let b = self.efficient.append(value); + fn append(&mut self, value: H, retention: Retention) -> bool { + let a = self.inefficient.append(value.clone(), retention.clone()); + let b = self.efficient.append(value, retention); assert_eq!(a, b); a } @@ -436,16 +450,6 @@ impl, E: Tree> Tree for Comb a } - fn mark(&mut self) -> Option { - let a = self.inefficient.mark(); - let b = self.efficient.mark(); - assert_eq!(a, b); - let apos = self.inefficient.marked_positions(); - let bpos = self.efficient.marked_positions(); - assert_eq!(apos, bpos); - a - } - fn marked_positions(&self) -> BTreeSet { let a = self.inefficient.marked_positions(); let b = self.efficient.marked_positions(); @@ -467,9 +471,11 @@ impl, E: Tree> Tree for Comb a } - fn checkpoint(&mut self) { - self.inefficient.checkpoint(); - self.efficient.checkpoint(); + fn checkpoint(&mut self, id: C) -> bool { + let a = self.inefficient.checkpoint(id.clone()); + let b = self.efficient.checkpoint(id); + assert_eq!(a, b); + a } fn rewind(&mut self) -> bool { @@ -479,38 +485,42 @@ impl, E: Tree> Tree for Comb a } } + // // Shared example tests // -pub fn check_root_hashes, F: Fn(usize) -> T>(new_tree: F) { +pub fn check_root_hashes, F: Fn(usize) -> T>(new_tree: F) { let mut tree = new_tree(100); assert_eq!(tree.root(0).unwrap(), "________________"); - tree.append("a".to_string()); + tree.append("a".to_string(), Retention::Ephemeral); assert_eq!(tree.root(0).unwrap().len(), 16); assert_eq!(tree.root(0).unwrap(), "a_______________"); - tree.append("b".to_string()); + tree.append("b".to_string(), Retention::Ephemeral); assert_eq!(tree.root(0).unwrap(), "ab______________"); - tree.append("c".to_string()); + tree.append("c".to_string(), Retention::Ephemeral); assert_eq!(tree.root(0).unwrap(), "abc_____________"); let mut t = new_tree(100); - t.append("a".to_string()); - t.checkpoint(); - t.mark(); - t.append("a".to_string()); - t.append("a".to_string()); - t.append("a".to_string()); + t.append( + "a".to_string(), + Retention::Checkpoint { + id: 1, + is_marked: true, + }, + ); + t.append("a".to_string(), Retention::Ephemeral); + t.append("a".to_string(), Retention::Ephemeral); + t.append("a".to_string(), Retention::Ephemeral); assert_eq!(t.root(0).unwrap(), "aaaa____________"); } -pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new_tree: F) { +pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new_tree: F) { let mut tree = new_tree(100); - tree.append("a".to_string()); - tree.mark(); + tree.append("a".to_string(), Retention::Marked); assert_eq!( tree.witness(Position::from(0), 0), Some(vec![ @@ -521,7 +531,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ]) ); - tree.append("b".to_string()); + tree.append("b".to_string(), Retention::Ephemeral); assert_eq!( tree.witness(0.into(), 0), Some(vec![ @@ -532,8 +542,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ]) ); - tree.append("c".to_string()); - tree.mark(); + tree.append("c".to_string(), Retention::Marked); assert_eq!( tree.witness(Position::from(2), 0), Some(vec![ @@ -544,7 +553,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ]) ); - tree.append("d".to_string()); + tree.append("d".to_string(), Retention::Ephemeral); assert_eq!( tree.witness(Position::from(2), 0), Some(vec![ @@ -555,7 +564,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ]) ); - tree.append("e".to_string()); + tree.append("e".to_string(), Retention::Ephemeral); assert_eq!( tree.witness(Position::from(2), 0), Some(vec![ @@ -567,13 +576,12 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ); let mut tree = new_tree(100); - tree.append("a".to_string()); - tree.mark(); - for c in 'b'..'h' { - tree.append(c.to_string()); + tree.append("a".to_string(), Retention::Marked); + for c in 'b'..'g' { + tree.append(c.to_string(), Retention::Ephemeral); } - tree.mark(); - tree.append("h".to_string()); + tree.append("g".to_string(), Retention::Marked); + tree.append("h".to_string(), Retention::Ephemeral); assert_eq!( tree.witness(0.into(), 0), @@ -586,17 +594,13 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ); let mut tree = new_tree(100); - tree.append("a".to_string()); - tree.mark(); - tree.append("b".to_string()); - tree.append("c".to_string()); - tree.append("d".to_string()); - tree.mark(); - tree.append("e".to_string()); - tree.mark(); - tree.append("f".to_string()); - tree.mark(); - tree.append("g".to_string()); + tree.append("a".to_string(), Retention::Marked); + tree.append("b".to_string(), Retention::Ephemeral); + tree.append("c".to_string(), Retention::Ephemeral); + tree.append("d".to_string(), Retention::Marked); + tree.append("e".to_string(), Retention::Marked); + tree.append("f".to_string(), Retention::Marked); + tree.append("g".to_string(), Retention::Ephemeral); assert_eq!( tree.witness(Position::from(5), 0), @@ -609,11 +613,11 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ); let mut tree = new_tree(100); - for c in 'a'..'l' { - tree.append(c.to_string()); + for c in 'a'..'k' { + tree.append(c.to_string(), Retention::Ephemeral); } - tree.mark(); - tree.append('l'.to_string()); + tree.append('k'.to_string(), Retention::Marked); + tree.append('l'.to_string(), Retention::Ephemeral); assert_eq!( tree.witness(Position::from(10), 0), @@ -626,16 +630,20 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ); let mut tree = new_tree(100); - tree.append('a'.to_string()); - tree.mark(); - tree.checkpoint(); + assert!(tree.append( + 'a'.to_string(), + Retention::Checkpoint { + id: 1, + is_marked: true + } + )); assert!(tree.rewind()); - for c in 'b'..'f' { - tree.append(c.to_string()); + for c in 'b'..'e' { + tree.append(c.to_string(), Retention::Ephemeral); } - tree.mark(); + tree.append("e".to_string(), Retention::Marked); for c in 'f'..'i' { - tree.append(c.to_string()); + tree.append(c.to_string(), Retention::Ephemeral); } assert_eq!( tree.witness(0.into(), 0), @@ -648,17 +656,20 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ); let mut tree = new_tree(100); - tree.append('a'.to_string()); - tree.append('b'.to_string()); - tree.append('c'.to_string()); - tree.mark(); - tree.append('d'.to_string()); - tree.append('e'.to_string()); - tree.append('f'.to_string()); - tree.append('g'.to_string()); - tree.mark(); - tree.checkpoint(); - tree.append('h'.to_string()); + tree.append('a'.to_string(), Retention::Ephemeral); + tree.append('b'.to_string(), Retention::Ephemeral); + tree.append('c'.to_string(), Retention::Marked); + tree.append('d'.to_string(), Retention::Ephemeral); + tree.append('e'.to_string(), Retention::Ephemeral); + tree.append('f'.to_string(), Retention::Ephemeral); + assert!(tree.append( + 'g'.to_string(), + Retention::Checkpoint { + id: 1, + is_marked: true + } + )); + tree.append('h'.to_string(), Retention::Ephemeral); assert!(tree.rewind()); assert_eq!( tree.witness(Position::from(2), 0), @@ -671,20 +682,18 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ); let mut tree = new_tree(100); - tree.append('a'.to_string()); - tree.append('b'.to_string()); - tree.mark(); + tree.append('a'.to_string(), Retention::Ephemeral); + tree.append('b'.to_string(), Retention::Marked); assert_eq!(tree.witness(Position::from(0), 0), None); let mut tree = new_tree(100); - for c in 'a'..'n' { - tree.append(c.to_string()); + for c in 'a'..'m' { + tree.append(c.to_string(), Retention::Ephemeral); } - tree.mark(); - tree.append('n'.to_string()); - tree.mark(); - tree.append('o'.to_string()); - tree.append('p'.to_string()); + tree.append('m'.to_string(), Retention::Marked); + tree.append('n'.to_string(), Retention::Marked); + tree.append('o'.to_string(), Retention::Ephemeral); + tree.append('p'.to_string(), Retention::Ephemeral); assert_eq!( tree.witness(Position::from(12), 0), @@ -698,10 +707,9 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new let ops = ('a'..='l') .into_iter() - .map(|c| Append(c.to_string())) - .chain(Some(Mark)) - .chain(Some(Append('m'.to_string()))) - .chain(Some(Append('n'.to_string()))) + .map(|c| Append(c.to_string(), Retention::Marked)) + .chain(Some(Append('m'.to_string(), Retention::Ephemeral))) + .chain(Some(Append('n'.to_string(), Retention::Ephemeral))) .chain(Some(Authpath(11usize.into(), 0))) .collect::>(); @@ -720,76 +728,66 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> T>(new ); } -pub fn check_checkpoint_rewind, F: Fn(usize) -> T>(new_tree: F) { +pub fn check_checkpoint_rewind, F: Fn(usize) -> T>(new_tree: F) { let mut t = new_tree(100); assert!(!t.rewind()); let mut t = new_tree(100); - t.checkpoint(); + t.checkpoint(1); assert!(t.rewind()); let mut t = new_tree(100); - t.append("a".to_string()); - t.checkpoint(); - t.append("b".to_string()); - t.mark(); + t.append("a".to_string(), Retention::Ephemeral); + t.checkpoint(1); + t.append("b".to_string(), Retention::Marked); assert!(t.rewind()); assert_eq!(Some(Position::from(0)), t.current_position()); let mut t = new_tree(100); - t.append("a".to_string()); - t.mark(); - t.checkpoint(); + t.append("a".to_string(), Retention::Marked); + t.checkpoint(1); assert!(t.rewind()); let mut t = new_tree(100); - t.append("a".to_string()); - t.checkpoint(); - t.mark(); - t.append("a".to_string()); + t.append("a".to_string(), Retention::Marked); + t.checkpoint(1); + t.append("a".to_string(), Retention::Ephemeral); assert!(t.rewind()); assert_eq!(Some(Position::from(0)), t.current_position()); let mut t = new_tree(100); - t.append("a".to_string()); - t.checkpoint(); - t.checkpoint(); + t.append("a".to_string(), Retention::Ephemeral); + t.checkpoint(1); + t.checkpoint(2); assert!(t.rewind()); - t.append("b".to_string()); + t.append("b".to_string(), Retention::Ephemeral); assert!(t.rewind()); - t.append("b".to_string()); + t.append("b".to_string(), Retention::Ephemeral); assert_eq!(t.root(0).unwrap(), "ab______________"); } -pub fn check_remove_mark, F: Fn(usize) -> T>(new_tree: F) { +pub fn check_remove_mark, F: Fn(usize) -> T>(new_tree: F) { let samples = vec![ vec![ - append_str("a"), - append_str("a"), - Checkpoint, - Mark, + append_str("a", Retention::Ephemeral), + append_str( + "a", + Retention::Checkpoint { + id: 1, + is_marked: true, + }, + ), witness(1, 1), ], vec![ - append_str("a"), - append_str("a"), - append_str("a"), - append_str("a"), - Mark, - Checkpoint, + append_str("a", Retention::Ephemeral), + append_str("a", Retention::Ephemeral), + append_str("a", Retention::Ephemeral), + append_str("a", Retention::Marked), + Checkpoint(1), unmark(3), witness(3, 0), ], - vec![ - append_str("a"), - append_str("a"), - Checkpoint, - Checkpoint, - Checkpoint, - Mark, - Checkpoint, - witness(1, 3), - ], ]; for (i, sample) in samples.iter().enumerate() { @@ -803,35 +801,25 @@ pub fn check_remove_mark, F: Fn(usize) -> T>(new_tree: F) { } } -pub fn check_rewind_remove_mark, F: Fn(usize) -> T>(new_tree: F) { +pub fn check_rewind_remove_mark, F: Fn(usize) -> T>(new_tree: F) { // rewinding doesn't remove a mark let mut tree = new_tree(100); - tree.append("e".to_string()); - tree.mark(); - tree.checkpoint(); + tree.append("e".to_string(), Retention::Marked); + tree.checkpoint(1); assert!(tree.rewind()); assert!(tree.remove_mark(0usize.into())); - // the order of checkpoint & mark does not matter - let mut tree = new_tree(100); - tree.append("e".to_string()); - tree.checkpoint(); - tree.mark(); - assert!(tree.rewind()); - assert!(!tree.remove_mark(0usize.into())); - // use a maximum number of checkpoints of 1 let mut tree = new_tree(1); - tree.append("e".to_string()); - tree.mark(); - tree.checkpoint(); + tree.append("e".to_string(), Retention::Marked); + tree.checkpoint(1); assert!(tree.marked_positions().contains(&0usize.into())); - tree.append("f".to_string()); + tree.append("f".to_string(), Retention::Ephemeral); // simulate a spend of `e` at `f` assert!(tree.remove_mark(0usize.into())); // even though the mark has been staged for removal, it's not gone yet assert!(tree.marked_positions().contains(&0usize.into())); - tree.checkpoint(); + tree.checkpoint(2); // the newest checkpoint will have caused the oldest to roll off, and // so the forgotten node will be unmarked assert!(!tree.marked_positions().contains(&0usize.into())); @@ -842,42 +830,41 @@ pub fn check_rewind_remove_mark, F: Fn(usize) -> T>(new_tree: F) // chain state restoration. let samples = vec![ - vec![append_str("x"), Checkpoint, Mark, Rewind, unmark(0)], vec![ - append_str("d"), - Checkpoint, - Mark, + append_str("x", Retention::Marked), + Checkpoint(1), + Rewind, + unmark(0), + ], + vec![ + append_str("d", Retention::Marked), + Checkpoint(1), unmark(0), Rewind, unmark(0), ], vec![ - append_str("o"), - Checkpoint, - Mark, - Checkpoint, + append_str("o", Retention::Marked), + Checkpoint(1), + Checkpoint(2), unmark(0), Rewind, Rewind, ], vec![ - append_str("s"), - Mark, - append_str("m"), - Checkpoint, + append_str("s", Retention::Marked), + append_str("m", Retention::Ephemeral), + Checkpoint(1), unmark(0), Rewind, unmark(0), unmark(0), ], vec![ - append_str("a"), - Mark, - Checkpoint, - Mark, + append_str("a", Retention::Marked), + Checkpoint(1), Rewind, - append_str("a"), - Mark, + append_str("a", Retention::Marked), ], ]; @@ -892,131 +879,121 @@ pub fn check_rewind_remove_mark, F: Fn(usize) -> T>(new_tree: F) } } -pub fn check_witness_consistency, F: Fn(usize) -> T>(new_tree: F) { +pub fn check_witness_consistency, F: Fn(usize) -> T>(new_tree: F) { let samples = vec![ // Reduced examples vec![ - append_str("a"), - append_str("b"), - Checkpoint, - Mark, + append_str("a", Retention::Ephemeral), + append_str("b", Retention::Marked), + Checkpoint(1), witness(0, 1), ], vec![ - append_str("c"), - append_str("d"), - Mark, - Checkpoint, + append_str("c", Retention::Ephemeral), + append_str("d", Retention::Marked), + Checkpoint(1), witness(1, 1), ], vec![ - append_str("e"), - Checkpoint, - Mark, - append_str("f"), + append_str("e", Retention::Marked), + Checkpoint(1), + append_str("f", Retention::Ephemeral), witness(0, 1), ], vec![ - append_str("g"), - Mark, - Checkpoint, + append_str("g", Retention::Marked), + Checkpoint(1), unmark(0), - append_str("h"), + append_str("h", Retention::Ephemeral), witness(0, 0), ], vec![ - append_str("i"), - Checkpoint, - Mark, + append_str("i", Retention::Marked), + Checkpoint(1), unmark(0), - append_str("j"), + append_str("j", Retention::Ephemeral), witness(0, 0), ], vec![ - append_str("i"), - Mark, - append_str("j"), - Checkpoint, - append_str("k"), + append_str("i", Retention::Marked), + append_str("j", Retention::Ephemeral), + Checkpoint(1), + append_str("k", Retention::Ephemeral), witness(0, 1), ], vec![ - append_str("l"), - Checkpoint, - Mark, - Checkpoint, - append_str("m"), - Checkpoint, + append_str("l", Retention::Marked), + Checkpoint(1), + Checkpoint(2), + append_str("m", Retention::Ephemeral), + Checkpoint(3), witness(0, 2), ], - vec![Checkpoint, append_str("n"), Mark, witness(0, 1)], vec![ - append_str("a"), - Mark, - Checkpoint, - unmark(0), - Checkpoint, - append_str("b"), + Checkpoint(1), + append_str("n", Retention::Marked), witness(0, 1), ], vec![ - append_str("a"), - Mark, - append_str("b"), + append_str("a", Retention::Marked), + Checkpoint(1), unmark(0), - Checkpoint, + Checkpoint(2), + append_str("b", Retention::Ephemeral), + witness(0, 1), + ], + vec![ + append_str("a", Retention::Marked), + append_str("b", Retention::Ephemeral), + unmark(0), + Checkpoint(1), witness(0, 0), ], vec![ - append_str("a"), - Mark, - Checkpoint, + append_str("a", Retention::Marked), + Checkpoint(1), unmark(0), - Checkpoint, + Checkpoint(2), Rewind, - append_str("b"), + append_str("b", Retention::Ephemeral), witness(0, 0), ], vec![ - append_str("a"), - Mark, - Checkpoint, - Checkpoint, + append_str("a", Retention::Marked), + Checkpoint(1), + Checkpoint(2), Rewind, - append_str("a"), + append_str("a", Retention::Ephemeral), unmark(0), witness(0, 1), ], // Unreduced examples vec![ - append_str("o"), - append_str("p"), - Mark, - append_str("q"), - Checkpoint, + append_str("o", Retention::Ephemeral), + append_str("p", Retention::Marked), + append_str("q", Retention::Ephemeral), + Checkpoint(1), unmark(1), witness(1, 1), ], vec![ - append_str("r"), - append_str("s"), - append_str("t"), - Mark, - Checkpoint, + append_str("r", Retention::Ephemeral), + append_str("s", Retention::Ephemeral), + append_str("t", Retention::Marked), + Checkpoint(1), unmark(2), - Checkpoint, + Checkpoint(2), witness(2, 2), ], vec![ - append_str("u"), - Mark, - append_str("v"), - append_str("w"), - Checkpoint, + append_str("u", Retention::Marked), + append_str("v", Retention::Ephemeral), + append_str("w", Retention::Ephemeral), + Checkpoint(1), unmark(0), - append_str("x"), - Checkpoint, - Checkpoint, + append_str("x", Retention::Ephemeral), + Checkpoint(2), + Checkpoint(3), witness(0, 3), ], ]; diff --git a/incrementalmerkletree/src/testing/complete_tree.rs b/incrementalmerkletree/src/testing/complete_tree.rs index 6c41d05..06e6d1f 100644 --- a/incrementalmerkletree/src/testing/complete_tree.rs +++ b/incrementalmerkletree/src/testing/complete_tree.rs @@ -1,8 +1,8 @@ //! Sample implementation of the Tree interface. use std::cmp::min; -use std::collections::{BTreeSet, VecDeque}; +use std::collections::{BTreeMap, BTreeSet}; -use crate::{testing::Tree, Hashable, Level, Position}; +use crate::{testing::Tree, Hashable, Level, Position, Retention}; pub(crate) fn root(leaves: &[H], depth: u8) -> H { let empty_leaf = H::empty_leaf(); @@ -38,7 +38,7 @@ pub(crate) fn root(leaves: &[H], depth: u8) -> H { #[derive(Clone, Debug, PartialEq, Eq)] pub struct Checkpoint { - /// The number of leaves that will be retained + /// The number of leaves in the tree when the checkpoint was created. leaves_len: usize, /// A set of the positions that have been marked during the period that this /// checkpoint is the current checkpoint. @@ -50,7 +50,7 @@ pub struct Checkpoint { } impl Checkpoint { - pub fn at_length(leaves_len: usize) -> Self { + fn at_length(leaves_len: usize) -> Self { Checkpoint { leaves_len, marked: BTreeSet::new(), @@ -60,71 +60,69 @@ impl Checkpoint { } #[derive(Clone, Debug)] -pub struct CompleteTree { +pub struct CompleteTree { leaves: Vec>, marks: BTreeSet, - checkpoints: VecDeque, + checkpoints: BTreeMap, max_checkpoints: usize, } -impl CompleteTree { +impl CompleteTree { /// Creates a new, empty binary tree - pub fn new(max_checkpoints: usize) -> Self { + pub fn new(max_checkpoints: usize, initial_checkpoint_id: C) -> Self { Self { leaves: vec![], marks: BTreeSet::new(), - checkpoints: VecDeque::from(vec![Checkpoint::at_length(0)]), + checkpoints: BTreeMap::from([(initial_checkpoint_id, Checkpoint::at_length(0))]), max_checkpoints, } } - fn append(&mut self, value: H) -> bool { - if self.leaves.len() == (1 << DEPTH) { - false - } else { - self.leaves.push(Some(value)); - true - } - } - - fn leaves_at_checkpoint_depth(&self, checkpoint_depth: usize) -> Option { - if checkpoint_depth == 0 { - Some(self.leaves.len()) - } else if checkpoint_depth <= self.checkpoints.len() { - self.checkpoints - .get(self.checkpoints.len() - checkpoint_depth) - .map(|c| c.leaves_len) - } else { - None - } - } -} - -impl CompleteTree { - /// Removes the oldest checkpoint. Returns true if successful and false if - /// there are fewer than `self.max_checkpoints` checkpoints. - fn drop_oldest_checkpoint(&mut self) -> bool { - if self.checkpoints.len() > self.max_checkpoints { - let c = self.checkpoints.pop_front().unwrap(); - for pos in c.forgotten.iter() { - self.marks.remove(pos); + /// Appends a new value to the tree at the next available slot. + /// + /// Returns true if successful and false if the tree is full or, for values with `Checkpoint` + /// retention, if a checkpoint id would be introduced that is less than or equal to the current + /// maximum checkpoint id. + fn append(&mut self, value: H, retention: Retention) -> Result<(), AppendError> { + fn append( + leaves: &mut Vec>, + value: H, + depth: u8, + ) -> Result<(), AppendError> { + if leaves.len() < (1 << depth) { + leaves.push(Some(value)); + Ok(()) + } else { + Err(AppendError::TreeFull) } - true - } else { - false } - } -} -impl Tree - for CompleteTree -{ - fn depth(&self) -> u8 { - DEPTH - } + match retention { + Retention::Marked => { + append(&mut self.leaves, value, DEPTH)?; + self.mark(); + } + Retention::Checkpoint { id, is_marked } => { + let latest_checkpoint = self.checkpoints.keys().rev().next(); + if Some(&id) > latest_checkpoint { + append(&mut self.leaves, value, DEPTH)?; + if is_marked { + self.mark(); + } + self.checkpoint(id, self.current_position()); + } else { + return Err(AppendError::CheckpointOutOfOrder { + current_max: latest_checkpoint.cloned(), + checkpoint: id, + }); + } + } + Retention::Ephemeral => { + append(&mut self.leaves, value, DEPTH)?; + } + } - fn append(&mut self, value: H) -> bool { - Self::append(self, value) + Ok(()) } fn current_position(&self) -> Option { @@ -135,6 +133,90 @@ impl Tree } } + fn mark(&mut self) -> Option { + match self.current_position() { + Some(pos) => { + if !self.marks.contains(&pos) { + self.marks.insert(pos); + self.checkpoints + .iter_mut() + .rev() + .next() + .unwrap() + .1 + .marked + .insert(pos); + } + Some(pos) + } + None => None, + } + } + + fn checkpoint(&mut self, id: C, pos: Option) { + self.checkpoints.insert( + id, + Checkpoint::at_length(pos.map_or_else(|| 0, |p| usize::from(p) + 1)), + ); + if self.checkpoints.len() > self.max_checkpoints { + self.drop_oldest_checkpoint(); + } + } + + fn leaves_at_checkpoint_depth(&self, checkpoint_depth: usize) -> Option { + if checkpoint_depth == 0 { + Some(self.leaves.len()) + } else { + self.checkpoints + .iter() + .rev() + .skip(checkpoint_depth - 1) + .map(|(_, c)| c.leaves_len) + .next() + } + } + + /// Removes the oldest checkpoint. Returns true if successful and false if + /// there are fewer than `self.max_checkpoints` checkpoints. + fn drop_oldest_checkpoint(&mut self) -> bool { + if self.checkpoints.len() > self.max_checkpoints { + let (id, c) = self.checkpoints.iter().next().unwrap(); + for pos in c.forgotten.iter() { + self.marks.remove(pos); + } + let id = id.clone(); // needed to avoid mutable/immutable borrow conflict + self.checkpoints.remove(&id); + true + } else { + false + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +enum AppendError { + TreeFull, + CheckpointOutOfOrder { + current_max: Option, + checkpoint: C, + }, +} + +impl Tree + for CompleteTree +{ + fn depth(&self) -> u8 { + DEPTH + } + + fn append(&mut self, value: H, retention: Retention) -> bool { + Self::append(self, value, retention).is_ok() + } + + fn current_position(&self) -> Option { + Self::current_position(self) + } + fn current_leaf(&self) -> Option<&H> { self.leaves.last().and_then(|opt: &Option| opt.as_ref()) } @@ -149,19 +231,6 @@ impl Tree } } - fn mark(&mut self) -> Option { - match self.current_position() { - Some(pos) => { - if !self.marks.contains(&pos) { - self.marks.insert(pos); - self.checkpoints.back_mut().unwrap().marked.insert(pos); - } - Some(pos) - } - None => None, - } - } - fn marked_positions(&self) -> BTreeSet { self.marks.clone() } @@ -173,18 +242,13 @@ impl Tree fn witness(&self, position: Position, checkpoint_depth: usize) -> Option> { if self.marks.contains(&position) && checkpoint_depth <= self.checkpoints.len() { - let checkpoint_idx = self.checkpoints.len() - checkpoint_depth; - let len = if checkpoint_depth == 0 { - self.leaves.len() - } else { - self.checkpoints[checkpoint_idx].leaves_len - }; - + let leaves_len = self.leaves_at_checkpoint_depth(checkpoint_depth)?; + let c_idx = self.checkpoints.len() - checkpoint_depth; if self .checkpoints .iter() - .skip(checkpoint_idx) - .any(|c| c.marked.contains(&position)) + .skip(c_idx) + .any(|(_, c)| c.marked.contains(&position)) { // The requested position was marked after the checkpoint was created, so we // cannot create a witness. @@ -195,8 +259,8 @@ impl Tree let mut leaf_idx: usize = position.into(); for bit in 0..DEPTH { leaf_idx ^= 1 << bit; - path.push(if leaf_idx < len { - let subtree_end = min(leaf_idx + (1 << bit), len); + path.push(if leaf_idx < leaves_len { + let subtree_end = min(leaf_idx + (1 << bit), leaves_len); root(&self.leaves[leaf_idx..subtree_end], bit)? } else { H::empty_root(Level::from(bit)) @@ -214,8 +278,11 @@ impl Tree fn remove_mark(&mut self, position: Position) -> bool { if self.marks.contains(&position) { self.checkpoints - .back_mut() + .iter_mut() + .rev() + .next() .unwrap() + .1 .forgotten .insert(position); true @@ -224,21 +291,24 @@ impl Tree } } - fn checkpoint(&mut self) { - self.checkpoints - .push_back(Checkpoint::at_length(self.leaves.len())); - if self.checkpoints.len() > self.max_checkpoints { - self.drop_oldest_checkpoint(); + fn checkpoint(&mut self, id: C) -> bool { + if Some(&id) > self.checkpoints.iter().rev().next().map(|(id, _)| id) { + Self::checkpoint(self, id, self.current_position()); + true + } else { + false } } fn rewind(&mut self) -> bool { if self.checkpoints.len() > 1 { - let c = self.checkpoints.pop_back().unwrap(); + let (id, c) = self.checkpoints.iter().rev().next().unwrap(); self.leaves.truncate(c.leaves_len); for pos in c.marked.iter() { self.marks.remove(pos); } + let id = id.clone(); // needed to avoid mutable/immutable borrow conflict + self.checkpoints.remove(&id); true } else { false @@ -256,7 +326,7 @@ mod tests { check_checkpoint_rewind, check_rewind_remove_mark, check_root_hashes, check_witnesses, compute_root_from_witness, SipHashable, Tree, }, - Hashable, Level, Position, + Hashable, Level, Position, Retention, }; #[test] @@ -267,7 +337,7 @@ mod tests { expected = SipHashable::combine(lvl.into(), &expected, &expected); } - let tree = CompleteTree::::new(100); + let tree = CompleteTree::::new(100, ()); assert_eq!(tree.root(0).unwrap(), expected); } @@ -276,11 +346,11 @@ mod tests { const DEPTH: u8 = 3; let values = (0..(1 << DEPTH)).into_iter().map(SipHashable); - let mut tree = CompleteTree::::new(100); + let mut tree = CompleteTree::::new(100, ()); for value in values { - assert!(tree.append(value)); + assert!(tree.append(value, Retention::Ephemeral).is_ok()); } - assert!(!tree.append(SipHashable(0))); + assert!(tree.append(SipHashable(0), Retention::Ephemeral).is_err()); let expected = SipHashable::combine( Level::from(2), @@ -301,25 +371,30 @@ mod tests { #[test] fn root_hashes() { - check_root_hashes(CompleteTree::::new); + check_root_hashes(|max_checkpoints| { + CompleteTree::::new(max_checkpoints, 0) + }); } #[test] fn witness() { - check_witnesses(CompleteTree::::new); + check_witnesses(|max_checkpoints| { + CompleteTree::::new(max_checkpoints, 0) + }); } #[test] fn correct_witness() { + use crate::{testing::Tree, Retention}; + const DEPTH: u8 = 3; let values = (0..(1 << DEPTH)).into_iter().map(SipHashable); - let mut tree = CompleteTree::::new(100); + let mut tree = CompleteTree::::new(100, ()); for value in values { - assert!(tree.append(value)); - tree.mark(); + assert!(Tree::append(&mut tree, value, Retention::Marked)); } - assert!(!tree.append(SipHashable(0))); + assert!(tree.append(SipHashable(0), Retention::Ephemeral).is_err()); let expected = SipHashable::combine( ::from(2), @@ -349,11 +424,15 @@ mod tests { #[test] fn checkpoint_rewind() { - check_checkpoint_rewind(CompleteTree::::new); + check_checkpoint_rewind(|max_checkpoints| { + CompleteTree::::new(max_checkpoints, 0) + }); } #[test] fn rewind_remove_mark() { - check_rewind_remove_mark(CompleteTree::::new); + check_rewind_remove_mark(|max_checkpoints| { + CompleteTree::::new(max_checkpoints, 0) + }); } } From 0e2329cebc0ae8ac1d1dab842b909295666450eb Mon Sep 17 00:00:00 2001 From: Kris Nuttycombe Date: Fri, 6 Jan 2023 10:58:35 -0700 Subject: [PATCH 4/6] Fix a missed Authpath -> Witness rename. --- incrementalmerkletree/src/testing.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/incrementalmerkletree/src/testing.rs b/incrementalmerkletree/src/testing.rs index fad30cf..641b31a 100644 --- a/incrementalmerkletree/src/testing.rs +++ b/incrementalmerkletree/src/testing.rs @@ -144,7 +144,7 @@ pub enum Operation { Unmark(Position), Checkpoint(C), Rewind, - Authpath(Position, usize), + Witness(Position, usize), GarbageCollect, } @@ -159,7 +159,7 @@ pub fn unmark(pos: usize) -> Operation { } pub fn witness(pos: usize, depth: usize) -> Operation { - Operation::Authpath(Position::from(pos), depth) + Operation::Witness(Position::from(pos), depth) } impl Operation { @@ -185,7 +185,7 @@ impl Operation { assert!(tree.rewind(), "rewind failed"); None } - Authpath(p, d) => tree.witness(*p, *d).map(|xs| (*p, xs)), + Witness(p, d) => tree.witness(*p, *d).map(|xs| (*p, xs)), GarbageCollect => None, } } @@ -211,7 +211,7 @@ impl Operation { Unmark(p) => Unmark(*p), Checkpoint(id) => Checkpoint(f(id)), Rewind => Rewind, - Authpath(p, d) => Authpath(*p, *d), + Witness(p, d) => Witness(*p, *d), GarbageCollect => GarbageCollect, } } @@ -250,7 +250,7 @@ where Just(Operation::Rewind), pos_gen .prop_flat_map(|i| (0usize..10) - .prop_map(move |depth| Operation::Authpath(Position::from(i), depth))), + .prop_map(move |depth| Operation::Witness(Position::from(i), depth))), ] } @@ -270,7 +270,7 @@ pub fn apply_operation>(tree: &mut T, op: Operation) { } CurrentPosition => {} CurrentLeaf => {} - Authpath(_, _) => {} + Witness(_, _) => {} MarkedLeaf(_) => {} MarkedPositions => {} GarbageCollect => {} @@ -331,7 +331,7 @@ pub fn check_operations>( tree_size = checkpointed_tree_size; } } - Authpath(position, depth) => { + Witness(position, depth) => { if let Some(path) = tree.witness(*position, *depth) { let value: H = tree_values[::from(*position)].clone(); let tree_root = tree.root(*depth); @@ -710,7 +710,7 @@ pub fn check_witnesses + std::fmt::Debug, F: Fn(usize) -> .map(|c| Append(c.to_string(), Retention::Marked)) .chain(Some(Append('m'.to_string(), Retention::Ephemeral))) .chain(Some(Append('n'.to_string(), Retention::Ephemeral))) - .chain(Some(Authpath(11usize.into(), 0))) + .chain(Some(Witness(11usize.into(), 0))) .collect::>(); let mut tree = new_tree(100); From 6867240e5ac9b552eb2732f3eb42fcd7a352628f Mon Sep 17 00:00:00 2001 From: Kris Nuttycombe Date: Tue, 27 Dec 2022 16:55:47 -0700 Subject: [PATCH 5/6] Remove the `current_leaf` method from `Tree` The notion of `current` doesn't make a lot of sense in the context of out-of-order insertion. --- bridgetree/src/lib.rs | 4 ---- incrementalmerkletree/src/testing.rs | 18 ------------------ .../src/testing/complete_tree.rs | 10 ++++------ 3 files changed, 4 insertions(+), 28 deletions(-) diff --git a/bridgetree/src/lib.rs b/bridgetree/src/lib.rs index 883c7e5..8d1d8e8 100644 --- a/bridgetree/src/lib.rs +++ b/bridgetree/src/lib.rs @@ -1324,10 +1324,6 @@ mod tests { BridgeTree::current_position(self) } - fn current_leaf(&self) -> Option<&H> { - BridgeTree::current_leaf(self) - } - fn get_marked_leaf(&self, position: Position) -> Option<&H> { BridgeTree::get_marked_leaf(self, position) } diff --git a/incrementalmerkletree/src/testing.rs b/incrementalmerkletree/src/testing.rs index 641b31a..de77b78 100644 --- a/incrementalmerkletree/src/testing.rs +++ b/incrementalmerkletree/src/testing.rs @@ -40,9 +40,6 @@ pub trait Tree { /// Returns the most recently appended leaf value. fn current_position(&self) -> Option; - /// Returns the most recently appended leaf value. - fn current_leaf(&self) -> Option<&H>; - /// Returns the leaf at the specified position if the tree can produce /// a witness for it. fn get_marked_leaf(&self, position: Position) -> Option<&H>; @@ -138,7 +135,6 @@ impl Hashable for Option { pub enum Operation { Append(A, Retention), CurrentPosition, - CurrentLeaf, MarkedLeaf(Position), MarkedPositions, Unmark(Position), @@ -170,7 +166,6 @@ impl Operation { None } CurrentPosition => None, - CurrentLeaf => None, MarkedLeaf(_) => None, MarkedPositions => None, Unmark(p) => { @@ -205,7 +200,6 @@ impl Operation { match self { Append(a, r) => Append(a.clone(), r.map(f)), CurrentPosition => CurrentPosition, - CurrentLeaf => CurrentLeaf, MarkedLeaf(l) => MarkedLeaf(*l), MarkedPositions => MarkedPositions, Unmark(p) => Unmark(*p), @@ -235,7 +229,6 @@ where prop_oneof![ (item_gen, arb_retention()).prop_map(|(i, r)| Operation::Append(i, r)), prop_oneof![ - Just(Operation::CurrentLeaf), Just(Operation::CurrentPosition), Just(Operation::MarkedPositions), ], @@ -269,7 +262,6 @@ pub fn apply_operation>(tree: &mut T, op: Operation) { tree.rewind(); } CurrentPosition => {} - CurrentLeaf => {} Witness(_, _) => {} MarkedLeaf(_) => {} MarkedPositions => {} @@ -307,9 +299,6 @@ pub fn check_operations>( prop_assert_eq!(tree_size - 1, pos.into()); } } - CurrentLeaf => { - prop_assert_eq!(tree_values.last(), tree.current_leaf()); - } MarkedLeaf(position) => { if tree.get_marked_leaf(*position).is_some() { prop_assert!(::from(*position) < tree_size); @@ -436,13 +425,6 @@ impl, E: Tree> a } - fn current_leaf(&self) -> Option<&H> { - let a = self.inefficient.current_leaf(); - let b = self.efficient.current_leaf(); - assert_eq!(a, b); - a - } - fn get_marked_leaf(&self, position: Position) -> Option<&H> { let a = self.inefficient.get_marked_leaf(position); let b = self.efficient.get_marked_leaf(position); diff --git a/incrementalmerkletree/src/testing/complete_tree.rs b/incrementalmerkletree/src/testing/complete_tree.rs index 06e6d1f..073da85 100644 --- a/incrementalmerkletree/src/testing/complete_tree.rs +++ b/incrementalmerkletree/src/testing/complete_tree.rs @@ -133,6 +133,8 @@ impl CompleteTr } } + /// Marks the current tree state leaf as a value that we're interested in + /// marking. Returns the current position if the tree is non-empty. fn mark(&mut self) -> Option { match self.current_position() { Some(pos) => { @@ -217,8 +219,8 @@ impl Option<&H> { - self.leaves.last().and_then(|opt: &Option| opt.as_ref()) + fn marked_positions(&self) -> BTreeSet { + self.marks.clone() } fn get_marked_leaf(&self, position: Position) -> Option<&H> { @@ -231,10 +233,6 @@ impl BTreeSet { - self.marks.clone() - } - fn root(&self, checkpoint_depth: usize) -> Option { self.leaves_at_checkpoint_depth(checkpoint_depth) .and_then(|len| root(&self.leaves[0..len], DEPTH)) From 0697df7d951f59614771db0d03c4fc56d87017f1 Mon Sep 17 00:00:00 2001 From: Kris Nuttycombe Date: Fri, 3 Mar 2023 17:02:03 -0700 Subject: [PATCH 6/6] Fix documentation errors. --- bridgetree/src/lib.rs | 10 ++++------ incrementalmerkletree/src/testing.rs | 8 +++----- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/bridgetree/src/lib.rs b/bridgetree/src/lib.rs index 8d1d8e8..3b39bc3 100644 --- a/bridgetree/src/lib.rs +++ b/bridgetree/src/lib.rs @@ -772,7 +772,7 @@ impl BridgeTree { } } - /// Removes the oldest checkpoin if there are more than `max_checkpoints`. Returns true if + /// Removes the oldest checkpoint if there are more than `max_checkpoints`. Returns true if /// successful and false if there are not enough checkpoints. fn drop_oldest_checkpoint(&mut self) -> bool { if self.checkpoints.len() > self.max_checkpoints { @@ -1111,11 +1111,9 @@ impl BridgeTree { } } - /// Obtains a witness to the value at the specified position, - /// as of the tree state corresponding to the given root. - /// Returns `None` if there is no available witness to that - /// position or if the root does not correspond to a checkpointed - /// root of the tree. + /// Obtains a witness for the value at the specified leaf position, as of the tree state at the + /// given checkpoint depth. Returns `None` if there is no witness information for the requested + /// position or if no checkpoint is available at the specified depth. pub fn witness( &self, position: Position, diff --git a/incrementalmerkletree/src/testing.rs b/incrementalmerkletree/src/testing.rs index de77b78..f511f51 100644 --- a/incrementalmerkletree/src/testing.rs +++ b/incrementalmerkletree/src/testing.rs @@ -53,11 +53,9 @@ pub trait Tree { /// requested checkpoint depth. fn root(&self, checkpoint_depth: usize) -> Option; - /// Obtains a witness to the value at the specified position, - /// as of the tree state corresponding to the given root. - /// Returns `None` if there is no available witness to that - /// position or if the root does not correspond to a checkpointed - /// root of the tree. + /// Obtains a witness for the value at the specified leaf position, as of the tree state at the + /// given checkpoint depth. Returns `None` if there is no witness information for the requested + /// position or if no checkpoint is available at the specified depth. fn witness(&self, position: Position, checkpoint_depth: usize) -> Option>; /// Marks the value at the specified position as a value we're no longer