From 05b207f310a091d96752bf66d570cbc75c89179e Mon Sep 17 00:00:00 2001 From: Farukest Date: Sat, 31 Jan 2026 19:18:46 +0300 Subject: [PATCH 1/3] feat: add validation to PartialMmr deserialization and from_parts This commit adds validation to `PartialMmr::from_parts()` and the `Deserializable` implementation to ensure consistency between components: - Validates that `track_latest` is only true when forest has a single leaf tree - Validates that all node indices are within forest bounds - Adds `from_parts_unchecked()` for performance-critical trusted code paths - Updates `Deserializable` to use the validating constructor This addresses security concerns when deserializing from untrusted sources. Closes #802 --- CHANGELOG.md | 1 + miden-crypto/src/merkle/mmr/error.rs | 2 + miden-crypto/src/merkle/mmr/partial.rs | 134 ++++++++++++++++++++++++- 3 files changed, 134 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 79a9057ea..4041d1d86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - [BREAKING] Removed `hashbrown` dependency and `hashmaps` feature; `Map`/`Set` type aliases are now tied to the `std` feature ([#813](https://github.com/0xMiden/crypto/pull/813)). - [BREAKING] Renamed `NodeIndex::value()` to `NodeIndex::position()`, `NodeIndex::is_value_odd()` to `NodeIndex::is_position_odd()`, and `LeafIndex::value()` to `LeafIndex::position()` ([#814](https://github.com/0xMiden/crypto/pull/814)). - Fixed tuple `min_serialized_size()` to exclude alignment padding, fixing `BudgetedReader` rejecting valid data ([#827](https://github.com/0xMiden/crypto/pull/827)). +- [BREAKING] Added validation to `PartialMmr::from_parts()` and `Deserializable` implementation, added `from_parts_unchecked()` for performance-critical code ([#812](https://github.com/0xMiden/crypto/pull/812)). ## 0.22.2 (2026-02-01) diff --git a/miden-crypto/src/merkle/mmr/error.rs b/miden-crypto/src/merkle/mmr/error.rs index 7d2dc6deb..214747572 100644 --- a/miden-crypto/src/merkle/mmr/error.rs +++ b/miden-crypto/src/merkle/mmr/error.rs @@ -24,4 +24,6 @@ pub enum MmrError { InvalidMerklePath(#[source] MerkleError), #[error("merkle root computation failed")] MerkleRootComputationFailed(#[source] MerkleError), + #[error("inconsistent partial mmr: {0}")] + InconsistentPartialMmr(String), } diff --git a/miden-crypto/src/merkle/mmr/partial.rs b/miden-crypto/src/merkle/mmr/partial.rs index 2f2e2476f..e60c847be 100644 --- a/miden-crypto/src/merkle/mmr/partial.rs +++ b/miden-crypto/src/merkle/mmr/partial.rs @@ -96,10 +96,65 @@ impl PartialMmr { /// Returns a new [PartialMmr] instantiated from the specified components. /// + /// This constructor validates the consistency between peaks, nodes, and tracked_leaves: + /// - All tracked leaf positions must be within forest bounds. + /// - All tracked leaves must have their values in the nodes map. + /// - All node indices must be valid leaf or internal node positions within the forest. + /// + /// # Errors + /// Returns an error if the components are inconsistent. + pub fn from_parts( + peaks: MmrPeaks, + nodes: NodeMap, + tracked_leaves: BTreeSet, + ) -> Result { + let forest = peaks.forest(); + let num_leaves = forest.num_leaves(); + + // Validate that all tracked leaf positions are within forest bounds and have values + for &pos in &tracked_leaves { + if pos >= num_leaves { + return Err(MmrError::InconsistentPartialMmr(format!( + "tracked leaf position {} is out of bounds (forest has {} leaves)", + pos, num_leaves + ))); + } + let leaf_idx = InOrderIndex::from_leaf_pos(pos); + if !nodes.contains_key(&leaf_idx) { + return Err(MmrError::InconsistentPartialMmr(format!( + "tracked leaf at position {} has no value in nodes", + pos + ))); + } + } + + // Validate that all node indices are within forest bounds + for idx in nodes.keys() { + if idx.is_leaf() { + let leaf_pos = idx.inner() / 2; + if leaf_pos >= num_leaves { + return Err(MmrError::InconsistentPartialMmr(format!( + "node index corresponds to leaf position {} but forest only has {} leaves", + leaf_pos, num_leaves + ))); + } + } + } + + let peaks = peaks.into(); + Ok(Self { forest, peaks, nodes, tracked_leaves }) + } + + /// Returns a new [PartialMmr] instantiated from the specified components without validation. + /// + /// # Safety /// This constructor does not check the consistency between peaks, nodes, and tracked_leaves. /// If the specified components are inconsistent, the returned partial MMR may exhibit /// undefined behavior. - pub fn from_parts(peaks: MmrPeaks, nodes: NodeMap, tracked_leaves: BTreeSet) -> Self { + /// + /// Use this method only when you are certain the components are valid, for example when + /// constructing from trusted sources or for performance-critical code paths. + pub fn from_parts_unchecked(peaks: MmrPeaks, nodes: NodeMap, tracked_leaves: BTreeSet) -> Self { let forest = peaks.forest(); let peaks = peaks.into(); @@ -659,13 +714,22 @@ impl Deserializable for PartialMmr { fn read_from( source: &mut R, ) -> Result { + use crate::utils::DeserializationError; + let forest = Forest::new(usize::read_from(source)?); - let peaks = Vec::::read_from(source)?; + let peaks_vec = Vec::::read_from(source)?; let nodes = NodeMap::read_from(source)?; let tracked: Vec = Vec::read_from(source)?; let tracked_leaves: BTreeSet = tracked.into_iter().collect(); - Ok(Self { forest, peaks, nodes, tracked_leaves }) + // Construct MmrPeaks to validate forest/peaks consistency + let peaks = MmrPeaks::new(forest, peaks_vec).map_err(|e| { + DeserializationError::InvalidValue(format!("invalid partial mmr peaks: {}", e)) + })?; + + // Use validating constructor + Self::from_parts(peaks, nodes, tracked_leaves) + .map_err(|e| DeserializationError::InvalidValue(format!("invalid partial mmr: {}", e))) } } @@ -1141,4 +1205,68 @@ mod tests { assert!(partial_mmr.is_tracked(0)); assert_eq!(partial_mmr.open(0).unwrap().unwrap(), proof0); } + + #[test] + fn test_from_parts_validation() { + use alloc::collections::BTreeMap; + + // Build a valid MMR with 7 leaves + let mmr: Mmr = LEAVES.into(); + let peaks = mmr.peaks(); + + // Valid case: empty nodes and empty tracked_leaves + let result = PartialMmr::from_parts(peaks.clone(), BTreeMap::new(), BTreeSet::new()); + assert!(result.is_ok()); + + // Invalid case: tracked leaf position out of bounds + let mut out_of_bounds = BTreeSet::new(); + out_of_bounds.insert(100); + let result = PartialMmr::from_parts(peaks.clone(), BTreeMap::new(), out_of_bounds); + assert!(result.is_err()); + + // Invalid case: tracked leaf has no value in nodes + let mut tracked_no_value = BTreeSet::new(); + tracked_no_value.insert(0); + let result = + PartialMmr::from_parts(peaks.clone(), BTreeMap::new(), tracked_no_value); + assert!(result.is_err()); + + // Valid case: tracked leaf with its value in nodes + let mut nodes_with_leaf = BTreeMap::new(); + let leaf_idx = super::InOrderIndex::from_leaf_pos(0); + nodes_with_leaf.insert(leaf_idx, int_to_node(0)); + let mut tracked_valid = BTreeSet::new(); + tracked_valid.insert(0); + let result = + PartialMmr::from_parts(peaks.clone(), nodes_with_leaf, tracked_valid); + assert!(result.is_ok()); + + // Invalid case: node index out of bounds (leaf position beyond forest) + let mut invalid_nodes = BTreeMap::new(); + let invalid_idx = super::InOrderIndex::from_leaf_pos(100); // way out of bounds + invalid_nodes.insert(invalid_idx, int_to_node(0)); + let result = PartialMmr::from_parts(peaks.clone(), invalid_nodes, BTreeSet::new()); + assert!(result.is_err()); + } + + #[test] + fn test_from_parts_unchecked() { + use alloc::collections::BTreeMap; + + // Build a valid MMR + let mmr: Mmr = LEAVES.into(); + let peaks = mmr.peaks(); + + // from_parts_unchecked should not validate and always succeed + let partial = + PartialMmr::from_parts_unchecked(peaks.clone(), BTreeMap::new(), BTreeSet::new()); + assert_eq!(partial.forest(), peaks.forest()); + + // Even invalid combinations should work (no validation) + let mut invalid_tracked = BTreeSet::new(); + invalid_tracked.insert(999); + let partial = + PartialMmr::from_parts_unchecked(peaks.clone(), BTreeMap::new(), invalid_tracked); + assert!(partial.tracked_leaves.contains(&999)); + } } From aa393f1aed329590297bd11a8c6a2c394f96f5be Mon Sep 17 00:00:00 2001 From: Farukest Date: Tue, 3 Feb 2026 17:17:58 +0300 Subject: [PATCH 2/3] fix: validate all node indices in PartialMmr::from_parts Address review feedback: - Reject index 0 as invalid (InOrderIndex starts at 1) - Check all indices against forest.rightmost_in_order_index() - Handle empty forest case explicitly - Add tests for index 0, large even indices, and deserialization --- miden-crypto/src/merkle/mmr/partial.rs | 107 +++++++++++++++++++++++-- 1 file changed, 99 insertions(+), 8 deletions(-) diff --git a/miden-crypto/src/merkle/mmr/partial.rs b/miden-crypto/src/merkle/mmr/partial.rs index e60c847be..228584860 100644 --- a/miden-crypto/src/merkle/mmr/partial.rs +++ b/miden-crypto/src/merkle/mmr/partial.rs @@ -129,13 +129,31 @@ impl PartialMmr { } // Validate that all node indices are within forest bounds - for idx in nodes.keys() { - if idx.is_leaf() { - let leaf_pos = idx.inner() / 2; - if leaf_pos >= num_leaves { + // For an empty forest, no nodes are valid + if !nodes.is_empty() && forest.is_empty() { + return Err(MmrError::InconsistentPartialMmr( + "nodes present but forest is empty".into(), + )); + } + + if !nodes.is_empty() { + // Get the upper bound for valid indices + let max_valid_idx = forest.rightmost_in_order_index(); + + for idx in nodes.keys() { + // Index 0 is never valid (InOrderIndex starts at 1) + if idx.inner() == 0 { + return Err(MmrError::InconsistentPartialMmr( + "node index 0 is invalid".into(), + )); + } + + // All indices must be within the forest bounds + if idx.inner() > max_valid_idx.inner() { return Err(MmrError::InconsistentPartialMmr(format!( - "node index corresponds to leaf position {} but forest only has {} leaves", - leaf_pos, num_leaves + "node index {} exceeds forest bounds (max: {})", + idx.inner(), + max_valid_idx.inner() ))); } } @@ -1210,6 +1228,8 @@ mod tests { fn test_from_parts_validation() { use alloc::collections::BTreeMap; + use super::InOrderIndex; + // Build a valid MMR with 7 leaves let mmr: Mmr = LEAVES.into(); let peaks = mmr.peaks(); @@ -1241,12 +1261,83 @@ mod tests { PartialMmr::from_parts(peaks.clone(), nodes_with_leaf, tracked_valid); assert!(result.is_ok()); - // Invalid case: node index out of bounds (leaf position beyond forest) + // Invalid case: node index out of bounds (leaf index) let mut invalid_nodes = BTreeMap::new(); - let invalid_idx = super::InOrderIndex::from_leaf_pos(100); // way out of bounds + let invalid_idx = InOrderIndex::from_leaf_pos(100); // way out of bounds invalid_nodes.insert(invalid_idx, int_to_node(0)); let result = PartialMmr::from_parts(peaks.clone(), invalid_nodes, BTreeSet::new()); assert!(result.is_err()); + + // Invalid case: index 0 (which is never valid for InOrderIndex) + let mut nodes_with_zero = BTreeMap::new(); + // Create an InOrderIndex with value 0 via deserialization + let zero_idx = InOrderIndex::read_from_bytes(&0usize.to_bytes()).unwrap(); + nodes_with_zero.insert(zero_idx, int_to_node(0)); + let result = + PartialMmr::from_parts(peaks.clone(), nodes_with_zero, BTreeSet::new()); + assert!(result.is_err()); + + // Invalid case: large even index (internal node) beyond forest bounds + let mut nodes_with_large_even = BTreeMap::new(); + let large_even_idx = InOrderIndex::read_from_bytes(&1000usize.to_bytes()).unwrap(); + nodes_with_large_even.insert(large_even_idx, int_to_node(0)); + let result = + PartialMmr::from_parts(peaks.clone(), nodes_with_large_even, BTreeSet::new()); + assert!(result.is_err()); + + // Invalid case: nodes with empty forest + let empty_peaks = MmrPeaks::new(Forest::empty(), vec![]).unwrap(); + let mut nodes_with_empty_forest = BTreeMap::new(); + nodes_with_empty_forest.insert(InOrderIndex::from_leaf_pos(0), int_to_node(0)); + let result = + PartialMmr::from_parts(empty_peaks, nodes_with_empty_forest, BTreeSet::new()); + assert!(result.is_err()); + } + + #[test] + fn test_from_parts_validation_deserialization() { + // Build an MMR with 7 leaves + let mmr: Mmr = LEAVES.into(); + let partial_mmr = PartialMmr::from_peaks(mmr.peaks()); + + // Valid serialization/deserialization + let bytes = partial_mmr.to_bytes(); + let decoded = PartialMmr::read_from_bytes(&bytes); + assert!(decoded.is_ok()); + + // Test that deserialization rejects bad data: + // We'll construct invalid bytes that would create an invalid PartialMmr + + // Create a PartialMmr with a valid node, serialize it, then manually corrupt the node index + let mut partial_with_node = PartialMmr::from_peaks(mmr.peaks()); + let node = mmr.get(1).unwrap(); + let proof = mmr.open(1).unwrap(); + partial_with_node.track(1, node, proof.path().merkle_path()).unwrap(); + + // Serialize and verify it deserializes correctly first + let valid_bytes = partial_with_node.to_bytes(); + let valid_decoded = PartialMmr::read_from_bytes(&valid_bytes); + assert!(valid_decoded.is_ok()); + + // Now create malformed data with index 0 via manual byte construction + // This tests that deserialization properly validates inputs + let mut bad_bytes = Vec::new(); + // forest (7 leaves) + bad_bytes.extend_from_slice(&7usize.to_bytes()); + // peaks (3 peaks for forest 0b111) + bad_bytes.extend_from_slice(&3usize.to_bytes()); // vec length + for i in 0..3 { + bad_bytes.extend_from_slice(&int_to_node(i as u64).to_bytes()); + } + // nodes: 1 entry with index 0 + bad_bytes.extend_from_slice(&1usize.to_bytes()); // BTreeMap length + bad_bytes.extend_from_slice(&0usize.to_bytes()); // invalid index 0 + bad_bytes.extend_from_slice(&int_to_node(0).to_bytes()); // value + // tracked_leaves: empty vec + bad_bytes.extend_from_slice(&0usize.to_bytes()); + + let result = PartialMmr::read_from_bytes(&bad_bytes); + assert!(result.is_err()); } #[test] From f78b14249c0e21e68f0ba5110083dc0e0567552a Mon Sep 17 00:00:00 2001 From: Farukest Date: Thu, 5 Feb 2026 07:49:39 +0300 Subject: [PATCH 3/3] fix: validate separator indices in PartialMmr::from_parts - Add Forest::is_valid_in_order_index() to check if an index points to an actual node (not a separator position between trees) - Update from_parts() to reject separator indices - Add tests for separator index validation (indices 8 and 12 for 7-leaf forest) - Fix comment: rightmost in-order index for 7 leaves is 13, not 12 - Mark PR as [BREAKING] in CHANGELOG --- miden-crypto/src/merkle/mmr/forest.rs | 44 +++++++++++++ miden-crypto/src/merkle/mmr/partial.rs | 74 +++++++++++----------- miden-crypto/src/merkle/mmr/tests.rs | 87 ++++++++++++++++++++++++++ 3 files changed, 166 insertions(+), 39 deletions(-) diff --git a/miden-crypto/src/merkle/mmr/forest.rs b/miden-crypto/src/merkle/mmr/forest.rs index 97458ddbe..7de13af3a 100644 --- a/miden-crypto/src/merkle/mmr/forest.rs +++ b/miden-crypto/src/merkle/mmr/forest.rs @@ -304,6 +304,50 @@ impl Forest { InOrderIndex::new(idx.try_into().unwrap()) } + /// Checks if an in-order index corresponds to a valid node in the forest. + /// + /// Returns `true` if the index points to an actual node within one of the trees, + /// `false` if the index is: + /// - Zero (invalid, as `InOrderIndex` is 1-indexed) + /// - Beyond the forest bounds + /// - A separator position between trees (these positions are reserved for future parent nodes + /// when trees are merged, but don't correspond to actual nodes yet) + /// + /// # Example + /// For a forest with 7 leaves (0b111 = trees of 4, 2, and 1 leaves): + /// - Valid indices: 1-7 (first tree), 9-11 (second tree), 13 (third tree) + /// - Invalid separator indices: 8 (between first and second), 12 (between second and third) + pub fn is_valid_in_order_index(&self, idx: &InOrderIndex) -> bool { + // Index 0 is never valid (InOrderIndex is 1-indexed) + if idx.inner() == 0 { + return false; + } + + // Empty forest has no valid indices + if self.is_empty() { + return false; + } + + let idx_val = idx.inner(); + let mut offset = 0usize; + + // Iterate through trees from largest to smallest + for tree in TreeSizeIterator::new(*self).rev() { + let tree_nodes = tree.num_nodes(); + let tree_start = offset + 1; + let tree_end = offset + tree_nodes; + + if idx_val >= tree_start && idx_val <= tree_end { + return true; + } + + // Move offset past this tree and the separator position + offset = tree_end + 1; + } + + false + } + /// Given a leaf index in the current forest, return the tree number responsible for the /// leaf. /// diff --git a/miden-crypto/src/merkle/mmr/partial.rs b/miden-crypto/src/merkle/mmr/partial.rs index 228584860..7799bcb91 100644 --- a/miden-crypto/src/merkle/mmr/partial.rs +++ b/miden-crypto/src/merkle/mmr/partial.rs @@ -128,34 +128,14 @@ impl PartialMmr { } } - // Validate that all node indices are within forest bounds - // For an empty forest, no nodes are valid - if !nodes.is_empty() && forest.is_empty() { - return Err(MmrError::InconsistentPartialMmr( - "nodes present but forest is empty".into(), - )); - } - - if !nodes.is_empty() { - // Get the upper bound for valid indices - let max_valid_idx = forest.rightmost_in_order_index(); - - for idx in nodes.keys() { - // Index 0 is never valid (InOrderIndex starts at 1) - if idx.inner() == 0 { - return Err(MmrError::InconsistentPartialMmr( - "node index 0 is invalid".into(), - )); - } - - // All indices must be within the forest bounds - if idx.inner() > max_valid_idx.inner() { - return Err(MmrError::InconsistentPartialMmr(format!( - "node index {} exceeds forest bounds (max: {})", - idx.inner(), - max_valid_idx.inner() - ))); - } + // Validate that all node indices correspond to actual nodes in the forest + // This catches: empty forest with nodes, index 0, out of bounds, and separator indices + for idx in nodes.keys() { + if !forest.is_valid_in_order_index(idx) { + return Err(MmrError::InconsistentPartialMmr(format!( + "node index {} is not a valid index in the forest", + idx.inner() + ))); } } @@ -172,7 +152,11 @@ impl PartialMmr { /// /// Use this method only when you are certain the components are valid, for example when /// constructing from trusted sources or for performance-critical code paths. - pub fn from_parts_unchecked(peaks: MmrPeaks, nodes: NodeMap, tracked_leaves: BTreeSet) -> Self { + pub fn from_parts_unchecked( + peaks: MmrPeaks, + nodes: NodeMap, + tracked_leaves: BTreeSet, + ) -> Self { let forest = peaks.forest(); let peaks = peaks.into(); @@ -1247,8 +1231,7 @@ mod tests { // Invalid case: tracked leaf has no value in nodes let mut tracked_no_value = BTreeSet::new(); tracked_no_value.insert(0); - let result = - PartialMmr::from_parts(peaks.clone(), BTreeMap::new(), tracked_no_value); + let result = PartialMmr::from_parts(peaks.clone(), BTreeMap::new(), tracked_no_value); assert!(result.is_err()); // Valid case: tracked leaf with its value in nodes @@ -1257,8 +1240,7 @@ mod tests { nodes_with_leaf.insert(leaf_idx, int_to_node(0)); let mut tracked_valid = BTreeSet::new(); tracked_valid.insert(0); - let result = - PartialMmr::from_parts(peaks.clone(), nodes_with_leaf, tracked_valid); + let result = PartialMmr::from_parts(peaks.clone(), nodes_with_leaf, tracked_valid); assert!(result.is_ok()); // Invalid case: node index out of bounds (leaf index) @@ -1273,24 +1255,38 @@ mod tests { // Create an InOrderIndex with value 0 via deserialization let zero_idx = InOrderIndex::read_from_bytes(&0usize.to_bytes()).unwrap(); nodes_with_zero.insert(zero_idx, int_to_node(0)); - let result = - PartialMmr::from_parts(peaks.clone(), nodes_with_zero, BTreeSet::new()); + let result = PartialMmr::from_parts(peaks.clone(), nodes_with_zero, BTreeSet::new()); assert!(result.is_err()); // Invalid case: large even index (internal node) beyond forest bounds let mut nodes_with_large_even = BTreeMap::new(); let large_even_idx = InOrderIndex::read_from_bytes(&1000usize.to_bytes()).unwrap(); nodes_with_large_even.insert(large_even_idx, int_to_node(0)); - let result = - PartialMmr::from_parts(peaks.clone(), nodes_with_large_even, BTreeSet::new()); + let result = PartialMmr::from_parts(peaks.clone(), nodes_with_large_even, BTreeSet::new()); assert!(result.is_err()); + // Invalid case: separator index between trees + // For 7 leaves (0b111 = 4+2+1), index 8 is a separator between the first tree (1-7) + // and the second tree (9-11). Similarly, index 12 is a separator between the second + // tree and the third tree (13). + let mut nodes_with_separator = BTreeMap::new(); + let separator_idx = InOrderIndex::read_from_bytes(&8usize.to_bytes()).unwrap(); + nodes_with_separator.insert(separator_idx, int_to_node(0)); + let result = PartialMmr::from_parts(peaks.clone(), nodes_with_separator, BTreeSet::new()); + assert!(result.is_err(), "separator index 8 should be rejected"); + + let mut nodes_with_separator_12 = BTreeMap::new(); + let separator_idx_12 = InOrderIndex::read_from_bytes(&12usize.to_bytes()).unwrap(); + nodes_with_separator_12.insert(separator_idx_12, int_to_node(0)); + let result = + PartialMmr::from_parts(peaks.clone(), nodes_with_separator_12, BTreeSet::new()); + assert!(result.is_err(), "separator index 12 should be rejected"); + // Invalid case: nodes with empty forest let empty_peaks = MmrPeaks::new(Forest::empty(), vec![]).unwrap(); let mut nodes_with_empty_forest = BTreeMap::new(); nodes_with_empty_forest.insert(InOrderIndex::from_leaf_pos(0), int_to_node(0)); - let result = - PartialMmr::from_parts(empty_peaks, nodes_with_empty_forest, BTreeSet::new()); + let result = PartialMmr::from_parts(empty_peaks, nodes_with_empty_forest, BTreeSet::new()); assert!(result.is_err()); } diff --git a/miden-crypto/src/merkle/mmr/tests.rs b/miden-crypto/src/merkle/mmr/tests.rs index aeabed8d9..515b9ce4b 100644 --- a/miden-crypto/src/merkle/mmr/tests.rs +++ b/miden-crypto/src/merkle/mmr/tests.rs @@ -199,6 +199,93 @@ fn test_forest_to_rightmost_index() { assert_eq!(Forest::new(0b1111).rightmost_in_order_index(), idx(29)); } +#[test] +fn test_is_valid_in_order_index() { + fn idx(pos: usize) -> InOrderIndex { + InOrderIndex::new(pos.try_into().unwrap()) + } + + // Empty forest has no valid indices + let empty = Forest::empty(); + assert!(!empty.is_valid_in_order_index(&idx(1))); + + // Single tree forests (power of 2 leaves) have no separators + // Forest with 1 leaf: valid indices are just 1 + let forest_1 = Forest::new(0b0001); + assert!(!forest_1.is_valid_in_order_index(&idx(2)), "index 0 is invalid"); + assert!(forest_1.is_valid_in_order_index(&idx(1))); + assert!(!forest_1.is_valid_in_order_index(&idx(2)), "beyond bounds"); + + // Forest with 2 leaves: valid indices are 1, 2, 3 + let forest_2 = Forest::new(0b0010); + assert!(forest_2.is_valid_in_order_index(&idx(1))); + assert!(forest_2.is_valid_in_order_index(&idx(2))); + assert!(forest_2.is_valid_in_order_index(&idx(3))); + assert!(!forest_2.is_valid_in_order_index(&idx(4)), "beyond bounds"); + + // Forest with 4 leaves: valid indices are 1-7 + let forest_4 = Forest::new(0b0100); + for i in 1..=7 { + assert!(forest_4.is_valid_in_order_index(&idx(i)), "index {} should be valid", i); + } + assert!(!forest_4.is_valid_in_order_index(&idx(8)), "beyond bounds"); + + // Multi-tree forest: 7 leaves (0b111 = 4 + 2 + 1) + // Tree 1 (4 leaves): indices 1-7 + // Separator: index 8 + // Tree 2 (2 leaves): indices 9-11 + // Separator: index 12 + // Tree 3 (1 leaf): index 13 + let forest_7 = Forest::new(0b0111); + + // Valid indices in first tree (4 leaves, 7 nodes) + for i in 1..=7 { + assert!( + forest_7.is_valid_in_order_index(&idx(i)), + "index {} should be valid in first tree", + i + ); + } + + // Separator between first and second tree + assert!(!forest_7.is_valid_in_order_index(&idx(8)), "index 8 is a separator"); + + // Valid indices in second tree (2 leaves, 3 nodes) + for i in 9..=11 { + assert!( + forest_7.is_valid_in_order_index(&idx(i)), + "index {} should be valid in second tree", + i + ); + } + + // Separator between second and third tree + assert!(!forest_7.is_valid_in_order_index(&idx(12)), "index 12 is a separator"); + + // Valid index in third tree (1 leaf) + assert!( + forest_7.is_valid_in_order_index(&idx(13)), + "index 13 should be valid in third tree" + ); + + // Beyond bounds + assert!(!forest_7.is_valid_in_order_index(&idx(14)), "index 14 is beyond bounds"); + + // Another multi-tree example: 6 leaves (0b110 = 4 + 2) + // Tree 1 (4 leaves): indices 1-7 + // Separator: index 8 + // Tree 2 (2 leaves): indices 9-11 + let forest_6 = Forest::new(0b0110); + for i in 1..=7 { + assert!(forest_6.is_valid_in_order_index(&idx(i)), "index {} should be valid", i); + } + assert!(!forest_6.is_valid_in_order_index(&idx(8)), "index 8 is a separator"); + for i in 9..=11 { + assert!(forest_6.is_valid_in_order_index(&idx(i)), "index {} should be valid", i); + } + assert!(!forest_6.is_valid_in_order_index(&idx(12)), "index 12 is beyond bounds"); +} + #[test] fn test_bit_position_iterator() { assert_eq!(TreeSizeIterator::new(Forest::empty()).count(), 0);