1use crate::range::prefix_upper_range;
6use crate::UserKey;
7use std::ops::RangeBounds;
8
9pub use crate::range::prefix_to_range;
10
11pub 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}