diff --git a/crates/rotortree/src/tree.rs b/crates/rotortree/src/tree.rs index 48940ca..abcc0e5 100644 --- a/crates/rotortree/src/tree.rs +++ b/crates/rotortree/src/tree.rs @@ -82,6 +82,9 @@ impl TreeSnapshot { if level > self.depth { return 0; } + if level == self.depth { + return usize::from(self.root.is_some()); + } self.levels[level].len() } @@ -93,6 +96,16 @@ impl TreeSnapshot { size: 0, }); } + // `levels` has no root slot at `MAX_DEPTH`, so read the root from `self.root`. + if level == self.depth { + return self + .root + .filter(|_| index == 0) + .ok_or(TreeError::IndexOutOfRange { + index: index as u64, + size: u64::from(self.root.is_some()), + }); + } self.levels[level].get(index) } } diff --git a/crates/rotortree/tests/tree.rs b/crates/rotortree/tests/tree.rs index 6486df0..a867f3d 100644 --- a/crates/rotortree/tests/tree.rs +++ b/crates/rotortree/tests/tree.rs @@ -137,6 +137,81 @@ fn snapshot_get_node_and_level_len() { } } +fn snapshots_at_max_depth( + mut check: impl FnMut(&rotortree::TreeSnapshot), +) { + let capacity = N.pow(DEPTH as u32); + for batched in [false, true] { + for count in (capacity / N + 1)..=capacity { + // given + let mut tree = LeanIMT::::new(XorHasher); + let leaves: Vec = (0..count as u32).map(leaf).collect(); + if batched { + tree.insert_many(&leaves).unwrap(); + } else { + for hash in leaves { + tree.insert(hash).unwrap(); + } + } + // when + let snapshot = tree.snapshot(); + + // then + assert_eq!(snapshot.depth(), DEPTH); + check(&snapshot); + } + } +} + +/// Regression: `level_len` and `get_node` panicked with index out of bounds at the +/// root level once the depth reached `MAX_DEPTH`. +#[test] +fn snapshot_root_level_accessors_return_root_at_max_depth() { + fn check( + snapshot: &rotortree::TreeSnapshot, + ) { + // then + let depth = snapshot.depth(); + assert_eq!(snapshot.level_len(depth), 1); + assert_eq!(snapshot.get_node(depth, 0), Ok(snapshot.root().unwrap())); + assert_eq!( + snapshot.get_node(depth, 1), + Err(TreeError::IndexOutOfRange { index: 1, size: 1 }) + ); + } + + // given + snapshots_at_max_depth::<2, 1>(check); + snapshots_at_max_depth::<2, 3>(check); + snapshots_at_max_depth::<3, 2>(check); +} + +/// Control: below `MAX_DEPTH` the root level reads the same as before the fix. +#[test] +fn snapshot_root_access_preserves_empty_and_single_leaf_behavior() { + // given + let mut tree = LeanIMT::::new(XorHasher); + + // when + let empty = tree.snapshot(); + tree.insert(leaf(7)).unwrap(); + let single = tree.snapshot(); + + // then + assert_eq!(empty.level_len(0), 0); + assert_eq!( + empty.get_node(0, 0), + Err(TreeError::IndexOutOfRange { index: 0, size: 0 }) + ); + assert_eq!(single.depth(), 0); + assert_eq!(single.level_len(0), 1); + assert_eq!(single.get_node(0, 0), Ok(leaf(7))); + assert_eq!( + single.get_node(0, 1), + Err(TreeError::IndexOutOfRange { index: 1, size: 1 }) + ); +} + #[test] fn proof_for_last_leaf() { // given