Fix range search bugs

fuzz/Cargo.toml

@@ -29,3 +29,4 @@ doc = false
 
 [profile.release]
 rustflags = ["-C", "target-cpu=native", "-Z", "tune-cpu=native"]
+debug-assertions = true

fuzz/fuzz_targets/tree_map_api.rs

@@ -1,17 +1,19 @@
 #![no_main]
 #![feature(btree_entry_insert)]
 
+use std::{
+    collections::{hash_map::RandomState, BTreeMap},
+    hash::BuildHasher,
+    mem,
+    ops::Bound,
+};
+
 use blart::{
     map::{PrefixEntry, PrefixOccupied},
     visitor::WellFormedChecker,
     TreeMap,
 };
 use libfuzzer_sys::arbitrary::{self, Arbitrary};
-use std::{
-    collections::{hash_map::RandomState, BTreeMap},
-    hash::BuildHasher,
-    mem,
-};
 
 #[derive(Arbitrary, Debug)]
 enum EntryAction {
@@ -45,6 +47,44 @@ enum RetainKind {
     Half,
 }
 
+#[derive(Arbitrary, Debug)]
+enum RangeBound {
+    Unbounded,
+    Included(Box<[u8]>),
+    Excluded(Box<[u8]>),
+}
+
+impl RangeBound {
+    fn as_bound(&self) -> Bound<&[u8]> {
+        match self {
+            RangeBound::Unbounded => Bound::Unbounded,
+            RangeBound::Included(k) => Bound::Included(k.as_ref()),
+            RangeBound::Excluded(k) => Bound::Excluded(k.as_ref()),
+        }
+    }
+
+    fn key_bytes(&self) -> Option<&[u8]> {
+        match self {
+            RangeBound::Unbounded => None,
+            RangeBound::Included(k) | RangeBound::Excluded(k) => Some(k.as_ref()),
+        }
+    }
+}
+
+fn range_is_valid(start: &RangeBound, end: &RangeBound) -> bool {
+    match (start.key_bytes(), end.key_bytes()) {
+        (None, _) | (_, None) => true,
+        (Some(s), Some(e)) => {
+            s < e
+                || (s == e
+                    && matches!(
+                        (start, end),
+                        (RangeBound::Included(_), RangeBound::Included(_))
+                    ))
+        },
+    }
+}
+
 #[derive(Arbitrary, Debug)]
 enum Action {
     Clear,
@@ -69,6 +109,7 @@ enum Action {
     Fuzzy(Box<[u8]>),
     Prefix(Box<[u8]>),
     IntoIter { take_front: usize, take_back: usize },
+    Range(RangeBound, RangeBound),
     Retain(RetainKind),
     SplitOff(Box<[u8]>),
 }
@@ -258,13 +299,31 @@ libfuzzer_sys::fuzz_target!(|actions: Vec<Action>| {
                 });
             },
             Action::Fuzzy(key) => {
+                // TODO: Provide an oracle implementation for fuzzy search (hard)
                 let v: Vec<_> = tree.fuzzy(&key, 3).collect();
                 std::hint::black_box(v);
             },
             Action::Prefix(key) => {
+                // TODO: Provide an oracle implementation for prefix search (easy)
                 let v: Vec<_> = tree.prefix(&key).collect();
                 std::hint::black_box(v);
             },
+            Action::Range(start, end) => {
+                if range_is_valid(&start, &end) {
+                    let bounds = (start.as_bound(), end.as_bound());
+                    let tree_result: Vec<_> = tree.range::<[u8], _>(bounds).collect();
+                    let oracle_result: Vec<_> = oracle.range::<[u8], _>(bounds).collect();
+                    assert_eq!(
+                        tree_result.len(),
+                        oracle_result.len(),
+                        "range result lengths differ for bounds {bounds:?}"
+                    );
+                    for ((tk, tv), (ok, ov)) in tree_result.iter().zip(oracle_result.iter()) {
+                        assert_eq!(tk.as_ref() as &[u8], ok.as_ref() as &[u8]);
+                        assert_eq!(tv, ov);
+                    }
+                }
+            },
             Action::IntoIter {
                 take_front,
                 take_back,

src/raw/operations/lookup.rs

@@ -1,5 +1,3 @@
-use core::cmp::Ordering;
-
 use crate::{
     raw::{
         match_concrete_node_ptr, AttemptOptimisticPrefixMatch, ConcreteNodePtr, InnerNode,
@@ -101,40 +99,6 @@ impl PrefixMatchBehavior {
         result
     }
 
-    /// This function will compare the key bytes against the key in the given
-    /// leaf node.
-    ///
-    /// Specifically:
-    ///  - If the current behavior is "optimistic", then the entire leaf key
-    ///    will be compared against the given key bytes
-    ///  - If the current behavior is "pessimistic", then only the key bytes
-    ///    that were not used during the lookup process will be compared against
-    ///    the corresponding leaf key bytes.
-    ///
-    /// This is a minor optimization to reduce the amount of work needed
-    /// confirming that a lookup found the right leaf node.
-    pub fn compare_leaf_key<K: AsBytes, V, const PREFIX_LEN: usize>(
-        self,
-        leaf: &LeafNode<K, V, PREFIX_LEN>,
-        key_bytes: &[u8],
-        current_depth: usize,
-    ) -> Ordering {
-        match self {
-            PrefixMatchBehavior::Pessimistic => {
-                let leaf_key_bytes = leaf.key_ref().as_bytes();
-                let current_depth = current_depth.min(leaf_key_bytes.len()).min(key_bytes.len());
-                // PANIC SAFETY: Since we limit `current_depth` to be the minimum of the lengths
-                // and the current depth we will at most get an empty slice, it
-                // should panic. I ran a small test to make sure that `&[1][1..] == &[][..]` and
-                // does not panic.
-                let leaf_key_bytes = &leaf_key_bytes[current_depth..];
-                let key_bytes = &key_bytes[current_depth..];
-                leaf_key_bytes.cmp(key_bytes)
-            },
-            PrefixMatchBehavior::Optimistic => leaf.key_ref().as_bytes().cmp(key_bytes),
-        }
-    }
-
     /// This function will test the key bytes against the key in the given leaf
     /// node.
     ///

src/raw/representation.rs

@@ -1,6 +1,7 @@
 //! Trie node representation
 
 use core::{
+    cmp::Ordering,
     fmt,
     iter::FusedIterator,
     ops::{Bound, RangeBounds, RangeInclusive},
@@ -138,6 +139,8 @@ pub struct ExplicitMismatch<K, V, const PREFIX_LEN: usize> {
     pub prefix_byte: u8,
     /// Pointer to the leaf if the prefix was reconstructed
     pub leaf_ptr: OptionalLeafPtr<K, V, PREFIX_LEN>,
+    /// Comparison between the full prefix and the relevant key segment
+    pub node_prefix_comparison_to_search_key_segment: Ordering,
 }
 
 impl<K, V, const PREFIX_LEN: usize> Clone for ExplicitMismatch<K, V, PREFIX_LEN> {
@@ -217,7 +220,7 @@ pub unsafe trait InnerNodeCommon<K, V, const PREFIX_LEN: usize>: Sized {
     /// `prefix_len` as the node prefix length.
     ///
     /// This is done because when a prefix mismatch happens
-    /// the length of the mismatch can be grater or equal to
+    /// the length of the mismatch can be greater or equal to
     /// prefix size, since we search for the first child of the
     /// node to recreate the prefix, that's why we don't use
     /// `prefix.len()` as the node prefix length
@@ -381,18 +384,28 @@ pub unsafe trait InnerNodeCommon<K, V, const PREFIX_LEN: usize>: Sized {
         }
 
         let (prefix, leaf_ptr) = self.read_full_prefix(current_depth);
-        let key = &key[current_depth..];
 
         let matched_bytes = prefix
             .iter()
-            .zip(key)
+            .zip(&key[current_depth..])
             .take_while(|(a, b)| **a == **b)
             .count();
         if matched_bytes < prefix.len() {
+            let upper_bound = (current_depth + prefix.len()).min(key.len());
+            let key_segment = &key[current_depth..upper_bound];
+            let node_prefix_comparison_to_search_key_segment = 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"
+            );
+
             Err(ExplicitMismatch {
                 matched_bytes,
                 prefix_byte: prefix[matched_bytes],
                 leaf_ptr,
+                node_prefix_comparison_to_search_key_segment,
             })
         } else {
             Ok(PrefixMatch { matched_bytes })
@@ -425,20 +438,18 @@ pub unsafe trait InnerNodeCommon<K, V, const PREFIX_LEN: usize>: Sized {
             let leaf = unsafe { leaf_ptr.as_ref() };
             let leaf = leaf.key_ref().as_bytes();
 
-            unsafe {
-                // SAFETY: Since we are iterating the key and prefixes, we
-                // expect that the depth never exceeds the key len.
-                // Because if this happens we ran out of bytes in the key to match
-                // and the whole process should be already finished
-                core::hint::assert_unchecked(current_depth <= leaf.len());
-
-                // SAFETY: By the construction of the prefix we know that this is inbounds
-                // since the prefix len guarantees it to us
-                core::hint::assert_unchecked(current_depth + len <= leaf.len());
-
-                // SAFETY: This can't overflow since len comes from a u32
-                core::hint::assert_unchecked(current_depth <= current_depth + len);
-            }
+            assert!(
+                current_depth <= leaf.len(),
+                "current_depth [{current_depth}] must not exceed leaf key length [{}]; leaf must \
+                 be a descendant of this node",
+                leaf.len()
+            );
+            assert!(
+                current_depth + len <= leaf.len(),
+                "current_depth [{current_depth}] + prefix_len [{len}] must not exceed leaf key \
+                 length [{}]; leaf must be a descendant of this node",
+                leaf.len()
+            );
             let leaf = &leaf[current_depth..(current_depth + len)];
             (leaf, Some(leaf_ptr))
         }

src/raw/representation/header.rs

@@ -103,16 +103,12 @@ impl<const PREFIX_LEN: usize> Header<PREFIX_LEN> {
         let begin = len;
         let end = begin + self.capped_prefix_len();
 
-        unsafe {
-            // SAFETY: This function is called when mismatch happened and
-            // we used the node to match the number of bytes,
-            // by this we know that len < prefix len, but since we + 1,
-            // to skip the key byte we have that len <= prefix len
-            core::hint::assert_unchecked(end <= self.prefix.len());
-
-            // SAFETY: This is by construction end = begin + len
-            core::hint::assert_unchecked(begin <= end);
-        }
+        assert!(
+            end <= self.prefix.len(),
+            "prefix copy range end [{end}] must not exceed prefix array length [{}]; ltrim_by \
+             must only be called when prefix_len <= PREFIX_LEN",
+            self.prefix.len()
+        );
         self.prefix.copy_within(begin..end, 0);
     }
 
@@ -179,16 +175,12 @@ impl<const PREFIX_LEN: usize> Header<PREFIX_LEN> {
         let end = begin + self.capped_prefix_len();
         let len = end - begin;
 
-        unsafe {
-            // SAFETY: This function is called a mismatch happened and
-            // we used the leaf to match the number of matching bytes,
-            // by this we know that len < prefix len, but since we + 1,
-            // to skip the key byte we have that len <= prefix len
-            core::hint::assert_unchecked(end <= leaf_key.len());
-
-            // SAFETY: This is by construction end = begin + len
-            core::hint::assert_unchecked(begin <= end);
-        }
+        assert!(
+            end <= leaf_key.len(),
+            "leaf key slice end [{end}] must not exceed leaf key length [{}]; leaf must be a \
+             descendant of this node so its key covers depth+prefix_len bytes",
+            leaf_key.len()
+        );
 
         let leaf_key = &leaf_key[begin..end];
         self.prefix[..len].copy_from_slice(leaf_key)