aatree/
iter.rs

1#![allow(missing_debug_implementations)]
2
3//! Iterator implementations for [`AATreeSet`](crate::AATreeSet) and [`AATreeMap`](crate::AATreeMap).
4
5use super::node::{AANode, Node};
6use alloc::vec::Vec;
7use core::{iter::FusedIterator, marker::PhantomData};
8
9/// This trait allows iterators to return elements other than that stored inside the tree. Useful
10/// for returning key-value-pairs from `AATreeMap`.
11pub trait IterContent<T> {
12	fn content(self) -> T;
13}
14
15impl<T> IterContent<T> for T {
16	fn content(self) -> Self {
17		self
18	}
19}
20
21/// The iterator produces from an reference of an AATree-based data structure when turned into an iterator.
22pub struct AAIter<'a, C, T> {
23	stack: Vec<(bool, &'a AANode<C>)>,
24	len: usize,
25	_ty: PhantomData<T>
26}
27
28impl<'a, C, T> AAIter<'a, C, T> {
29	pub(super) fn new(root: &'a AANode<C>, len: usize) -> Self {
30		let mut stack = Vec::with_capacity(root.level() as usize * 2 + 1);
31		stack.push((false, root));
32		Self {
33			stack,
34			len,
35			_ty: PhantomData::default()
36		}
37	}
38}
39
40impl<'a, C, T> Iterator for AAIter<'a, C, T>
41where
42	&'a C: IterContent<T>
43{
44	type Item = T;
45
46	fn next(&mut self) -> Option<T> {
47		loop {
48			let (visited_left, last) = self.stack.pop()?;
49			if let Some(Node { left_child, .. }) = last.as_ref() {
50				self.stack.push((true, last));
51				if !visited_left && !left_child.is_nil() {
52					self.stack.push((false, left_child));
53				} else {
54					break;
55				}
56			}
57		}
58
59		let (_, last) = self.stack.pop()?;
60		match last.as_ref() {
61			None => unreachable!(),
62			Some(Node {
63				content,
64				right_child,
65				..
66			}) => {
67				if !right_child.is_nil() {
68					self.stack.push((false, right_child));
69				}
70				self.len -= 1;
71				Some(content.content())
72			}
73		}
74	}
75
76	fn size_hint(&self) -> (usize, Option<usize>) {
77		(self.len, Some(self.len))
78	}
79}
80
81impl<'a, C, T> ExactSizeIterator for AAIter<'a, C, T> where &'a C: IterContent<T> {}
82
83impl<'a, C, T> FusedIterator for AAIter<'a, C, T> where &'a C: IterContent<T> {}
84
85/// The iterator produces from an AATree-based data structure when turned into an iterator.
86pub struct AAIntoIter<C, T> {
87	stack: Vec<AANode<C>>,
88	len: usize,
89	_ty: PhantomData<T>
90}
91
92impl<C, T> AAIntoIter<C, T> {
93	pub(super) fn new(root: AANode<C>, len: usize) -> Self {
94		let mut stack = Vec::with_capacity(root.level() as usize * 2 + 1);
95		stack.push(root);
96		Self {
97			stack,
98			len,
99			_ty: PhantomData::default()
100		}
101	}
102}
103
104impl<C, T> Iterator for AAIntoIter<C, T>
105where
106	C: IterContent<T>
107{
108	type Item = T;
109
110	fn next(&mut self) -> Option<T> {
111		loop {
112			let last = self.stack.pop()?;
113			if let Some(Node {
114				level,
115				content,
116				left_child,
117				right_child
118			}) = last.unbox()
119			{
120				self.stack.push(
121					Node {
122						level,
123						content,
124						left_child: AANode::new(),
125						right_child
126					}
127					.into()
128				);
129				if !left_child.is_nil() {
130					self.stack.push(left_child);
131				} else {
132					break;
133				}
134			}
135		}
136
137		let last = self.stack.pop()?;
138		match last.unbox() {
139			None => unreachable!(),
140			Some(Node {
141				content,
142				right_child,
143				..
144			}) => {
145				if !right_child.is_nil() {
146					self.stack.push(right_child);
147				}
148				self.len -= 1;
149				Some(content.content())
150			}
151		}
152	}
153
154	fn size_hint(&self) -> (usize, Option<usize>) {
155		(self.len, Some(self.len))
156	}
157}
158
159impl<C, T> ExactSizeIterator for AAIntoIter<C, T> where C: IterContent<T> {}
160
161impl<C, T> FusedIterator for AAIntoIter<C, T> where C: IterContent<T> {}
aatree/iter.rs

aatree/
iter.rs
