rbt_rs/
iter.rs

1use std::{iter::FusedIterator, mem::ManuallyDrop};
2
3use crate::{
4    Node,
5    NodeKind::{self},
6    NodePtr, RedBlackTree,
7};
8
9/// An iterator over shared references of the nodes of a [RedBlackTree].
10///
11/// This `struct` is created by the [`iter_node`] method on [RedBlackTree] or [Node].
12///
13/// [`iter_node`]: RedBlackTree::iter_node
14#[derive(Debug)]
15pub struct IterNode<'a, T> {
16    pub(crate) stack: Vec<(&'a Node<T>, bool)>,
17}
18
19impl<'a, T> Iterator for IterNode<'a, T> {
20    type Item = &'a Node<T>;
21
22    fn next(&mut self) -> Option<Self::Item> {
23        if let Some((top, recursed)) = self.stack.pop() {
24            let top = if !recursed {
25                self.stack.push((top, true));
26
27                let mut x = top;
28                while let NodeKind::Node(ref left) = x.left {
29                    x = left;
30                    self.stack.push((left, true));
31                }
32
33                self.stack.pop().unwrap().0
34            } else {
35                top
36            };
37
38            if let NodeKind::Node(ref right) = top.right {
39                self.stack.push((right, false));
40            }
41            return Some(top);
42        }
43        None
44    }
45}
46
47impl<'a, T> FusedIterator for IterNode<'a, T> {}
48
49/// An iterator over shared references of the values of a [RedBlackTree].
50///
51/// This `struct` is created by the [`iter`] method on [RedBlackTree] or [Node].
52///
53/// [`iter`]: RedBlackTree::iter
54#[derive(Debug)]
55pub struct Iter<'a, T> {
56    pub(crate) iter_node: IterNode<'a, T>,
57}
58
59impl<'a, T> Iterator for Iter<'a, T> {
60    type Item = &'a T;
61
62    fn next(&mut self) -> Option<Self::Item> {
63        self.iter_node.next().map(|n| &n.t)
64    }
65}
66
67impl<T> FusedIterator for Iter<'_, T> {}
68
69impl<'a, T> IntoIterator for &'a Node<T> {
70    type Item = &'a T;
71    type IntoIter = Iter<'a, T>;
72
73    fn into_iter(self) -> Self::IntoIter {
74        self.iter()
75    }
76}
77
78impl<'a, T> IntoIterator for &'a RedBlackTree<T> {
79    type Item = &'a T;
80    type IntoIter = Iter<'a, T>;
81
82    fn into_iter(self) -> Self::IntoIter {
83        self.iter()
84    }
85}
86
87/// An iterator over owned values of the values of a [RedBlackTree].
88///
89/// This `struct` is created by the [`into_iter`] method on [RedBlackTree] (provided by the [`IntoIterator`] trait).
90///
91/// [`into_iter`]: IntoIterator::into_iter
92#[derive(Debug)]
93pub struct IntoIter<T> {
94    stack: Vec<(NodePtr<T>, bool)>,
95    // ManuallyDrop stops RedBlackTree::drop as well as the drop-glue for RedBlackTree from being called.
96    // Since RedBlackTree contains only stack allocations apart from the nodes, this is fine.
97    _tree: ManuallyDrop<RedBlackTree<T>>,
98}
99
100// SAFETY: We don't access the T itself.
101unsafe impl<#[may_dangle] T> Drop for IntoIter<T> {
102    fn drop(&mut self) {
103        // Simply exhaust the iterator, which drops the unvisted allocated nodes.
104        while self.next().is_some() {}
105    }
106}
107
108impl<T> Iterator for IntoIter<T> {
109    type Item = T;
110
111    fn next(&mut self) -> Option<Self::Item> {
112        if let Some((top, recursed)) = self.stack.pop() {
113            let top = if !recursed {
114                self.stack.push((top.clone(), true));
115
116                let mut x = top;
117                while let NodeKind::Node(left) = x.left.clone() {
118                    x = left.clone();
119                    self.stack.push((left, true));
120                }
121
122                self.stack.pop().unwrap().0
123            } else {
124                top
125            };
126
127            let Node { right, t, .. } = Box::into_inner(unsafe { Box::from_raw(top.0.as_ptr()) });
128            if let Some(right) = right.node() {
129                self.stack.push((right, false));
130            }
131            return Some(t);
132        }
133        None
134    }
135}
136
137impl<T> FusedIterator for IntoIter<T> {}
138
139impl<T> IntoIterator for RedBlackTree<T> {
140    type Item = T;
141    type IntoIter = IntoIter<T>;
142
143    fn into_iter(self) -> Self::IntoIter {
144        if let Some(root) = self.root.clone().node() {
145            IntoIter {
146                stack: vec![(root, false)],
147                _tree: ManuallyDrop::new(self),
148            }
149        } else {
150            IntoIter {
151                stack: vec![],
152                _tree: ManuallyDrop::new(self),
153            }
154        }
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn works_iter() {
164        let tree = RedBlackTree::from_iter([243, 116, 212, 255, 177]);
165        assert_eq!(
166            tree.iter().map(|n| n).cloned().collect::<Vec<i32>>(),
167            &[116, 177, 212, 243, 255]
168        );
169    }
170
171    #[test]
172    fn works_into_iter() {
173        let tree = RedBlackTree::from_iter([15, 6, 3, 2, 4, 18]);
174        assert_eq!(
175            tree.into_iter().map(|n| n).collect::<Vec<i32>>(),
176            &[2, 3, 4, 6, 15, 18]
177        );
178    }
179
180    #[test]
181    fn no_double_drop_into_iter() {
182        let mut iter = RedBlackTree::from_iter([243, 116, 212, 255, 177]).into_iter();
183        iter.next();
184        iter.next();
185    }
186}
rbt_rs/iter.rs

rbt_rs/
iter.rs
