lsm_tree/
util.rs

1// Copyright (c) 2024-present, fjall-rs
2// This source code is licensed under both the Apache 2.0 and MIT License
3// (found in the LICENSE-* files in the repository)
4
5use crate::range::prefix_upper_range;
6use crate::UserKey;
7use std::ops::RangeBounds;
8
9pub use crate::range::prefix_to_range;
10
11/// Helper function to create a prefixed range.
12///
13/// Made for phil.
14///
15/// # Panics
16///
17/// Panics if the prefix is empty.
18pub fn prefixed_range<P: AsRef<[u8]>, K: AsRef<[u8]>, R: RangeBounds<K>>(
19    prefix: P,
20    range: R,
21) -> impl RangeBounds<UserKey> {
22    use std::ops::Bound::{Excluded, Included, Unbounded};
23
24    let prefix = prefix.as_ref();
25
26    assert!(!prefix.is_empty(), "prefix may not be empty");
27
28    match (range.start_bound(), range.end_bound()) {
29        (Unbounded, Unbounded) => prefix_to_range(prefix),
30        (lower, Unbounded) => {
31            let lower = lower.map(|k| UserKey::fused(prefix, k.as_ref()));
32            let upper = prefix_upper_range(prefix);
33            (lower, upper)
34        }
35        (Unbounded, upper) => {
36            let upper = match upper {
37                Included(k) => Included(UserKey::fused(prefix, k.as_ref())),
38                Excluded(k) => Excluded(UserKey::fused(prefix, k.as_ref())),
39                Unbounded => unreachable!(),
40            };
41
42            (Included(prefix.into()), upper)
43        }
44        (lower, upper) => {
45            let lower = match lower {
46                Included(k) => Included(UserKey::fused(prefix, k.as_ref())),
47                Excluded(k) => Excluded(UserKey::fused(prefix, k.as_ref())),
48                Unbounded => unreachable!(),
49            };
50
51            let upper = match upper {
52                Included(k) => Included(UserKey::fused(prefix, k.as_ref())),
53                Excluded(k) => Excluded(UserKey::fused(prefix, k.as_ref())),
54                Unbounded => unreachable!(),
55            };
56
57            (lower, upper)
58        }
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::prefixed_range;
65    use crate::UserKey;
66    use std::ops::Bound::{Excluded, Included};
67    use std::ops::RangeBounds;
68    use test_log::test;
69
70    #[test]
71    fn prefixed_range_1() {
72        let prefix = "abc";
73        let min = 5u8.to_be_bytes();
74        let max = 9u8.to_be_bytes();
75
76        let range = prefixed_range(prefix, min..=max);
77
78        assert_eq!(
79            range.start_bound(),
80            Included(&UserKey::new(&[b'a', b'b', b'c', 5]))
81        );
82        assert_eq!(
83            range.end_bound(),
84            Included(&UserKey::new(&[b'a', b'b', b'c', 9]))
85        );
86    }
87
88    #[test]
89    fn prefixed_range_2() {
90        let prefix = "abc";
91        let min = 5u8.to_be_bytes();
92        let max = 9u8.to_be_bytes();
93
94        let range = prefixed_range(prefix, min..max);
95
96        assert_eq!(
97            range.start_bound(),
98            Included(&UserKey::new(&[b'a', b'b', b'c', 5]))
99        );
100        assert_eq!(
101            range.end_bound(),
102            Excluded(&UserKey::new(&[b'a', b'b', b'c', 9]))
103        );
104    }
105
106    #[test]
107    fn prefixed_range_3() {
108        let prefix = "abc";
109        let min = 5u8.to_be_bytes();
110
111        let range = prefixed_range(prefix, min..);
112
113        assert_eq!(
114            range.start_bound(),
115            Included(&UserKey::new(&[b'a', b'b', b'c', 5]))
116        );
117        assert_eq!(range.end_bound(), Excluded(&UserKey::new(b"abd")));
118    }
119
120    #[test]
121    fn prefixed_range_4() {
122        let prefix = "abc";
123        let max = 9u8.to_be_bytes();
124
125        let range = prefixed_range(prefix, ..max);
126
127        assert_eq!(range.start_bound(), Included(&UserKey::new(b"abc")));
128        assert_eq!(
129            range.end_bound(),
130            Excluded(&UserKey::new(&[b'a', b'b', b'c', 9]))
131        );
132    }
133
134    #[test]
135    fn prefixed_range_5() {
136        let prefix = "abc";
137        let max = u8::MAX.to_be_bytes();
138
139        let range = prefixed_range(prefix, ..=max);
140
141        assert_eq!(range.start_bound(), Included(&UserKey::new(b"abc")));
142        assert_eq!(
143            range.end_bound(),
144            Included(&UserKey::new(&[b'a', b'b', b'c', u8::MAX]))
145        );
146    }
147
148    #[test]
149    fn prefixed_range_6() {
150        let prefix = "abc";
151        let max = u8::MAX.to_be_bytes();
152
153        let range = prefixed_range(prefix, ..max);
154
155        assert_eq!(range.start_bound(), Included(&UserKey::new(b"abc")));
156        assert_eq!(
157            range.end_bound(),
158            Excluded(&UserKey::new(&[b'a', b'b', b'c', u8::MAX]))
159        );
160    }
161}