use core::iter;
use alloc::collections::VecDeque;
use crate::node::FrozenNode;
use crate::traverse::{DftEvent, StableShallowDepthFirstTraverser};
#[derive(Debug)]
pub struct NonAllocatingBreadthFirstTraverser<T> {
inner: StableShallowDepthFirstTraverser<T>,
root: Option<FrozenNode<T>>,
current_depth: usize,
}
impl<T> Clone for NonAllocatingBreadthFirstTraverser<T> {
#[inline]
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
root: self.root.clone(),
current_depth: self.current_depth,
}
}
}
impl<T> NonAllocatingBreadthFirstTraverser<T> {
#[must_use]
pub fn with_toplevel(toplevel: FrozenNode<T>, start_depth: usize) -> Self {
Self {
inner: StableShallowDepthFirstTraverser::with_toplevel(
Some(toplevel.clone()),
Some(start_depth),
),
root: Some(toplevel),
current_depth: start_depth,
}
}
fn next_in_current_depth(&mut self) -> Option<FrozenNode<T>> {
self.inner
.by_ref()
.filter_map(DftEvent::into_open)
.find_map(|(node, depth)| {
if depth == self.current_depth {
Some(node)
} else {
None
}
})
}
}
impl<T> Iterator for NonAllocatingBreadthFirstTraverser<T> {
type Item = (FrozenNode<T>, usize);
fn next(&mut self) -> Option<Self::Item> {
let _ = self.root.as_ref()?;
if let Some(node) = self.next_in_current_depth() {
return Some((node, self.current_depth));
}
self.current_depth += 1;
self.inner = StableShallowDepthFirstTraverser::with_toplevel(
self.root.clone(),
Some(self.current_depth),
);
if let Some(node) = self.next_in_current_depth() {
return Some((node, self.current_depth));
}
self.root = None;
None
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
if self.root.is_none() {
(0, Some(0))
} else {
(0, None)
}
}
}
impl<T> iter::FusedIterator for NonAllocatingBreadthFirstTraverser<T> {}
#[derive(Debug)]
enum QueuedEvent<T> {
Node(FrozenNode<T>),
IncrementDepth,
}
impl<T> Clone for QueuedEvent<T> {
#[inline]
fn clone(&self) -> Self {
match self {
Self::Node(node) => Self::Node(node.clone()),
Self::IncrementDepth => Self::IncrementDepth,
}
}
}
#[derive(Debug)]
pub struct AllocatingBreadthFirstTraverser<T> {
events: VecDeque<QueuedEvent<T>>,
current_depth: usize,
}
impl<T> Clone for AllocatingBreadthFirstTraverser<T> {
#[inline]
fn clone(&self) -> Self {
Self {
events: self.events.clone(),
current_depth: self.current_depth,
}
}
}
impl<T> AllocatingBreadthFirstTraverser<T> {
#[must_use]
pub fn with_toplevel(toplevel: FrozenNode<T>, start_depth: usize) -> Self {
let mut events = VecDeque::new();
if start_depth == 0 {
events.push_back(QueuedEvent::Node(toplevel));
events.push_back(QueuedEvent::IncrementDepth);
} else {
let nodes =
StableShallowDepthFirstTraverser::with_toplevel(Some(toplevel), Some(start_depth))
.filter_map(DftEvent::into_open)
.filter_map(|(node, depth)| {
if depth == start_depth {
Some(node)
} else {
None
}
})
.map(QueuedEvent::Node);
events.extend(nodes);
events.push_back(QueuedEvent::IncrementDepth);
}
Self {
events,
current_depth: start_depth,
}
}
#[inline]
#[must_use]
pub fn peek(&self) -> Option<(FrozenNode<T>, usize)> {
match self.events.front().cloned()? {
QueuedEvent::Node(next) => Some((next, self.current_depth)),
QueuedEvent::IncrementDepth => match self.events.get(1).cloned()? {
QueuedEvent::Node(next) => Some((next, self.current_depth + 1)),
QueuedEvent::IncrementDepth => unreachable!(
"[consistency] `IncrementDepth` must not appear right after `IncrementDepth`"
),
},
}
}
#[inline]
#[must_use]
pub fn capacity(&self) -> usize {
self.events.capacity()
}
#[inline]
pub fn shrink_to_fit(&mut self) {
self.events.shrink_to_fit()
}
#[inline]
pub fn shrink_to(&mut self, min_capacity: usize) {
self.events.shrink_to(min_capacity)
}
}
impl<T> Iterator for AllocatingBreadthFirstTraverser<T> {
type Item = (FrozenNode<T>, usize);
fn next(&mut self) -> Option<Self::Item> {
while let Some(ev) = self.events.pop_front() {
let next = match ev {
QueuedEvent::Node(v) => v,
QueuedEvent::IncrementDepth => {
if self.events.is_empty() {
self.events = Default::default();
return None;
}
self.current_depth += 1;
self.events.push_back(QueuedEvent::IncrementDepth);
continue;
}
};
self.events
.extend(next.children_stable().map(QueuedEvent::Node));
return Some((next, self.current_depth));
}
None
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.events.len();
if len <= 1 {
(0, Some(0))
} else {
(len - 1, None)
}
}
}