use crate::{
iter::{
iter_between,
IterBetween,
},
tree::{
mut_gen::{
WalkMut,
WalkMutPath,
},
MutPath,
Node,
TreeProperties,
WalkedDirection,
},
};
use super::walk::OwnedTreeMarker;
unsafe fn extract_from_node<'r, TP: TreeProperties>(
node: &mut Node<TP>,
) -> (
&'r TP::Key,
&'r mut TP::Value,
Option<&'r mut TP::LeafValue>,
) {
let key: *const <TP as TreeProperties>::Key = node.get_key() as *const _;
let value = node.get_value_mut() as *mut _;
let leaf_value = node.get_leaf_value_mut().map(|v| v as *mut _);
(
unsafe { &*key },
unsafe { &mut *value },
leaf_value.map(|v| unsafe { &mut *v }),
)
}
pub struct IterMutPath<'r, TP: TreeProperties> {
path: MutPath<'r, TP>,
}
impl<'r, TP: TreeProperties> IterMutPath<'r, TP> {
pub(in crate::tree) fn new(path: MutPath<'r, TP>) -> Self {
Self { path }
}
}
impl<'r, TP: TreeProperties> Iterator for IterMutPath<'r, TP> {
type Item = (
&'r TP::Key,
&'r mut TP::Value,
Option<&'r mut TP::LeafValue>,
);
fn next(&mut self) -> Option<Self::Item> {
let node = self.path.next()?;
Some(unsafe { extract_from_node(node) })
}
}
pub struct IterWalkMutPath<'r, 'w, TP, O, D = ()>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, D, ()>,
{
path: WalkMutPath<'r, 'w, TP, O, D>,
}
impl<'r, 'w, TP, O, D> IterWalkMutPath<'r, 'w, TP, O, D>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, D, ()>,
{
pub(in crate::tree) fn new(path: WalkMutPath<'r, 'w, TP, O, D>) -> Self {
Self { path }
}
}
impl<'r, TP, O, D> Iterator for IterWalkMutPath<'r, '_, TP, O, D>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, D, ()>,
D: From<WalkedDirection>,
{
type Item = (
&'r TP::Key,
&'r mut TP::Value,
Option<&'r mut TP::LeafValue>,
);
fn next(&mut self) -> Option<Self::Item> {
let node = self.path.next()?;
Some(unsafe { extract_from_node(node) })
}
}
pub(in crate::tree) struct IterMutPreOrder<'r, TP, O>
where
TP: TreeProperties,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
pub(super) walk: WalkMut<'r, TP, O, WalkedDirection>,
}
impl<'r, TP, O> Iterator for IterMutPreOrder<'r, TP, O>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
type Item = (
&'r TP::Key,
&'r mut TP::Value,
Option<&'r mut TP::LeafValue>,
);
fn next(&mut self) -> Option<Self::Item> {
let node = self.walk.next_pre_order()?;
Some(unsafe { extract_from_node(node) })
}
}
pub(in crate::tree) struct IterMutInOrder<'r, TP, O>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
pub(super) walk: WalkMut<'r, TP, O, WalkedDirection>,
}
impl<'r, TP, O> Iterator for IterMutInOrder<'r, TP, O>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
type Item = (
&'r TP::Key,
&'r mut TP::Value,
Option<&'r mut TP::LeafValue>,
);
fn next(&mut self) -> Option<Self::Item> {
let node = self.walk.next_in_order()?;
Some(unsafe { extract_from_node(node) })
}
}
pub(in crate::tree) struct IterMutPostOrder<'r, TP, O>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
pub(super) walk: WalkMut<'r, TP, O, WalkedDirection>,
}
impl<'r, TP, O> Iterator for IterMutPostOrder<'r, TP, O>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
type Item = (
&'r TP::Key,
&'r mut TP::Value,
Option<&'r mut TP::LeafValue>,
);
fn next(&mut self) -> Option<Self::Item> {
let node = self.walk.next_post_order()?;
Some(unsafe { extract_from_node(node) })
}
}
pub(in crate::tree) struct IterMutLeaf<'r, TP, O>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
pub(super) walk: WalkMut<'r, TP, O, WalkedDirection>,
}
impl<'r, TP, O> Iterator for IterMutLeaf<'r, TP, O>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
type Item = (&'r TP::Key, &'r mut TP::LeafValue);
fn next(&mut self) -> Option<Self::Item> {
let node = self.walk.next_leaf()?;
let (key, _, leaf_value) = unsafe { extract_from_node(node) };
Some((key, leaf_value.expect("leaf node")))
}
}
pub(in crate::tree) struct IterMutLeafFull<'r, TP, O>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
walk: Option<WalkMut<'r, TP, O, WalkedDirection>>,
previous_key: Option<TP::Key>,
uncovered: IterBetween<TP::Key>,
next: Option<(TP::Key, &'r mut TP::LeafValue)>,
}
impl<'r, TP, O> IterMutLeafFull<'r, TP, O>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
pub(in crate::tree) fn new(walk: WalkMut<'r, TP, O, WalkedDirection>) -> Self {
Self {
walk: Some(walk),
previous_key: None,
uncovered: Default::default(),
next: None,
}
}
}
impl<'r, TP, O> Iterator for IterMutLeafFull<'r, TP, O>
where
TP: TreeProperties + 'r,
O: OwnedTreeMarker<'r, TP, WalkedDirection>,
{
type Item = (TP::Key, Option<&'r mut TP::LeafValue>);
fn next(&mut self) -> Option<Self::Item> {
loop {
if let Some(k) = self.uncovered.next() {
return Some((k, None));
}
if let Some((key, value)) = self.next.take() {
return Some((key, Some(value)));
}
match self.walk.as_mut()?.next_leaf() {
None => {
self.walk = None;
self.uncovered = iter_between(self.previous_key.clone(), None);
},
Some(node) => {
let (key, _, leaf_value) = unsafe { extract_from_node(node) };
let key = key.clone();
let leaf_value = leaf_value.expect("leaf node");
let start = core::mem::replace(&mut self.previous_key, Some(key.clone()));
self.uncovered = iter_between(start, Some(key.clone()));
self.next = Some((key, leaf_value));
},
}
}
}
}