use alloc::vec::Vec;
use core::ops::Range;
use crate::node::Node;
pub(crate) trait PruningOracle<R, V> {
fn visit_right(&self, subtree_root: &Node<R, V>, query: &Range<R>) -> bool;
fn filter_yield(&self, n: &Node<R, V>, query: &Range<R>) -> bool;
}
#[derive(Debug)]
pub(crate) struct PruningIter<'a, 'b, R, V, T> {
query: &'b Range<R>,
stack: Vec<&'a Node<R, V>>,
pruner: T,
}
impl<'a, 'b, R, V, T> PruningIter<'a, 'b, R, V, T>
where
R: Ord,
T: PruningOracle<R, V>,
{
pub(crate) fn new(root: &'a Node<R, V>, query: &'b Range<R>, pruner: T) -> Self {
let mut this = Self {
stack: vec![],
query,
pruner,
};
this.push_subtree(root);
this
}
fn push_subtree(&mut self, subtree_root: &'a Node<R, V>) {
let mut ptr = Some(subtree_root);
while let Some(v) = ptr {
self.stack.push(v);
ptr = v.left();
}
}
}
impl<'a, 'b, R, V, T> Iterator for PruningIter<'a, 'b, R, V, T>
where
R: Ord,
T: PruningOracle<R, V>,
{
type Item = &'a Node<R, V>;
fn next(&mut self) -> Option<Self::Item> {
loop {
let v = self.stack.pop()?;
if !self.pruner.visit_right(v, self.query) {
continue;
}
if let Some(right) = v.right() {
self.push_subtree(right);
}
if self.pruner.filter_yield(v, self.query) {
return Some(v);
}
}
}
}