Fix regression in prefix iterator

declanvk · declanvk · commit eb60a11a324d · 2026-04-10T09:17:17.000-07:00
# Description
The issue was that in the case that prefix search terminated in an inner
node due, it was confusing the cases of (1) running out of key bytes or
(2) an actual mismatch between bytes of the inner node prefix and the key
bytes. In case (1) we know that all the nodes of the subtree should be
returned by the iterator, since the key is a prefix of all those leaves.
In case (2), no nodes of the subtree should be returned, since there was
a mismatch all nodes will not be prefixed by the key.

The fix was to make the `InnerNodeSearchResultReason` have another
variant which reports when the search ran out of bytes.

# Testing
Regression tests were added in a previous commit and labeled
`should_panic`, this commit just removed that label and all tests still
passed.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -20,6 +20,12 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
  - `TreeMap::{get_prefix, get_prefix_key_value, get_prefix_mut, force_insert}` were renamed to `TreeMap::{prefix_get, prefix_get_key_value, prefix_get_mut, prefix_insert}`. This makes it clear that they're all centered around the theme of manipulating or querying key prefixes.
  - Update MSRV to 1.85
 
+### Fixed
+
+ - `TreeMap::prefix` iterator was broken (see [issue 66]) because a refactor stopped distinguishing between different cases of inner nodes matching or not matching the prefix.
+
+[issue 66]: https://github.com/declanvk/blart/issues/66
+
 ## [0.4.0] - 2025-07-10
 
 ### Added
diff --git a/src/collections/map/iterators/prefix.rs b/src/collections/map/iterators/prefix.rs
@@ -6,7 +6,7 @@ use super::{
 };
 use crate::{
     allocator::{Allocator, Global},
-    map::DEFAULT_PREFIX_LEN,
+    map::{NonEmptyTree, DEFAULT_PREFIX_LEN},
     raw::{maximum_unchecked, minimum_unchecked, RawIterator},
     AsBytes, TreeMap,
 };
@@ -40,48 +40,11 @@ macro_rules! implement_prefix_iter {
                     };
                 };
 
-                // SAFETY: Since we have a (shared or mutable) reference to the original
-                // TreeMap, we know there will be no concurrent mutation
-                let search_result = unsafe { find_terminating_node(tree_state.root, prefix) };
-                let (start, end) = match search_result {
-                    TerminatingNodeSearchResult::Leaf { leaf_ptr, .. } => {
-                        // SAFETY: Its safe to create a shared reference to the leaf since we hold
-                        // either a shared or mutable reference to the owning TreeMap, which prevents
-                        // other concurrent mutable references.
-                        let leaf = unsafe { leaf_ptr.as_ref() };
-
-                        // Only include the item in the iterator if the prefix actually matches
-                        if leaf.key_ref().as_bytes().starts_with(prefix) {
-                            (leaf_ptr, leaf_ptr)
-                        } else {
-                            return Self {
-                                _tree: tree,
-                                inner: RawIterator::empty(),
-                            };
-                        }
-                    },
-                    TerminatingNodeSearchResult::InnerNode(InnerNodeSearchResult { node_ptr, reason, .. }) => {
-                        if matches!(reason, InnerNodeSearchResultReason::MissingChild) {
-                            // if the child is missing, then there is nothing to be the prefix of
-                            return Self {
-                                _tree: tree,
-                                inner: RawIterator::empty(),
-                            };
-                        }
-
-                        // SAFETY: Its safe to create a shared reference to the leaf since we hold
-                        // either a shared or mutable reference to the owning TreeMap, which prevents
-                        // other concurrent mutable references.
-                        unsafe { (minimum_unchecked(node_ptr), maximum_unchecked(node_ptr)) }
-                    },
-                };
+                let inner = prefix_iter_constructor(tree_state, prefix);
 
                 Self {
                     _tree: tree,
-                    // SAFETY: `start` is guaranteed to be less than or equal to `end` in the iteration
-                    // order because of the check we do on the bytes of the resolved leaf pointers, just
-                    // above this line
-                    inner: unsafe { RawIterator::new(start, end) },
+                    inner,
                 }
             }
         }
@@ -120,6 +83,53 @@ macro_rules! implement_prefix_iter {
     };
 }
 
+fn prefix_iter_constructor<K: AsBytes, V, const PREFIX_LEN: usize>(
+    tree_state: &NonEmptyTree<K, V, PREFIX_LEN>,
+    prefix: &[u8],
+) -> RawIterator<K, V, PREFIX_LEN> {
+    // SAFETY: Since we have a (shared or mutable) reference to the original
+    // TreeMap, we know there will be no concurrent mutation
+    let search_result = unsafe { find_terminating_node(tree_state.root, prefix) };
+    let (start, end) = match search_result {
+        TerminatingNodeSearchResult::Leaf { leaf_ptr, .. } => {
+            // SAFETY: Its safe to create a shared reference to the leaf since we hold
+            // either a shared or mutable reference to the owning TreeMap, which prevents
+            // other concurrent mutable references.
+            let leaf = unsafe { leaf_ptr.as_ref() };
+
+            // Only include the item in the iterator if the prefix actually matches
+            if leaf.key_ref().as_bytes().starts_with(prefix) {
+                (leaf_ptr, leaf_ptr)
+            } else {
+                return RawIterator::empty();
+            }
+        },
+        TerminatingNodeSearchResult::InnerNode(InnerNodeSearchResult {
+            node_ptr, reason, ..
+        }) => {
+            match reason {
+                InnerNodeSearchResultReason::PrefixMismatch
+                | InnerNodeSearchResultReason::MissingChild => {
+                    // if the child is missing OR there is a prefix mismatch, then there is nothing
+                    // to be the prefix of
+                    return RawIterator::empty();
+                },
+                InnerNodeSearchResultReason::InsufficientBytes => {
+                    // SAFETY: Its safe to create a shared reference to the leaf since we hold
+                    // either a shared or mutable reference to the owning TreeMap, which prevents
+                    // other concurrent mutable references.
+                    unsafe { (minimum_unchecked(node_ptr), maximum_unchecked(node_ptr)) }
+                },
+            }
+        },
+    };
+
+    // SAFETY: `start` is guaranteed to be less than or equal to `end` in the
+    // iteration order because of the check we do on the bytes of the resolved
+    // leaf pointers, just above this line
+    unsafe { RawIterator::new(start, end) }
+}
+
 implement_prefix_iter!(
     /// An iterator over a range of entries that all have the same key prefix in a [`TreeMap`].
     ///
@@ -438,8 +448,7 @@ mod tests {
     /// the prefix. In blart 0.3.0 and 0.4.0 it incorrectly returns both
     /// keys.
     #[test]
-    #[should_panic(expected = "`left == right` failed")]
-    fn test_prefix_iterator_no_false_matches() {
+    fn prefix_iterator_no_false_matches() {
         // Construct a map with the following keys
         // [0, 0]
         // [0, 1]
diff --git a/src/collections/map/iterators/range.rs b/src/collections/map/iterators/range.rs
@@ -14,7 +14,7 @@ use crate::{
     raw::{
         match_concrete_node_ptr, maximum_unchecked, minimum_unchecked,
         AttemptOptimisticPrefixMatch, ConcreteNodePtr, InnerNode, LeafNode, NodePtr, OpaqueNodePtr,
-        PrefixMatchBehavior, RawIterator,
+        PessimisticMismatch, PrefixMatchBehavior, RawIterator,
     },
     AsBytes, TreeMap,
 };
@@ -55,9 +55,11 @@ impl<K, V, const PREFIX_LEN: usize> fmt::Debug for InnerNodeSearchResult<K, V, P
 #[derive(Debug)]
 pub(crate) enum InnerNodeSearchResultReason {
     /// This variant means the search terminated in a mismatched prefix of an
-    /// inner node OR the prefix matched, but there were insufficient key bytes
-    /// to lookup a child of the inner node.
-    PrefixMismatchOrInsufficientBytes,
+    /// inner node.
+    PrefixMismatch,
+    /// This variant means the search terminated because the key was too short,
+    /// not because of a mismatch with prefix or child key byte.
+    InsufficientBytes,
     /// This variant means the search terminated in an inner node, when there
     /// was no corresponding child for a key byte.
     MissingChild,
@@ -136,17 +138,34 @@ pub(crate) unsafe fn find_terminating_node<K: AsBytes, V, const PREFIX_LEN: usiz
         let match_prefix =
             prefix_match_behavior.match_prefix(inner_node, &key_bytes[*current_depth..]);
         match match_prefix {
-            Err(_) => {
+            Err(PessimisticMismatch { matched_bytes, .. }) => {
                 let (full_prefix, _) = inner_node.read_full_prefix(*current_depth);
                 let upper_bound = (*current_depth + full_prefix.len()).min(key_bytes.len());
                 let key_segment = &key_bytes[(*current_depth)..upper_bound];
 
                 let node_prefix_comparison_to_search_key_segment = full_prefix.cmp(key_segment);
 
+                debug_assert_ne!(
+                    node_prefix_comparison_to_search_key_segment,
+                    Ordering::Equal,
+                    "if there was a mismatch, the prefix must not be equal"
+                );
+
+                // We report a prefix mismatch if the prefix is longer that the remaining
+                // portion of the key bytes. However, in the context of the
+                // `InnerNodeSearchResultReason` running out of bytes is really a different
+                // thing than a mismatch. We should only report mismatch as the reason if there
+                // was an actual byte difference.
+                let reason = if matched_bytes == key_bytes[*current_depth..].len() {
+                    InnerNodeSearchResultReason::InsufficientBytes
+                } else {
+                    InnerNodeSearchResultReason::PrefixMismatch
+                };
+
                 ControlFlow::Break(InnerNodeSearchResult {
                     node_ptr: inner_ptr.to_opaque(),
                     current_depth: *current_depth,
-                    reason: InnerNodeSearchResultReason::PrefixMismatchOrInsufficientBytes,
+                    reason,
                     node_prefix_comparison_to_search_key_segment,
                 })
             },
@@ -163,7 +182,7 @@ pub(crate) unsafe fn find_terminating_node<K: AsBytes, V, const PREFIX_LEN: usiz
                     return ControlFlow::Break(InnerNodeSearchResult {
                         node_ptr: inner_ptr.to_opaque(),
                         current_depth: *current_depth,
-                        reason: InnerNodeSearchResultReason::PrefixMismatchOrInsufficientBytes,
+                        reason: InnerNodeSearchResultReason::InsufficientBytes,
                         node_prefix_comparison_to_search_key_segment: Ordering::Equal,
                     });
                 };
@@ -308,7 +327,8 @@ pub(crate) unsafe fn find_leaf_pointer_for_bound<K: AsBytes, V, const PREFIX_LEN
                     reason,
                     node_prefix_comparison_to_search_key_segment,
                 }) => match reason {
-                    InnerNodeSearchResultReason::PrefixMismatchOrInsufficientBytes => {
+                    InnerNodeSearchResultReason::PrefixMismatch
+                    | InnerNodeSearchResultReason::InsufficientBytes => {
                         match (is_start, node_prefix_comparison_to_search_key_segment) {
                             (true, Ordering::Equal | Ordering::Greater) => unsafe {
                                 // SAFETY: Covered by function safety doc
