fix(rotortree): read snapshot root at maximum depth

crates/rotortree/src/tree.rs

@@ -82,6 +82,9 @@ impl<const N: usize, const MAX_DEPTH: usize> TreeSnapshot<N, MAX_DEPTH> {
         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<const N: usize, const MAX_DEPTH: usize> TreeSnapshot<N, MAX_DEPTH> {
                 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)
     }
 }

crates/rotortree/tests/tree.rs

@@ -137,6 +137,81 @@ fn snapshot_get_node_and_level_len() {
     }
 }
 
+fn snapshots_at_max_depth<const N: usize, const DEPTH: usize>(
+    mut check: impl FnMut(&rotortree::TreeSnapshot<N, DEPTH>),
+) {
+    let capacity = N.pow(DEPTH as u32);
+    for batched in [false, true] {
+        for count in (capacity / N + 1)..=capacity {
+            // given
+            let mut tree = LeanIMT::<XorHasher, N, DEPTH>::new(XorHasher);
+            let leaves: Vec<Hash> = (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<const N: usize, const DEPTH: usize>(
+        snapshot: &rotortree::TreeSnapshot<N, DEPTH>,
+    ) {
+        // 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::<XorHasher, 2, 1>::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