prefix_range/
lib.rs

1//! Code for computing an upper bound for BTreeMapRange for a string prefix.
2//!
3//! While this crate doesn't need std, it does need alloc.
4//!
5//! Taken from https://www.thecodedmessage.com/posts/prefix-ranges/ (by Jimmy Hartzell)
6//! Extended with documentation.
7//!
8//! Licensed under: MIT
9
10#![no_std]
11
12extern crate alloc;
13use alloc::collections::BTreeSet;
14use alloc::{format, string::String};
15
16/// Find the key to use as an upper bound for range query
17///
18/// ```
19/// # use std::collections::BTreeMap;
20/// # use std::ops::Bound;
21/// let mut data: BTreeMap<String, i32> = BTreeMap::new();
22/// data.insert("bbabb".to_owned(), 1);
23/// data.insert("bbbcc".to_owned(), 2);
24/// data.insert("bbbab".to_owned(), 3);
25/// data.insert("bbb".to_owned(), 4);
26/// data.insert("bbc".to_owned(), 5);
27///
28/// let lower = Bound::Included("bbb".to_owned());
29/// let upper = Bound::Excluded(prefix_range::upper_bound_from_prefix("bbb").unwrap());
30/// let in_prefix: Vec<_> = data.range((lower, upper)).map(|(_, v)| *v).collect();
31/// assert_eq!(in_prefix, vec![4, 3, 2]);
32/// ```
33pub fn upper_bound_from_prefix(prefix: &str) -> Option<String> {
34    for i in (0..prefix.len()).rev() {
35        if let Some(last_char_str) = prefix.get(i..) {
36            let rest_of_prefix = {
37                debug_assert!(prefix.is_char_boundary(i));
38                &prefix[0..i]
39            };
40
41            let last_char = last_char_str
42                .chars()
43                .next()
44                .expect("last_char_str will contain at least one char");
45            let Some(last_char_incr) = (last_char..=char::MAX).nth(1) else {
46                // Last character is highest possible code point.
47                // Go to second-to-last character instead.
48                continue;
49            };
50
51            let new_string = format!("{rest_of_prefix}{last_char_incr}");
52
53            return Some(new_string);
54        }
55    }
56
57    None
58}
59
60/// Create a new BTreeSet that contains all the keys with the given prefix
61pub fn prefixed_set(mut set: BTreeSet<String>, prefix: &str) -> BTreeSet<String> {
62    let mut set = set.split_off(prefix);
63
64    if let Some(not_in_prefix) = upper_bound_from_prefix(prefix) {
65        set.split_off(&not_in_prefix);
66    }
67
68    set
69}
70
71#[cfg(test)]
72mod tests {
73
74    use alloc::string::ToString;
75
76    use super::*;
77
78    #[test]
79    fn it_works() {
80        let set = {
81            let mut set = BTreeSet::new();
82            set.insert("Hi".to_string());
83            set.insert("Hey".to_string());
84            set.insert("Hello".to_string());
85            set.insert("heyyy".to_string());
86            set.insert("".to_string());
87            set.insert("H".to_string());
88            set
89        };
90        let set = prefixed_set(set, "H");
91        assert_eq!(set.len(), 4);
92        assert!(!set.contains("heyyy"));
93    }
94
95    #[test]
96    fn maxicode() {
97        let set = {
98            let mut set = BTreeSet::new();
99            set.insert("Hi".to_string());
100            set.insert("Hey".to_string());
101            set.insert("Hello".to_string());
102            set.insert("heyyy".to_string());
103            set.insert("H\u{10FFFF}eyyy".to_string());
104            set.insert("H\u{10FFFF}".to_string());
105            set.insert("I".to_string());
106            set.insert("".to_string());
107            set.insert("H".to_string());
108            set
109        };
110        let set = prefixed_set(set, "H\u{10FFFF}");
111        assert_eq!(set.len(), 2);
112        assert!(!set.contains("I"));
113    }
114}
prefix_range/lib.rs

prefix_range/
lib.rs
