1#![no_std]
6#![warn(clippy::cargo)]
7#![deny(clippy::cargo_common_metadata)]
8#![deny(rustdoc::broken_intra_doc_links)]
9#![deny(clippy::all)]
10#![deny(clippy::pedantic)]
11#![allow(clippy::missing_panics_doc)]
12#![cfg_attr(
13 not(test),
14 warn(
15 missing_debug_implementations,
16 trivial_casts,
17 trivial_numeric_casts,
18 unused_extern_crates,
19 unused_import_braces,
20 unused_qualifications,
21 unused_results
22 )
23)]
24#![cfg_attr(
25 test,
26 deny(
27 missing_debug_implementations,
28 trivial_casts,
29 trivial_numeric_casts,
30 unused_extern_crates,
31 unused_import_braces,
32 unused_qualifications,
33 unused_must_use,
34 unused_results
35 )
36)]
37#![cfg_attr(
38 test,
39 deny(
40 bad_style,
41 dead_code,
42 improper_ctypes,
43 non_shorthand_field_patterns,
44 no_mangle_generic_items,
45 overflowing_literals,
46 path_statements,
47 patterns_in_fns_without_body,
48 private_in_public,
49 unconditional_recursion,
50 unused,
51 unused_allocation,
52 unused_comparisons,
53 unused_parens,
54 while_true
55 )
56)]
57
58#[macro_use]
59pub extern crate alloc;
60
61use alloc::boxed::Box;
62use core::{
63 cmp::{Ord, Ordering},
64 ops::Range,
65};
66#[cfg(feature = "serde")]
67use serde::{Deserialize, Serialize};
68
69mod node;
70use node::Node;
71
72mod interval;
73pub use interval::Interval;
74
75mod iterators;
76pub use iterators::{Entry, EntryMut, IntervalTreeIterator, IntervalTreeIteratorMut};
77
78#[derive(Clone, Debug, Default)]
79#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
80pub struct IntervalTree<T: Ord + Clone, V> {
81 root: Option<Box<Node<T, V>>>,
82}
83
84impl<T: Ord + Clone, V> IntervalTree<T, V> {
85 #[must_use]
86 pub fn new() -> Self {
87 IntervalTree { root: None }
88 }
89
90 #[must_use]
91 pub fn is_empty(&self) -> bool {
92 self.root.is_none()
93 }
94
95 #[must_use]
96 pub fn size(&self) -> usize {
97 Node::size(&self.root)
98 }
99
100 #[must_use]
101 pub fn height(&self) -> i64 {
102 Node::height(&self.root)
103 }
104
105 #[must_use]
106 pub fn query<I: Into<Interval<T>>>(&self, interval: I) -> IntervalTreeIterator<'_, T, V> {
107 if let Some(ref n) = self.root {
108 IntervalTreeIterator {
109 nodes: vec![n],
110 interval: interval.into(),
111 }
112 } else {
113 let nodes = vec![];
114 IntervalTreeIterator {
115 nodes,
116 interval: interval.into(),
117 }
118 }
119 }
120
121 #[must_use]
122 pub fn query_mut<I: Into<Interval<T>>>(
123 &mut self,
124 interval: I,
125 ) -> IntervalTreeIteratorMut<'_, T, V> {
126 if let Some(ref mut n) = self.root {
127 IntervalTreeIteratorMut {
128 nodes: vec![n],
129 interval: interval.into(),
130 }
131 } else {
132 let nodes = vec![];
133 IntervalTreeIteratorMut {
134 nodes,
135 interval: interval.into(),
136 }
137 }
138 }
139
140 pub fn insert<I: Into<Interval<T>>>(&mut self, interval: I, value: V) {
141 let interval = interval.into();
142 let max = interval.end.clone();
143
144 self.root = Some(IntervalTree::insert_helper(
145 self.root.take(),
146 interval,
147 value,
148 max,
149 ));
150 }
151
152 #[allow(clippy::unnecessary_box_returns)]
153 fn insert_helper(
154 node: Option<Box<Node<T, V>>>,
155 interval: Interval<T>,
156 value: V,
157 max: T,
158 ) -> Box<Node<T, V>> {
159 if node.is_none() {
160 return Box::new(Node::new(interval, value, max, 0, 1));
161 }
162
163 let mut node_ref = node.unwrap();
164
165 match interval.cmp(&node_ref.interval) {
166 Ordering::Less => {
167 node_ref.left_child = Some(IntervalTree::insert_helper(
168 node_ref.left_child,
169 interval,
170 value,
171 max,
172 ));
173 }
174 Ordering::Greater => {
175 node_ref.right_child = Some(IntervalTree::insert_helper(
176 node_ref.right_child,
177 interval,
178 value,
179 max,
180 ));
181 }
182 Ordering::Equal => return node_ref,
183 }
184
185 node_ref.update_height();
186 node_ref.update_size();
187 node_ref.update_max();
188
189 IntervalTree::balance(node_ref)
190 }
191
192 #[allow(clippy::unnecessary_box_returns)]
193 fn balance(mut node: Box<Node<T, V>>) -> Box<Node<T, V>> {
194 if Node::balance_factor(&node) < -1 {
195 if Node::balance_factor(node.right_child.as_ref().unwrap()) > 0 {
196 node.right_child = Some(IntervalTree::rotate_right(node.right_child.unwrap()));
197 }
198 node = IntervalTree::rotate_left(node);
199 } else if Node::balance_factor(&node) > 1 {
200 if Node::balance_factor(node.left_child.as_ref().unwrap()) < 0 {
201 node.left_child = Some(IntervalTree::rotate_left(node.left_child.unwrap()));
202 }
203 node = IntervalTree::rotate_right(node);
204 }
205 node
206 }
207
208 #[allow(clippy::unnecessary_box_returns)]
209 fn rotate_right(mut node: Box<Node<T, V>>) -> Box<Node<T, V>> {
210 let mut y = node.left_child.unwrap();
211 node.left_child = y.right_child;
212 y.size = node.size;
213 node.update_height();
214 node.update_size();
215 node.update_max();
216
217 y.right_child = Some(node);
218 y.update_height();
219 y.update_max();
220
221 y
222 }
223
224 #[allow(clippy::unnecessary_box_returns)]
225 fn rotate_left(mut node: Box<Node<T, V>>) -> Box<Node<T, V>> {
226 let mut y = node.right_child.unwrap();
227 node.right_child = y.left_child;
228 y.size = node.size;
229
230 node.update_height();
231 node.update_size();
232 node.update_max();
233
234 y.left_child = Some(node);
235 y.update_height();
236 y.update_max();
237
238 y
239 }
240
241 pub fn delete<I: Into<Interval<T>>>(&mut self, interval: I) {
242 if !self.is_empty() {
243 let interval = interval.into();
244 self.root = IntervalTree::delete_helper(self.root.take(), &interval);
245 }
246 }
247
248 fn delete_helper(
249 node: Option<Box<Node<T, V>>>,
250 interval: &Interval<T>,
251 ) -> Option<Box<Node<T, V>>> {
252 match node {
253 None => None,
254 Some(mut node) => {
255 if *interval < node.interval {
256 node.left_child = IntervalTree::delete_helper(node.left_child.take(), interval);
257 } else if *interval > node.interval {
258 node.right_child =
259 IntervalTree::delete_helper(node.right_child.take(), interval);
260 } else if node.left_child.is_none() {
261 return node.right_child;
262 } else if node.right_child.is_none() {
263 return node.left_child;
264 } else {
265 let mut y = node;
266 node = IntervalTree::min(&mut y.right_child);
267 node.right_child = IntervalTree::delete_min_helper(y.right_child.unwrap());
268 node.left_child = y.left_child;
269 }
270
271 node.update_height();
272 node.update_size();
273 node.update_max();
274 Some(IntervalTree::balance(node))
275 }
276 }
277 }
278
279 #[allow(clippy::unnecessary_box_returns)]
280 fn min(node: &mut Option<Box<Node<T, V>>>) -> Box<Node<T, V>> {
281 match node {
282 Some(node) => {
283 if node.left_child.is_none() {
284 Box::new(Node::new(
285 node.interval.clone(),
286 node.value.take().unwrap(),
287 node.max.clone(),
288 0,
289 1,
290 ))
291 } else {
292 IntervalTree::min(&mut node.left_child)
293 }
294 }
295 None => panic!("Called min on None node"),
296 }
297 }
298
299 pub fn delete_min(&mut self) {
300 if !self.is_empty() {
301 self.root = IntervalTree::delete_min_helper(self.root.take().unwrap());
302 }
303 }
304
305 fn delete_min_helper(mut node: Box<Node<T, V>>) -> Option<Box<Node<T, V>>> {
306 if node.left_child.is_none() {
307 return node.right_child.take();
308 }
309
310 node.left_child = IntervalTree::delete_min_helper(node.left_child.unwrap());
311
312 node.update_height();
313 node.update_size();
314 node.update_max();
315
316 Some(IntervalTree::balance(node))
317 }
318
319 pub fn delete_max(&mut self) {
320 if !self.is_empty() {
321 self.root = IntervalTree::delete_max_helper(self.root.take().unwrap());
322 }
323 }
324
325 fn delete_max_helper(mut node: Box<Node<T, V>>) -> Option<Box<Node<T, V>>> {
326 if node.right_child.is_none() {
327 return node.left_child.take();
328 }
329
330 node.right_child = IntervalTree::delete_max_helper(node.right_child.unwrap());
331
332 node.update_height();
333 node.update_size();
334 node.update_max();
335
336 Some(IntervalTree::balance(node))
337 }
338
339 pub fn clear(&mut self) {
340 self.root = None;
341 }
342}
343
344impl<T: Ord + Clone, V> FromIterator<(Range<T>, V)> for IntervalTree<T, V> {
345 fn from_iter<I: IntoIterator<Item = (Range<T>, V)>>(iter: I) -> Self {
346 let mut ret = IntervalTree::new();
347 for (interval, value) in iter {
348 ret.insert(interval, value);
349 }
350 ret
351 }
352}
353
354#[cfg(test)]
355mod tests {
356 use alloc::vec::Vec;
357
358 use super::*;
359
360 #[test]
361 fn query_1() {
362 let mut tree = IntervalTree::<usize, bool>::new();
363 for i in 0..10 {
364 tree.insert((i * 10)..(i * 10 + 10), false);
365 }
366
367 let mut cnt = 0;
368 for _ in tree.query(0..10000) {
369 cnt += 1;
370 }
371 assert_eq!(cnt, 10);
372 }
373
374 #[test]
375 fn query_2() {
376 let mut tree = IntervalTree::<usize, bool>::new();
377 for i in 0..10 {
378 tree.insert((i * 10)..(i * 10 + 10), false);
379 }
380
381 let mut cnt = 0;
382 for _ in tree.query(0..30) {
383 cnt += 1;
384 }
385 assert_eq!(cnt, 3);
386 }
387
388 fn verify(tree: &IntervalTree<u32, u32>, i: u32, expected: &[u32]) {
389 let mut v1: Vec<_> = tree.query(i..=i).map(|x| *x.value).collect();
390 v1.sort();
391 let mut v2: Vec<_> = tree.query(i..(i + 1)).map(|x| *x.value).collect();
392 v2.sort();
393 assert_eq!(v1, expected);
394 assert_eq!(v2, expected);
395 }
396
397 #[test]
398 fn it_works() {
399 let tree: IntervalTree<u32, u32> = [
400 (0..3, 1),
401 (1..4, 2),
402 (2..5, 3),
403 (3..6, 4),
404 (4..7, 5),
405 (5..8, 6),
406 (4..5, 7),
407 (2..7, 8),
408 ]
409 .iter()
410 .cloned()
411 .collect();
412
413 verify(&tree, 0, &[1]);
414 verify(&tree, 1, &[1, 2]);
415 verify(&tree, 2, &[1, 2, 3, 8]);
416 verify(&tree, 3, &[2, 3, 4, 8]);
417 verify(&tree, 4, &[3, 4, 5, 7, 8]);
418 verify(&tree, 5, &[4, 5, 6, 8]);
419 verify(&tree, 6, &[5, 6, 8]);
420 verify(&tree, 7, &[6]);
421 verify(&tree, 8, &[]);
422 verify(&tree, 9, &[]);
423
424 assert_eq!(tree.query(1..1).next(), None);
425 assert_eq!(tree.query(1..=2).next().unwrap().interval.end, 4);
426 assert_eq!(tree.query(1..=2).next().unwrap().value, &2);
427 assert_eq!(tree.query(1..=5).next().unwrap().value, &4);
428 }
429
430 #[test]
431 fn empty() {
432 let tree: IntervalTree<u32, u32> = IntervalTree::new();
433 verify(&tree, 42, &[]);
434 }
435}