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) + }); } }