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/bridgetree.txt deleted file mode 100644 index afa608a..0000000 --- a/bridgetree/proptest-regressions/bridgetree.txt +++ /dev/null @@ -1,8 +0,0 @@ -# Seeds for failure cases proptest has generated in the past. It is -# automatically read and these particular cases re-run before any -# novel cases are generated. -# -# It is recommended to check this file in to source control so that -# 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 } diff --git a/bridgetree/proptest-regressions/lib.txt b/bridgetree/proptest-regressions/lib.txt index 093b42f..be5eefa 100644 --- a/bridgetree/proptest-regressions/lib.txt +++ b/bridgetree/proptest-regressions/lib.txt @@ -4,5 +4,11 @@ # # It is recommended to check this file in to source control so that # everyone who runs the test benefits from these saved cases. -cc 190d23d28fc081e651e779d6209951363ee8a21752233cb72471626d14dd8bad # shrinks to ops = [Append("a"), Append("a"), Append("a")] -cc 81533fad8faadbdfdc9547e07f0491e40172a2f5fe4e8768d289389ed2e3cbcb # shrinks to ops = [Append(SipHashable(0)), Append(SipHashable(0)), Append(SipHashable(0))] +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] +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 daf39f4..3b39bc3 100644 --- a/bridgetree/src/lib.rs +++ b/bridgetree/src/lib.rs @@ -31,13 +31,13 @@ //! //! 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; 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. @@ -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(); } @@ -614,48 +614,51 @@ 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 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( + id: usize, bridges_len: usize, - is_marked: bool, marked: BTreeSet, - forgotten: BTreeMap, + forgotten: BTreeSet, ) -> Self { Self { + id, 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, id: usize) -> Self { Checkpoint { + id, bridges_len, - is_marked, marked: BTreeSet::new(), - forgotten: BTreeMap::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. /// @@ -665,15 +668,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 +676,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 +708,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 +723,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 +758,34 @@ pub enum BridgeTreeError { impl BridgeTree { /// Construct an empty BridgeTree value with the specified maximum number of checkpoints. - pub fn new(max_checkpoints: usize) -> Self { + /// + /// Panics if `max_checkpoints < 1` because mark/rewind logic depends upon the presence + /// of checkpoints to function. + 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: vec![], + checkpoints: VecDeque::from(vec![Checkpoint::at_length(0, initial_checkpoint_id)]), max_checkpoints, } } - /// Removes the oldest checkpoint. Returns true if successful and false if - /// there are no checkpoints. + /// 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.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 +806,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 +836,7 @@ impl BridgeTree { frontier, )), saved: BTreeMap::new(), - checkpoints: vec![], + checkpoints: VecDeque::new(), max_checkpoints, } } @@ -843,7 +847,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 +880,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 +973,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 +999,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,51 +1032,54 @@ 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 } } - /// 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) => { - let is_marked = self.get_marked_leaf(cur_b.position()).is_some(); + /// 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); + } - // 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(Checkpoint::at_length(self.prior_bridges.len(), is_marked)); } - None => { - self.checkpoints.push(Checkpoint::at_length(0, false)); - } - } - 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 } } @@ -1083,80 +1089,72 @@ 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 } } - /// 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. - 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> { + /// 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, + 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 +1196,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 +1210,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 +1220,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 +1284,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, }; @@ -1317,27 +1300,32 @@ 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 { + DEPTH } fn current_position(&self) -> Option { 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) } - fn mark(&mut self) -> Option { - BridgeTree::mark(self) - } - fn marked_positions(&self) -> BTreeSet { BridgeTree::marked_positions(self) } @@ -1346,16 +1334,16 @@ 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 { 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 { @@ -1460,13 +1448,29 @@ 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())) } 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 i in 0..tree.max_checkpoints { + tree.checkpoint(i + 1); + } + + 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, @@ -1476,9 +1480,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 }, @@ -1506,70 +1510,53 @@ 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); + 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 t = BridgeTree::::new(10); + let mut tree: BridgeTree = BridgeTree::new(1000, 0); + let empty_root = tree.root(0); + tree.append("a".to_string()); + 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, 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(); @@ -1586,9 +1573,9 @@ 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_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), }) @@ -1596,48 +1583,31 @@ 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, &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); + assert_eq!(witness, retained_witness); } // Combined tree tests fn new_combined_tree( max_checkpoints: usize, - ) -> CombinedTree, BridgeTree> { + ) -> CombinedTree, BridgeTree> { CombinedTree::new( - CompleteTree::new(4, max_checkpoints), - BridgeTree::::new(max_checkpoints), + CompleteTree::::new(max_checkpoints, 0), + BridgeTree::::new(max_checkpoints, 0), ) } #[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 +1621,8 @@ mod tests { ) ) { let tree = new_combined_tree(100); - check_operations(tree, 4, &ops)?; + let indexed_ops = ops.iter().enumerate().map(|(i, op)| op.map_checkpoint_id(|_| i + 1)).collect::>(); + check_operations(tree, &indexed_ops)?; } #[test] @@ -1662,7 +1633,8 @@ mod tests { ) ) { let tree = new_combined_tree(100); - check_operations(tree, 4, &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/CHANGELOG.md b/incrementalmerkletree/CHANGELOG.md index cfb25a4..f938ab8 100644 --- a/incrementalmerkletree/CHANGELOG.md +++ b/incrementalmerkletree/CHANGELOG.md @@ -41,7 +41,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 a773d10..08382d4 100644 --- a/incrementalmerkletree/src/lib.rs +++ b/incrementalmerkletree/src/lib.rs @@ -10,6 +10,39 @@ 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, + } + } + + 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. #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] #[repr(transparent)] @@ -373,7 +406,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..f511f51 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,28 +28,22 @@ 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; - /// 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>; - /// 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; @@ -59,22 +53,23 @@ 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. - fn witness(&self, position: Position, as_of_root: &H) -> Option>; + /// 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 /// 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; - /// 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 @@ -117,95 +112,121 @@ 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 // #[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), + Witness(Position, usize), GarbageCollect, } 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 { - Operation::Authpath(Position::from(pos), depth) +pub fn witness(pos: usize, depth: usize) -> Operation { + Operation::Witness(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 => { assert!(tree.rewind(), "rewind failed"); None } - Authpath(p, d) => tree - .root(*d) - .and_then(|root| tree.witness(*p, &root)) - .map(|xs| (*p, xs)), + Witness(p, d) => tree.witness(*p, *d).map(|xs| (*p, xs)), GarbageCollect => None, } } - 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, + MarkedLeaf(l) => MarkedLeaf(*l), + MarkedPositions => MarkedPositions, + Unmark(p) => Unmark(*p), + Checkpoint(id) => Checkpoint(f(id)), + Rewind => Rewind, + Witness(p, d) => Witness(*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), Just(Operation::MarkedPositions), ], @@ -216,44 +237,39 @@ 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) - .prop_map(move |depth| Operation::Authpath(Position::from(i), depth))), + .prop_map(move |depth| Operation::Witness(Position::from(i), depth))), ] } -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(); } CurrentPosition => {} - CurrentLeaf => {} - Authpath(_, _) => {} + Witness(_, _) => {} MarkedLeaf(_) => {} MarkedPositions => {} GarbageCollect => {} } } -pub fn check_operations>( +pub fn check_operations>( mut tree: T, - tree_depth: u8, - ops: &[Operation], + ops: &[Operation], ) -> Result<(), TestCaseError> { let mut tree_size = 0; let mut tree_values: Vec = vec![]; @@ -263,13 +279,16 @@ 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()) { - prop_assert!(tree_size < (1 << tree_depth)); + 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); + prop_assert_eq!(tree_size, 1 << tree.depth()); } } CurrentPosition => { @@ -278,16 +297,6 @@ pub fn check_operations>( prop_assert_eq!(tree_size - 1, pos.into()); } } - 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); @@ -297,9 +306,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() { @@ -309,8 +318,8 @@ pub fn check_operations>( tree_size = checkpointed_tree_size; } } - Authpath(position, depth) => { - if let Some(path) = tree.root(*depth).and_then(|r| tree.witness(*position, &r)) { + 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); @@ -324,11 +333,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!( @@ -368,26 +376,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 { - fn append(&mut self, value: H) -> bool { - let a = self.inefficient.append(value.clone()); - let b = self.efficient.append(value); +impl, E: Tree> Tree + for CombinedTree +{ + fn depth(&self) -> u8 { + self.inefficient.depth() + } + + 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 } @@ -406,13 +423,6 @@ impl, E: Tree> Tree for Comb 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); @@ -420,16 +430,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(); @@ -437,9 +437,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 } @@ -451,9 +451,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 { @@ -463,40 +465,44 @@ 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), &tree.root(0).unwrap()), + tree.witness(Position::from(0), 0), Some(vec![ "_".to_string(), "__".to_string(), @@ -505,9 +511,9 @@ 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(), &tree.root(0).unwrap()), + tree.witness(0.into(), 0), Some(vec![ "b".to_string(), "__".to_string(), @@ -516,10 +522,9 @@ 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), &tree.root(0).unwrap()), + tree.witness(Position::from(2), 0), Some(vec![ "_".to_string(), "ab".to_string(), @@ -528,9 +533,9 @@ 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), &tree.root(0).unwrap()), + tree.witness(Position::from(2), 0), Some(vec![ "d".to_string(), "ab".to_string(), @@ -539,9 +544,9 @@ 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), &tree.root(0).unwrap()), + tree.witness(Position::from(2), 0), Some(vec![ "d".to_string(), "ab".to_string(), @@ -551,16 +556,15 @@ 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(), &tree.root(0).unwrap()), + tree.witness(0.into(), 0), Some(vec![ "b".to_string(), "cd".to_string(), @@ -570,20 +574,16 @@ 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), &tree.root(0).unwrap()), + tree.witness(Position::from(5), 0), Some(vec![ "e".to_string(), "g_".to_string(), @@ -593,14 +593,14 @@ 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), &tree.root(0).unwrap()), + tree.witness(Position::from(10), 0), Some(vec![ "l".to_string(), "ij".to_string(), @@ -610,20 +610,23 @@ 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(), &tree.root(0).unwrap()), + tree.witness(0.into(), 0), Some(vec![ "b".to_string(), "cd".to_string(), @@ -633,21 +636,23 @@ 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), &tree.root(0).unwrap()), + tree.witness(Position::from(2), 0), Some(vec![ "d".to_string(), "ab".to_string(), @@ -657,26 +662,21 @@ 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(); - assert_eq!( - tree.witness(Position::from(0), &tree.root(0).unwrap()), - None - ); + 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), &tree.root(0).unwrap()), + tree.witness(Position::from(12), 0), Some(vec![ "n".to_string(), "op".to_string(), @@ -687,11 +687,10 @@ 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()))) - .chain(Some(Authpath(11usize.into(), 0))) + .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(Witness(11usize.into(), 0))) .collect::>(); let mut tree = new_tree(100); @@ -709,91 +708,101 @@ 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_rewind_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", Retention::Ephemeral), + append_str( + "a", + Retention::Checkpoint { + id: 1, + is_marked: true, + }, + ), + witness(1, 1), + ], + vec![ + 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), + ], + ]; + + 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(); - tree.checkpoint(); + tree.append("e".to_string(), Retention::Marked); + tree.checkpoint(1); assert!(tree.rewind()); assert!(tree.remove_mark(0usize.into())); - let mut tree = new_tree(100); - tree.append("e".to_string()); - tree.checkpoint(); - tree.mark(); - assert!(tree.rewind()); - assert!(!tree.remove_mark(0usize.into())); - - let mut tree = new_tree(100); - tree.append("e".to_string()); - tree.mark(); - tree.checkpoint(); + // use a maximum number of checkpoints of 1 + let mut tree = new_tree(1); + tree.append("e".to_string(), Retention::Marked); + tree.checkpoint(1); + assert!(tree.marked_positions().contains(&0usize.into())); + tree.append("f".to_string(), Retention::Ephemeral); + // 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())); - 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()); + // even though the mark has been staged for removal, it's not gone yet + assert!(tree.marked_positions().contains(&0usize.into())); + 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())); assert!(!tree.remove_mark(0usize.into())); // The following check_operations tests cover errors where the @@ -801,38 +810,46 @@ 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", Retention::Marked), + Checkpoint(1), + Rewind, + append_str("a", Retention::Marked), + ], ]; 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 {}: {:?}", @@ -842,179 +859,127 @@ 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), ], ]; 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..073da85 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::{BTreeMap, BTreeSet}; -use crate::{ - testing::{Frontier, Tree}, - Hashable, Level, Position, -}; +use crate::{testing::Tree, Hashable, Level, Position, Retention}; -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,277 @@ pub(crate) fn root(mut leaves: Vec) -> H { leaves[0].clone() } +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Checkpoint { + /// 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. + 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 { + 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, + checkpoints: BTreeMap, + max_checkpoints: usize, } -impl TreeState { - /// Creates a new, empty binary tree of specified depth. - pub fn new(depth: usize) -> Self { +impl CompleteTree { + /// Creates a new, empty binary tree + pub fn new(max_checkpoints: usize, initial_checkpoint_id: C) -> Self { Self { - leaves: vec![H::empty_leaf(); 1 << depth], - current_offset: 0, + leaves: vec![], 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 + checkpoints: BTreeMap::from([(initial_checkpoint_id, Checkpoint::at_length(0))]), + max_checkpoints, } } - /// Obtains the current root of this Merkle tree. - fn root(&self) -> H { - root(self.leaves.clone()) - } -} + /// 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) + } + } + + 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)?; + } + } + + Ok(()) + } -impl TreeState { fn current_position(&self) -> Option { - if self.current_offset == 0 { + if self.leaves.is_empty() { 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 + Some((self.leaves.len() - 1).into()) } } /// 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>, - max_checkpoints: usize, -} - -impl CompleteTree { - /// Creates a new, empty binary tree of specified depth. - pub fn new(depth: usize, max_checkpoints: usize) -> Self { - Self { - tree_state: TreeState::new(depth), - checkpoints: vec![], - 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() { - false - } else { - self.checkpoints.remove(0); - 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 { - self.checkpoints - .get(self.checkpoints.len() - checkpoint_depth) - } - } -} - -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) - } - - /// Returns the most recently appended leaf value. - fn current_position(&self) -> Option { - self.tree_state.current_position() - } - - fn current_leaf(&self) -> Option<&H> { - self.tree_state.current_leaf() - } - - fn get_marked_leaf(&self, position: Position) -> Option<&H> { - self.tree_state.get_marked_leaf(position) - } - - fn mark(&mut self) -> Option { - self.tree_state.mark() - } - - fn marked_positions(&self) -> BTreeSet { - self.tree_state.marks.clone() - } - - fn root(&self, checkpoint_depth: usize) -> Option { - self.tree_state_at_checkpoint_depth(checkpoint_depth) - .map(|s| s.root()) - } - - 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 + 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 remove_mark(&mut self, position: Position) -> bool { - self.tree_state.remove_mark(position) - } - - fn checkpoint(&mut self) { - self.checkpoints.push(self.tree_state.clone()); + 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 marked_positions(&self) -> BTreeSet { + self.marks.clone() + } + + fn get_marked_leaf(&self, position: Position) -> Option<&H> { + if self.marks.contains(&position) { + self.leaves + .get(usize::from(position)) + .and_then(|opt: &Option| opt.as_ref()) + } else { + None + } + } + + fn root(&self, checkpoint_depth: usize) -> Option { + self.leaves_at_checkpoint_depth(checkpoint_depth) + .and_then(|len| root(&self.leaves[0..len], DEPTH)) + } + + fn witness(&self, position: Position, checkpoint_depth: usize) -> Option> { + if self.marks.contains(&position) && checkpoint_depth <= self.checkpoints.len() { + let leaves_len = self.leaves_at_checkpoint_depth(checkpoint_depth)?; + let c_idx = self.checkpoints.len() - checkpoint_depth; + if self + .checkpoints + .iter() + .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. + 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 < 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)) + }); + leaf_idx &= usize::MAX << (bit + 1); + } + + Some(path) + } + } else { + None + } + } + + fn remove_mark(&mut self, position: Position) -> bool { + if self.marks.contains(&position) { + self.checkpoints + .iter_mut() + .rev() + .next() + .unwrap() + .1 + .forgotten + .insert(position); + true + } else { + false + } + } + + 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 let Some(checkpointed_state) = self.checkpoints.pop() { - self.tree_state = checkpointed_state; + if self.checkpoints.len() > 1 { + 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 @@ -257,7 +324,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] @@ -268,20 +335,20 @@ 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)); + 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), @@ -302,25 +369,30 @@ mod tests { #[test] fn root_hashes() { - check_root_hashes(|max_c| CompleteTree::::new(4, max_c)); + check_root_hashes(|max_checkpoints| { + CompleteTree::::new(max_checkpoints, 0) + }); } #[test] - fn witnesss() { - check_witnesses(|max_c| CompleteTree::::new(4, max_c)); + fn witness() { + check_witnesses(|max_checkpoints| { + CompleteTree::::new(max_checkpoints, 0) + }); } #[test] fn correct_witness() { - const DEPTH: usize = 3; + use crate::{testing::Tree, Retention}; + + 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(); + 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), @@ -340,7 +412,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 +422,15 @@ mod tests { #[test] fn checkpoint_rewind() { - check_checkpoint_rewind(|max_c| CompleteTree::::new(4, max_c)); + check_checkpoint_rewind(|max_checkpoints| { + CompleteTree::::new(max_checkpoints, 0) + }); } #[test] fn rewind_remove_mark() { - check_rewind_remove_mark(|max_c| CompleteTree::::new(4, max_c)); + check_rewind_remove_mark(|max_checkpoints| { + CompleteTree::::new(max_checkpoints, 0) + }); } }