use alloc::vec::Vec;
use bitstring::BitString;
use crate::walk_mut::NodeOrTree;
use super::{
goto::{
GotoStepResult,
NodeRef,
},
InsertPosition,
Node,
Tree,
TreeProperties,
WalkedDirection,
};
pub struct Walk<'r, TP: TreeProperties, D = (), A = ()> {
tree: Option<&'r Node<TP>>,
stack: Vec<(&'r Node<TP>, (D, A))>,
}
impl<'r, TP: TreeProperties, D, A> Walk<'r, TP, D, A> {
pub(in crate::tree) fn new(tree: &'r Tree<TP>) -> Self {
Self {
tree: tree.node.as_ref(),
stack: Vec::new(),
}
}
pub fn up(&mut self) -> Option<D> {
Some(self.stack.pop()?.1 .0)
}
pub fn up_with(&mut self) -> Option<(D, A)> {
Some(self.stack.pop()?.1)
}
pub fn current(&self) -> NodeOrTree<Option<&'r Node<TP>>, &'r Node<TP>> {
match self.stack.last() {
Some(&(node, _)) => NodeOrTree::Node(node),
None => NodeOrTree::Tree(self.tree),
}
}
}
impl<'r, TP: TreeProperties, D, A> Walk<'r, TP, D, A>
where
D: From<WalkedDirection>,
{
pub fn down_root_with(&mut self, add: A) -> bool {
if self.stack.is_empty() {
if let Some(root) = self.tree {
self.stack.push((root, (WalkedDirection::Down.into(), add)));
return true;
}
}
false
}
pub fn down_left_with(&mut self, add: A) -> bool {
self.down_with(false, add)
}
pub fn down_right_with(&mut self, add: A) -> bool {
self.down_with(true, add)
}
pub fn down_with(&mut self, side: bool, add: A) -> bool {
if let Some(&(node, _)) = self.stack.last() {
if let Some(child) = node.get_child(side) {
self.stack
.push((child, (WalkedDirection::from_side(side).into(), add)));
return true;
}
}
false
}
}
impl<'r, TP: TreeProperties, D> Walk<'r, TP, D, ()>
where
D: From<WalkedDirection>,
{
pub fn down_root(&mut self) -> bool {
self.down_root_with(())
}
pub fn down_left(&mut self) -> bool {
self.down_left_with(())
}
pub fn down_right(&mut self) -> bool {
self.down_right_with(())
}
pub fn down(&mut self, side: bool) -> bool {
self.down_with(side, ())
}
}
impl<'r, TP: TreeProperties, D> Walk<'r, TP, D>
where
D: From<WalkedDirection>,
{
fn goto_clean(&mut self, key: &TP::Key) {
let key_len = key.len();
while let Some(&(node, _)) = self.stack.last() {
if node._is_prefix_of(key, key_len) {
return;
}
self.stack.pop();
}
}
fn goto_insert_step(
&mut self,
key: &TP::Key,
key_len: usize,
) -> Result<(), Option<InsertPosition>> {
if let Some(&(node, _)) = self.stack.last() {
match node.goto_insert_step(key, key_len) {
GotoStepResult::Final(r) => Err(Some(r.into())),
GotoStepResult::Continue(node, dir) => {
self.stack.push((node, (dir.into(), ())));
Ok(())
},
}
} else if let Some(root) = self.tree {
self.stack.push((root, (WalkedDirection::Down.into(), ())));
Ok(())
} else {
Err(None)
}
}
fn goto_insert_down(&mut self, key: &TP::Key) -> Option<InsertPosition> {
let key_len = key.len();
loop {
match self.goto_insert_step(key, key_len) {
Ok(()) => (), Err(r) => return r, }
}
}
pub fn goto_insert(&mut self, key: &TP::Key) -> Option<InsertPosition> {
self.goto_clean(key);
self.goto_insert_down(key)
}
}
impl<'r, TP: TreeProperties> Walk<'r, TP, WalkedDirection> {
pub fn next_pre_order(&mut self) -> Option<&'r Node<TP>> {
match self.current() {
NodeOrTree::Tree(_) => {
self.down_root();
},
NodeOrTree::Node(node) => {
if node.is_leaf() {
loop {
match self.up()? {
WalkedDirection::Down => {
return None;
}, WalkedDirection::Left => {
self.down_right();
break;
},
WalkedDirection::Right => (), }
}
} else {
self.down_left();
}
},
}
return self.current().node();
}
pub fn next_in_order(&mut self) -> Option<&'r Node<TP>> {
match self.current() {
NodeOrTree::Tree(_) => {
self.down_root();
while self.down_left() {}
},
NodeOrTree::Node(node) => {
if node.is_leaf() {
loop {
match self.up()? {
WalkedDirection::Down => {
return None;
}, WalkedDirection::Left => {
break;
},
WalkedDirection::Right => (), }
}
} else {
self.down_right();
while self.down_left() {}
}
},
}
return self.current().node();
}
pub fn next_leaf(&mut self) -> Option<&'r Node<TP>> {
match self.current() {
NodeOrTree::Tree(_) => {
self.down_root();
while self.down_left() {}
},
NodeOrTree::Node(_) => {
loop {
match self.up()? {
WalkedDirection::Down => {
return None; },
WalkedDirection::Left => {
self.down_right();
while self.down_left() {}
break;
},
WalkedDirection::Right => (), }
}
},
}
return self.current().node();
}
pub fn next_post_order(&mut self) -> Option<&'r Node<TP>> {
match self.current() {
NodeOrTree::Tree(_) => {
self.down_root();
while self.down_left() {}
},
NodeOrTree::Node(_) => {
match self.up()? {
WalkedDirection::Down => {
return None;
}, WalkedDirection::Left => {
self.down_right();
while self.down_left() {}
},
WalkedDirection::Right => (),
}
},
}
return self.current().node();
}
}