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