store_interval_tree/
iterators.rs

1use alloc::vec::Vec;
2use core::{cmp::Ord, fmt::Debug};
3
4use crate::{interval::Interval, node::Node};
5
6/// A `find` query on the interval tree does not directly return references to the nodes in the tree, but
7/// wraps the fields `interval` and `data` in an `Entry`.
8#[derive(PartialEq, Eq, Debug)]
9pub struct Entry<'a, T: Ord, V> {
10    value: &'a V,
11    interval: &'a Interval<T>,
12}
13
14impl<'a, T: Ord + 'a, V: 'a> Entry<'a, T, V> {
15    /// Get a reference to the data for this entry
16    #[must_use]
17    pub fn value(&self) -> &'a V {
18        self.value
19    }
20
21    /// Get a reference to the interval for this entry
22    #[must_use]
23    pub fn interval(&self) -> &'a Interval<T> {
24        self.interval
25    }
26}
27
28/// An `IntervalTreeIterator` is returned by `Intervaltree::find` and iterates over the entries
29/// overlapping the query
30#[derive(Debug)]
31pub struct IntervalTreeIterator<'v, 'i, T: Ord, V> {
32    pub(crate) nodes: Vec<&'v Node<T, V>>,
33    pub(crate) interval: &'i Interval<T>,
34}
35
36impl<'v, 'i, T: Ord + 'i, V: 'v> Iterator for IntervalTreeIterator<'v, 'i, T, V> {
37    type Item = Entry<'v, T, V>;
38
39    fn next(&mut self) -> Option<Entry<'v, T, V>> {
40        loop {
41            let node_ref = match self.nodes.pop() {
42                None => return None,
43                Some(node) => node,
44            };
45
46            if node_ref.right_child.is_some() {
47                self.nodes.push(node_ref.right_child.as_ref().unwrap());
48            }
49            if node_ref.left_child.is_some()
50                && Node::<T, V>::is_ge(
51                    &node_ref.left_child.as_ref().unwrap().get_max(),
52                    &self.interval.get_low(),
53                )
54            {
55                self.nodes.push(node_ref.left_child.as_ref().unwrap());
56            }
57
58            if Interval::overlaps(node_ref.interval(), self.interval) {
59                return Some(Entry {
60                    value: node_ref.value(),
61                    interval: node_ref.interval(),
62                });
63            }
64        }
65    }
66}
67
68/// A `find_mut` query on the interval tree does not directly return references to the nodes in the tree, but
69/// wraps the fields `interval` and `data` in an `EntryMut`. Only the data part can be mutably accessed
70/// using the `data` method
71#[derive(PartialEq, Eq, Debug)]
72pub struct EntryMut<'a, T: Ord, V> {
73    value: &'a mut V,
74    interval: &'a Interval<T>,
75}
76
77impl<'a, T: Ord + 'a, V: 'a> EntryMut<'a, T, V> {
78    /// Get a mutable reference to the data for this entry
79    pub fn value(&'a mut self) -> &'a mut V {
80        self.value
81    }
82
83    /// Get a reference to the interval for this entry
84    #[must_use]
85    pub fn interval(&self) -> &'a Interval<T> {
86        self.interval
87    }
88}
89
90/// An `IntervalTreeIteratorMut` is returned by `Intervaltree::find_mut` and iterates over the entries
91/// overlapping the query allowing mutable access to the data `D`, not the `Interval`.
92#[derive(Debug)]
93pub struct IntervalTreeIteratorMut<'v, 'i, T: Ord, V> {
94    pub(crate) nodes: Vec<&'v mut Node<T, V>>,
95    pub(crate) interval: &'i Interval<T>,
96}
97
98impl<'v, 'i, T: Ord + 'i, V: 'v> Iterator for IntervalTreeIteratorMut<'v, 'i, T, V> {
99    type Item = EntryMut<'v, T, V>;
100
101    fn next(&mut self) -> Option<EntryMut<'v, T, V>> {
102        loop {
103            let node_ref = match self.nodes.pop() {
104                None => return None,
105                Some(node) => node,
106            };
107
108            let overlaps = Interval::overlaps(node_ref.interval(), self.interval);
109
110            if node_ref.right_child.is_some() {
111                self.nodes.push(node_ref.right_child.as_mut().unwrap());
112            }
113            if node_ref.left_child.is_some()
114                && Node::<T, V>::is_ge(
115                    &node_ref.left_child.as_ref().unwrap().get_max(),
116                    &self.interval.get_low(),
117                )
118            {
119                self.nodes.push(node_ref.left_child.as_mut().unwrap());
120            }
121
122            if overlaps {
123                return Some(EntryMut {
124                    value: node_ref.value.as_mut().unwrap(),
125                    interval: node_ref.interval.as_ref().unwrap(),
126                });
127            }
128        }
129    }
130}