Skip to main content

embed_avl/
iter.rs

1use super::*;
2use core::marker::PhantomData;
3use core::ptr::null_mut;
4use pointers::Pointer;
5
6pub struct AvlIntoIter<P: Pointer, Tag>
7where
8    P: Pointer,
9    P::Target: AvlItem<Tag>,
10{
11    node: *mut P::Target,
12    dir: Dir,
13    _phan: PhantomData<fn(&Tag)>,
14}
15
16impl<P, Tag> AvlIntoIter<P, Tag>
17where
18    P: Pointer,
19    P::Target: AvlItem<Tag>,
20{
21    pub(crate) fn new(root: *mut P::Target, dir: Dir) -> Self {
22        Self { node: root, dir, _phan: Default::default() }
23    }
24}
25
26impl<P, Tag> Iterator for AvlIntoIter<P, Tag>
27where
28    P: Pointer,
29    P::Target: AvlItem<Tag>,
30{
31    type Item = P;
32
33    #[inline]
34    fn next(&mut self) -> Option<Self::Item> {
35        if !self.node.is_null() {
36            unsafe {
37                // self.node was the parent of last visited node
38                let item =
39                    AvlTree::<P, Tag>::bottom_child(NonNull::new_unchecked(self.node), self.dir);
40                let _node = item.as_ref().get_node();
41                let right_child = _node.get_child(self.dir.reverse());
42                let parent = _node.get_parent();
43                _node.detach();
44                if !parent.is_null() {
45                    if right_child.is_null() {
46                        (*parent).get_node().set_child(self.dir, null_mut());
47                        self.node = parent;
48                    } else {
49                        (*parent).get_node().set_child2::<P>(self.dir, right_child, parent);
50                        self.node = right_child;
51                    }
52                } else {
53                    // I'm the root
54                    if !right_child.is_null() {
55                        self.node = right_child;
56                        (*right_child).get_node().parent = null_mut();
57                    } else {
58                        self.node = null_mut();
59                    }
60                }
61                Some(P::from_raw(item.as_ptr()))
62            }
63        } else {
64            None
65        }
66    }
67}
68
69impl<P, Tag> Drop for AvlIntoIter<P, Tag>
70where
71    P: Pointer,
72    P::Target: AvlItem<Tag>,
73{
74    #[inline]
75    fn drop(&mut self) {
76        if self.node.is_null() {
77            return;
78        }
79        for _node in self.by_ref() {}
80    }
81}
82
83/// Return a reference of P (you can P.clone() of P has Clone)
84///
85/// # Example
86///
87/// ```rust
88/// use embed_avl::{AvlTree, AvlItem, AvlNode, impl_ord_with_borrow};
89/// use core::cell::UnsafeCell;
90/// use core::cmp::Ordering;
91/// use core::borrow::Borrow;
92/// extern crate alloc;
93/// use alloc::sync::Arc;
94///
95/// #[derive(format_attr::DebugAttr)]
96/// #[fmt("MyNode {value}")]
97/// struct MyNode {
98///     value: i32,
99///     avl_node: UnsafeCell<AvlNode<MyNode, ()>>,
100/// }
101///
102/// impl Borrow<i32> for MyNode {
103///     fn borrow(&self) -> &i32 {
104///         &self.value
105///     }
106/// }
107///
108/// unsafe impl AvlItem<()> for MyNode {
109///     fn get_node(&self) -> &mut AvlNode<MyNode, ()> {
110///         unsafe { &mut *self.avl_node.get() }
111///     }
112/// }
113/// impl_ord_with_borrow!(MyNode, i32);
114///
115/// let mut tree = AvlTree::<Arc<MyNode>, ()>::new();
116///
117/// let new_node = | n: i32| {
118///     Arc::new(MyNode { value: n, avl_node: UnsafeCell::new(Default::default()) })
119/// };
120/// tree.add(new_node(1)).expect("add");
121/// tree.add(new_node(2)).expect("add");
122/// tree.add(new_node(3)).expect("add");
123/// assert_eq!(tree.len(), 3);
124///
125/// let mut iter = tree.iter_rev();
126/// let mut all_nodes_rev = Vec::with_capacity(tree.len());
127/// while let Some(node) = iter.next() {
128///     all_nodes_rev.push(node.clone_ref());
129/// }
130/// ```
131pub struct AvlIter<'a, P, Tag>
132where
133    P: Pointer,
134    P::Target: AvlItem<Tag>,
135{
136    cur: Option<NonNull<P::Target>>,
137    dir: Dir,
138    left: usize,
139    _phan: PhantomData<fn(&'a P, &'a Tag)>,
140}
141
142/// A lending iterator for AvlTree
143///
144/// NOTE: The iterator yield [AvlNodeRef]. You can clone the reference with [AvlNodeRef::clone_ref] if it's `Arc`
145impl<'a, P, Tag> AvlIter<'a, P, Tag>
146where
147    P: Pointer,
148    P::Target: AvlItem<Tag>,
149{
150    #[inline]
151    pub(crate) fn new(
152        tree: &'a AvlTree<P, Tag>, init: Option<NonNull<P::Target>>, dir: Dir,
153    ) -> Self {
154        Self { cur: init, dir, left: tree.len(), _phan: Default::default() }
155    }
156}
157
158unsafe impl<'a, P, Tag> Send for AvlIter<'a, P, Tag>
159where
160    P: Pointer + Send,
161    P::Target: AvlItem<Tag>,
162{
163}
164
165impl<'a, P, Tag> Iterator for AvlIter<'a, P, Tag>
166where
167    P: Pointer,
168    P::Target: AvlItem<Tag>,
169{
170    type Item = AvlNodeRef<'a, P>;
171
172    #[inline]
173    fn next(&mut self) -> Option<Self::Item> {
174        let cur: NonNull<P::Target> = self.cur.take()?;
175        self.left -= 1;
176        self.cur = AvlTree::<P, Tag>::_walk_dir(cur, self.dir);
177        Some(AvlNodeRef::<'a, P>::new(cur.as_ptr()))
178    }
179}
180
181impl<'a, P, Tag> ExactSizeIterator for AvlIter<'a, P, Tag>
182where
183    P: Pointer,
184    P::Target: AvlItem<Tag>,
185{
186    #[inline]
187    fn len(&self) -> usize {
188        self.left
189    }
190}