intervaltree/
lib.rs

1#![no_std]
2#![warn(missing_docs)]
3//! A simple and generic implementation of an immutable interval tree.
4
5#[cfg(not(feature = "std"))]
6extern crate alloc;
7#[cfg(feature = "serde")]
8extern crate serde;
9#[cfg(feature = "std")]
10extern crate std;
11
12#[cfg(not(feature = "std"))]
13use alloc::vec::{IntoIter, Vec};
14use core::cmp;
15use core::fmt::{Debug, Formatter, Result as FmtResult};
16use core::iter::FromIterator;
17use core::ops::Range;
18use core::slice::Iter;
19#[cfg(feature = "serde")]
20use serde::{Deserialize, Serialize};
21use smallvec::SmallVec;
22#[cfg(feature = "std")]
23use std::vec::{IntoIter, Vec};
24
25/// An element of an interval tree.
26#[derive(Debug, Clone, PartialEq, Eq, Hash)]
27#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
28pub struct Element<K, V> {
29    /// The range associated with this element.
30    pub range: Range<K>,
31    /// The value associated with this element.
32    pub value: V,
33}
34
35impl<K, V> From<(Range<K>, V)> for Element<K, V> {
36    fn from(tup: (Range<K>, V)) -> Element<K, V> {
37        let (range, value) = tup;
38        Element { range, value }
39    }
40}
41
42#[derive(Clone, Debug, Hash)]
43#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
44struct Node<K, V> {
45    element: Element<K, V>,
46    max: K,
47}
48
49/// A simple and generic implementation of an immutable interval tree.
50///
51/// To build it, always use `FromIterator`. This is not very optimized
52/// as it takes `O(log n)` stack (it uses recursion) but runs in `O(n log n)`.
53#[derive(Clone, Debug, Hash)]
54#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
55pub struct IntervalTree<K, V> {
56    data: Vec<Node<K, V>>,
57}
58
59impl<K: Ord + Clone, V, I: Into<Element<K, V>>> FromIterator<I> for IntervalTree<K, V> {
60    fn from_iter<T: IntoIterator<Item = I>>(iter: T) -> Self {
61        let mut nodes: Vec<_> = iter.into_iter().map(|i| i.into())
62            .map(|element| Node { max: element.range.end.clone(), element }).collect();
63
64        nodes.sort_unstable_by(|a, b| a.element.range.start.cmp(&b.element.range.start));
65
66        if !nodes.is_empty() {
67            Self::update_max(&mut nodes);
68        }
69
70        IntervalTree { data: nodes }
71    }
72}
73
74/// An iterator over all the elements in the tree (in no particular order).
75pub struct TreeIter<'a, K: 'a, V: 'a>(Iter<'a, Node<K, V>>);
76
77impl<'a, K: 'a, V: 'a> Iterator for TreeIter<'a, K, V> {
78    type Item = &'a Element<K, V>;
79
80    fn next(&mut self) -> Option<Self::Item> {
81        self.0.next().map(|x| &x.element)
82    }
83}
84
85impl<'a, K: 'a + Ord, V: 'a> IntoIterator for &'a IntervalTree<K, V> {
86    type Item = &'a Element<K, V>;
87    type IntoIter = TreeIter<'a, K, V>;
88
89    fn into_iter(self) -> TreeIter<'a, K, V> {
90        self.iter()
91    }
92}
93
94/// An iterator that moves out of an interval tree.
95pub struct TreeIntoIter<K, V>(IntoIter<Node<K, V>>);
96
97impl<K, V> IntoIterator for IntervalTree<K, V> {
98    type Item = Element<K, V>;
99    type IntoIter = TreeIntoIter<K, V>;
100
101    fn into_iter(self) -> TreeIntoIter<K, V> {
102        TreeIntoIter(self.data.into_iter())
103    }
104}
105
106impl<K, V> Iterator for TreeIntoIter<K, V> {
107    type Item = Element<K, V>;
108
109    fn next(&mut self) -> Option<Element<K, V>> {
110        self.0.next().map(|x| x.element)
111    }
112}
113
114impl<K: Ord + Clone, V> IntervalTree<K, V> {
115    fn update_max(nodes: &mut [Node<K, V>]) -> K {
116        assert!(!nodes.is_empty());
117        let i = nodes.len() / 2;
118        if nodes.len() > 1 {
119            {
120                let (left, rest) = nodes.split_at_mut(i);
121                if !left.is_empty() {
122                    rest[0].max = cmp::max(rest[0].max.clone(), Self::update_max(left));
123                }
124            }
125
126            {
127                let (rest, right) = nodes.split_at_mut(i + 1);
128                if !right.is_empty() {
129                    rest[i].max = cmp::max(rest[i].max.clone(), Self::update_max(right));
130                }
131            }
132        }
133
134        nodes[i].max.clone()
135    }
136}
137
138impl<K: Ord, V> IntervalTree<K, V> {
139    fn todo(&self) -> TodoVec {
140        let mut todo = SmallVec::new();
141        if !self.data.is_empty() {
142            todo.push((0, self.data.len()));
143        }
144        todo
145    }
146
147    /// Queries the interval tree for all elements overlapping a given interval.
148    ///
149    /// This runs in `O(log n + m)`.
150    pub fn query(&self, range: Range<K>) -> QueryIter<K, V> {
151        QueryIter {
152            todo: self.todo(),
153            tree: self,
154            query: Query::Range(range),
155        }
156    }
157
158    /// Queries the interval tree for all elements containing a given point.
159    ///
160    /// This runs in `O(log n + m)`.
161    pub fn query_point(&self, point: K) -> QueryIter<K, V> {
162        QueryIter {
163            todo: self.todo(),
164            tree: self,
165            query: Query::Point(point),
166        }
167    }
168
169    /// Returns an iterator over all elements in the tree (in no particular order).
170    pub fn iter(&self) -> TreeIter<K, V> {
171        TreeIter(self.data.iter())
172    }
173
174    /// Returns an iterator over all elements in the tree, sorted by `Element.range.start`.
175    ///
176    /// This is currently identical to `IntervalTree::iter` because the internal structure
177    /// is already sorted this way, but may not be in the future.
178    pub fn iter_sorted(&self) -> impl Iterator<Item=&Element<K, V>> {
179        TreeIter(self.data.iter())
180    }
181}
182
183#[derive(Clone)]
184enum Query<K> {
185    Point(K),
186    Range(Range<K>),
187}
188
189impl<K: Ord> Query<K> {
190    fn point(&self) -> &K {
191        match *self {
192            Query::Point(ref k) => k,
193            Query::Range(ref r) => &r.start,
194        }
195    }
196
197    fn go_right(&self, start: &K) -> bool {
198        match *self {
199            Query::Point(ref k) => k >= start,
200            Query::Range(ref r) => &r.end > start,
201        }
202    }
203
204    fn intersect(&self, range: &Range<K>) -> bool {
205        match *self {
206            Query::Point(ref k) => k < &range.end,
207            Query::Range(ref r) => r.end > range.start && r.start < range.end,
208        }
209    }
210}
211
212type TodoVec = SmallVec<[(usize, usize); 16]>;
213
214/// Iterator for query results.
215pub struct QueryIter<'a, K: 'a, V: 'a> {
216    tree: &'a IntervalTree<K, V>,
217    todo: TodoVec,
218    query: Query<K>,
219}
220
221impl<'a, K: Ord + Clone, V> Clone for QueryIter<'a, K, V> {
222    fn clone(&self) -> Self {
223        QueryIter {
224            tree: self.tree,
225            todo: self.todo.clone(),
226            query: self.query.clone(),
227        }
228    }
229}
230
231impl<'a, K: Ord + Clone + Debug, V: Debug> Debug for QueryIter<'a, K, V> {
232    fn fmt(&self, fmt: &mut Formatter) -> FmtResult {
233        let v: Vec<_> = (*self).clone().collect();
234        write!(fmt, "{:?}", v)
235    }
236}
237
238impl<'a, K: Ord, V> Iterator for QueryIter<'a, K, V> {
239    type Item = &'a Element<K, V>;
240
241    fn next(&mut self) -> Option<&'a Element<K, V>> {
242        while let Some((s, l)) = self.todo.pop() {
243            let i = s + l/2;
244
245            let node = &self.tree.data[i];
246            if self.query.point() < &node.max {
247                // push left
248                {
249                    let leftsz = i - s;
250                    if leftsz > 0 {
251                        self.todo.push((s, leftsz));
252                    }
253                }
254
255                if self.query.go_right(&node.element.range.start) {
256                    // push right
257                    {
258                        let rightsz = l + s - i - 1;
259                        if rightsz > 0 {
260                            self.todo.push((i + 1, rightsz));
261                        }
262                    }
263
264                    // finally, search this
265                    if self.query.intersect(&node.element.range) {
266                        return Some(&node.element);
267                    }
268                }
269            }
270        }
271        None
272    }
273}
274
275#[cfg(test)]
276mod tests {
277    use core::iter;
278    use super::*;
279
280    fn verify(tree: &IntervalTree<u32, u32>, i: u32, expected: &[u32]) {
281        let mut v1: Vec<_> = tree.query_point(i).map(|x| x.value).collect();
282        v1.sort();
283        let mut v2: Vec<_> = tree.query(i..(i+1)).map(|x| x.value).collect();
284        v2.sort();
285        assert_eq!(v1, expected);
286        assert_eq!(v2, expected);
287    }
288
289    #[test]
290    fn it_works() {
291        let tree: IntervalTree<u32, u32> = [
292            (0..3, 1),
293            (1..4, 2),
294            (2..5, 3),
295            (3..6, 4),
296            (4..7, 5),
297            (5..8, 6),
298            (4..5, 7),
299            (2..7, 8),
300        ].iter().cloned().collect();
301
302
303        verify(&tree, 0, &[1]);
304        verify(&tree, 1, &[1, 2]);
305        verify(&tree, 2, &[1, 2, 3, 8]);
306        verify(&tree, 3, &[2, 3, 4, 8]);
307        verify(&tree, 4, &[3, 4, 5, 7, 8]);
308        verify(&tree, 5, &[4, 5, 6, 8]);
309        verify(&tree, 6, &[5, 6, 8]);
310        verify(&tree, 7, &[6]);
311        verify(&tree, 8, &[]);
312        verify(&tree, 9, &[]);
313    }
314
315    #[test]
316    fn empty() {
317        let tree: IntervalTree<u32, u32> = iter::empty::<Element<u32, u32>>().collect();
318        verify(&tree, 42, &[]);
319    }
320}
intervaltree/lib.rs

intervaltree/
lib.rs
