embed_avl/

entry.rs

use super::*;
use core::cmp::Ordering;
use core::fmt::Debug;
use core::ops::Deref;
use pointers::{CloneablePointer, Pointer};

/// Result of a search operation in an [`AvlTree`].
pub enum Entry<'a, P, Tag, K>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    Vacant(VacantEntry<'a, P, Tag, K>),
    Occupied(OccupiedEntry<'a, P, Tag>),
}

impl<'a, P, Tag, K> Entry<'a, P, Tag, K>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    #[inline]
    pub fn unwrap_occupied(self) -> OccupiedEntry<'a, P, Tag>
    where
        K: Debug,
    {
        match self {
            Self::Occupied(ent) => ent,
            Self::Vacant(ent) => {
                panic!("called unwrap_occupied on vacent entry {:?}", ent.key);
            }
        }
    }

    #[inline]
    pub fn unwrap_vacant(self) -> VacantEntry<'a, P, Tag, K>
    where
        P::Target: Debug,
    {
        match self {
            Self::Vacant(ent) => ent,
            Self::Occupied(ent) => {
                panic!("called unwrap_occupied on vacent entry {:?}", ent.as_ref());
            }
        }
    }

    #[inline]
    pub fn try_occupied(self) -> Option<OccupiedEntry<'a, P, Tag>> {
        match self {
            Self::Occupied(ent) => Some(ent),
            Self::Vacant(_ent) => None,
        }
    }

    #[inline]
    pub fn try_vacant(self) -> Option<VacantEntry<'a, P, Tag, K>> {
        match self {
            Self::Vacant(ent) => Some(ent),
            Self::Occupied(_ent) => None,
        }
    }
}

impl<'a, P, Tag, K> From<Entry<'a, P, Tag, K>> for Option<OccupiedEntry<'a, P, Tag>>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    #[inline]
    fn from(ent: Entry<'a, P, Tag, K>) -> Option<OccupiedEntry<'a, P, Tag>> {
        ent.try_occupied()
    }
}

impl<'a, P, Tag, K> From<Entry<'a, P, Tag, K>> for Option<VacantEntry<'a, P, Tag, K>>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    #[inline]
    fn from(ent: Entry<'a, P, Tag, K>) -> Option<VacantEntry<'a, P, Tag, K>> {
        ent.try_vacant()
    }
}

/// Use to mark last entry
pub struct Last();

impl fmt::Debug for Last {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "last")
    }
}

impl<'a, P, Tag> Entry<'a, P, Tag, Last>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    /// Append node to the last entry, return the new entry.
    ///
    /// For efficient batch insert without redundant seeking to the last,
    /// only rotation will be triggered after insert.
    ///
    /// # Safety
    ///
    /// You have to make sure the node append in total ordering
    pub unsafe fn append_unchecked(self, node: P) -> Self {
        match self {
            Self::Occupied(ent) => Self::Occupied(unsafe { ent.insert_move_next_unchecked(node) }),
            Self::Vacant(ent) => {
                debug_assert!(
                    ent.node.is_null(),
                    "can only append to VacantEntry when tree is empty"
                );
                debug_assert_eq!(ent.direction, DIR_RIGHT);
                let p = node.into_raw() as *mut P::Target;
                ent.tree._insert(p, null_mut(), DIR_RIGHT);
                Self::Occupied(OccupiedEntry {
                    node: unsafe { NonNull::new_unchecked(p) },
                    tree: ent.tree,
                })
            }
        }
    }

    /// Append node to the last entry with check, return the new entry.
    ///
    /// For efficient batch insert without redundant seeking to the last,
    /// only rotation will be triggered after insert.
    ///
    /// # Arguments:
    /// - `node`: the node to insert
    /// - `cmp_f`: compare `P::Target` of the entry and `node`.
    ///   If it returns Ordering::Less, insert `node` and return `Ok`; otherwise return `Err(node)`
    pub fn append<F>(self, node: P, cmp_f: F) -> Result<Self, P>
    where
        F: FnOnce(&P::Target, &P::Target) -> Ordering,
    {
        match self {
            Self::Occupied(ent) => {
                if cmp_f(ent.get(), node.as_ref()) == Ordering::Less {
                    Ok(Self::Occupied(unsafe { ent.insert_move_next_unchecked(node) }))
                } else {
                    Err(node)
                }
            }
            Self::Vacant(ent) => {
                debug_assert!(
                    ent.node.is_null(),
                    "can only append to VacantEntry when tree is empty"
                );
                debug_assert_eq!(ent.direction, DIR_RIGHT);
                let p = node.into_raw() as *mut P::Target;
                ent.tree._insert(p, null_mut(), DIR_RIGHT);
                Ok(Self::Occupied(OccupiedEntry {
                    node: unsafe { NonNull::new_unchecked(p) },
                    tree: ent.tree,
                }))
            }
        }
    }
}

/// Use to mark first entry
pub struct First();

impl fmt::Debug for First {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "first")
    }
}

impl<'a, P, Tag> Entry<'a, P, Tag, First>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    /// Prepend node to the first entry, return the new entry.
    ///
    /// For efficient batch insert without redundant seeking to the first,
    /// only rotation will be triggered after insert.
    ///
    /// # Safety
    ///
    /// You have to make sure the node append in total ordering
    pub unsafe fn prepend_unchecked(self, node: P) -> Self {
        match self {
            Self::Occupied(ent) => Self::Occupied(unsafe { ent.insert_move_prev_unchecked(node) }),
            Self::Vacant(ent) => {
                debug_assert!(
                    ent.node.is_null(),
                    "can only append to VacantEntry when tree is empty"
                );
                debug_assert_eq!(ent.direction, DIR_LEFT);
                let p = node.into_raw() as *mut P::Target;
                ent.tree._insert(p, null_mut(), DIR_LEFT);
                Self::Occupied(OccupiedEntry {
                    node: unsafe { NonNull::new_unchecked(p) },
                    tree: ent.tree,
                })
            }
        }
    }

    /// Prepend node to the first entry with check, return the new entry.
    ///
    /// For efficient batch insert without redundant seeking to the first,
    /// only rotation will be triggered after insert.
    ///
    /// # Arguments:
    /// - `node`: the node to insert
    /// - `cmp_f`: compare `P::Target` of the `node` and entry.
    ///   If it returns Ordering::Less, insert `node` and return `Ok`; otherwise return `Err(node)`
    pub fn prepend<F>(self, node: P, cmp_f: F) -> Result<Self, P>
    where
        F: FnOnce(&P::Target, &P::Target) -> Ordering,
    {
        match self {
            Self::Occupied(ent) => {
                if cmp_f(node.as_ref(), ent.get()) == Ordering::Less {
                    Ok(Self::Occupied(unsafe { ent.insert_move_prev_unchecked(node) }))
                } else {
                    Err(node)
                }
            }
            Self::Vacant(ent) => {
                debug_assert!(
                    ent.node.is_null(),
                    "can only append to VacantEntry when tree is empty"
                );
                debug_assert_eq!(ent.direction, DIR_LEFT);
                let p = node.into_raw() as *mut P::Target;
                ent.tree._insert(p, null_mut(), DIR_LEFT);
                Ok(Self::Occupied(OccupiedEntry {
                    node: unsafe { NonNull::new_unchecked(p) },
                    tree: ent.tree,
                }))
            }
        }
    }
}

pub struct OccupiedEntry<'a, P, Tag>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    /// The matching node or the parent for insertion.
    pub(crate) node: NonNull<P::Target>,
    pub(crate) tree: &'a mut AvlTree<P, Tag>,
}

impl<'a, P, Tag> OccupiedEntry<'a, P, Tag>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    #[inline]
    pub fn get(&self) -> &'a P::Target {
        unsafe { self.node.as_ref() }
    }

    #[inline]
    pub fn clone_ref(&self) -> P
    where
        P: CloneablePointer,
    {
        unsafe { CloneablePointer::from_raw_and_clone(self.node.as_ptr()) }
    }

    /// Remove the node from the tree
    #[inline]
    pub fn remove(self) -> P {
        let p = self.node.as_ptr();
        unsafe {
            self.tree._remove(p);
            P::from_raw(p)
        }
    }

    /// Remove the next node if it exists and return the reference, without changing the current position.
    #[inline]
    pub fn remove_prev(&mut self) -> Option<P> {
        if let Some(prev) = AvlTree::<P, Tag>::_walk_dir(self.node, DIR_LEFT) {
            unsafe {
                self.tree._remove(prev.as_ptr());
                return Some(P::from_raw(prev.as_ptr()));
            }
        }
        None
    }

    /// Remove the next node if it exists and return the reference, without changing the current position.
    #[inline]
    pub fn remove_next(&mut self) -> Option<P> {
        if let Some(next) = AvlTree::<P, Tag>::_walk_dir(self.node, DIR_RIGHT) {
            unsafe {
                self.tree._remove(next.as_ptr());
                return Some(P::from_raw(next.as_ptr()));
            }
        }
        None
    }

    #[inline]
    pub fn peek_prev(&self) -> Option<&'a P::Target> {
        if let Some(p) = AvlTree::<P, Tag>::_walk_dir(self.node, DIR_LEFT) {
            Some(unsafe { p.as_ref() })
        } else {
            None
        }
    }

    #[inline]
    pub fn peek_next(&self) -> Option<&'a P::Target> {
        if let Some(p) = AvlTree::<P, Tag>::_walk_dir(self.node, DIR_RIGHT) {
            Some(unsafe { p.as_ref() })
        } else {
            None
        }
    }

    #[inline]
    pub fn move_prev(mut self) -> Option<Self> {
        if let Some(p) = AvlTree::<P, Tag>::_walk_dir(self.node, DIR_LEFT) {
            self.node = p;
            Some(self)
        } else {
            None
        }
    }

    #[inline]
    pub fn move_next(mut self) -> Option<Self> {
        if let Some(p) = AvlTree::<P, Tag>::_walk_dir(self.node, DIR_RIGHT) {
            self.node = p;
            Some(self)
        } else {
            None
        }
    }

    /// Insert before the current node.
    /// The entry position is unchanged.
    ///
    /// # Safety
    ///
    /// Caller should make sure the new_node is exactly before the current node.
    ///
    /// otherwise the new_node is way smaller than its ancestor
    /// (with value not adjacent to current node) will break the tree total ordering.
    #[inline]
    pub unsafe fn insert_prev_unchecked(&mut self, new_node: P) {
        unsafe {
            let new_ptr = new_node.into_raw() as *mut P::Target;
            self.tree._insert_here(new_ptr, self.node, DIR_LEFT);
        }
    }

    /// Insert before the current node,
    /// with entry position is changed to newly inserted node.
    ///
    /// # Safety
    ///
    /// Caller should make sure the new_node is exactly before the current node.
    ///
    /// otherwise the new_node is way smaller than its ancestor
    /// (with value not adjacent to current node) will break the tree total ordering.
    #[inline]
    pub unsafe fn insert_move_prev_unchecked(self, new_node: P) -> Self {
        unsafe {
            let new_ptr = new_node.into_raw() as *mut P::Target;
            self.tree._insert_here(new_ptr, self.node, DIR_LEFT);
            Self { node: NonNull::new_unchecked(new_ptr), tree: self.tree }
        }
    }

    /// Insert after the current node.
    /// The entry position is unchanged.
    ///
    /// # Safety
    ///
    /// Caller should make sure the new_node is exactly after the current node.
    ///
    /// Otherwise the new_node is way smaller than its ancestor
    /// (with value not adjacent to current node) will break the tree total ordering.
    #[inline]
    pub unsafe fn insert_next_unchecked(&mut self, new_node: P) {
        unsafe {
            let new_ptr = new_node.into_raw() as *mut P::Target;
            self.tree._insert_here(new_ptr, self.node, DIR_RIGHT);
        }
    }

    /// Insert after the current node,
    /// with entry position is changed to newly inserted node.
    ///
    /// # Safety
    ///
    /// Caller should make sure the new_node is exactly after the current node.
    ///
    /// Otherwise the new_node is way smaller than its ancestor
    /// (with value not adjacent to current node) will break the tree total ordering.
    #[inline]
    pub unsafe fn insert_move_next_unchecked(self, new_node: P) -> Self {
        unsafe {
            let new_ptr = new_node.into_raw() as *mut P::Target;
            self.tree._insert_here(new_ptr, self.node, DIR_RIGHT);
            Self { node: NonNull::new_unchecked(new_ptr), tree: self.tree }
        }
    }
}

impl<'a, P, Tag> Deref for OccupiedEntry<'a, P, Tag>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    type Target = P::Target;

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.as_ref()
    }
}

impl<'a, P, Tag> AsRef<P::Target> for OccupiedEntry<'a, P, Tag>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    #[inline]
    fn as_ref(&self) -> &P::Target {
        Self::get(self)
    }
}

/// # Examples
///
/// ```rust
/// use embed_avl::{AvlTree, AvlItem, AvlNode, Entry};
/// use core::cell::UnsafeCell;
/// use core::borrow::Borrow;
/// extern crate alloc;
/// use alloc::sync::Arc;
///
/// struct MyNode {
///     value: i32,
///     avl_node: UnsafeCell<AvlNode<MyNode, ()>>,
/// }
///
/// unsafe impl AvlItem<()> for MyNode {
///
///     fn get_node(&self) -> &mut AvlNode<MyNode, ()> {
///         unsafe { &mut *self.avl_node.get() }
///     }
/// }
/// impl Borrow<i32> for MyNode {
///     fn borrow(&self) -> &i32 {
///         &self.value
///     }
/// }
///
/// let mut tree = AvlTree::<Arc<MyNode>, ()>::new();
/// let key = 42;
/// if let Entry::Vacant(ent) = tree.entry(key) {
///     ent.insert_node(Arc::new(MyNode {
///         value: key,
///         avl_node: UnsafeCell::new(Default::default()),
///     }));
/// }
/// ```
pub struct VacantEntry<'a, P, Tag, K>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    pub(crate) node: *mut P::Target,
    pub(crate) direction: Dir,
    pub(crate) tree: &'a mut AvlTree<P, Tag>,
    pub(crate) key: K,
}

impl<'a, P, Tag, K> VacantEntry<'a, P, Tag, K>
where
    P: Pointer,
    P::Target: AvlItem<Tag>,
{
    /// Return the nearest node in the search result
    #[inline(always)]
    pub fn nearest(&self) -> Option<&P::Target> {
        if !self.node.is_null() { unsafe { self.node.as_ref() } } else { None }
    }

    #[inline]
    pub fn peek_prev(&self) -> Option<&P::Target> {
        if !self.node.is_null() {
            unsafe {
                if self.direction == DIR_RIGHT {
                    return Some(&*self.node);
                } else if let Some(node) =
                    AvlTree::<P, Tag>::_walk_dir(NonNull::new_unchecked(self.node), DIR_LEFT)
                {
                    return Some(node.as_ref());
                }
            }
        }
        None
    }

    #[inline]
    pub fn key(&self) -> &K {
        &self.key
    }

    #[inline]
    pub fn peek_next(&self) -> Option<&P::Target> {
        if !self.node.is_null() {
            unsafe {
                if self.direction == DIR_LEFT {
                    return Some(&*self.node);
                } else if let Some(node) =
                    AvlTree::<P, Tag>::_walk_dir(NonNull::new_unchecked(self.node), DIR_RIGHT)
                {
                    return Some(node.as_ref());
                }
            }
        }
        None
    }

    #[inline]
    pub fn move_prev(self) -> Option<OccupiedEntry<'a, P, Tag>> {
        if !self.node.is_null() {
            unsafe {
                let p = NonNull::new_unchecked(self.node);
                if self.direction == DIR_RIGHT {
                    return Some(OccupiedEntry { node: p, tree: self.tree });
                } else if let Some(node) = AvlTree::<P, Tag>::_walk_dir(p, DIR_LEFT) {
                    return Some(OccupiedEntry { node, tree: self.tree });
                }
            }
        }
        None
    }

    #[inline]
    pub fn move_next(self) -> Option<OccupiedEntry<'a, P, Tag>> {
        if !self.node.is_null() {
            unsafe {
                let p = NonNull::new_unchecked(self.node);
                if self.direction == DIR_LEFT {
                    return Some(OccupiedEntry { node: p, tree: self.tree });
                } else if let Some(node) = AvlTree::<P, Tag>::_walk_dir(p, DIR_RIGHT) {
                    return Some(OccupiedEntry { node, tree: self.tree });
                }
            }
        }
        None
    }

    /// # Safety
    ///
    /// Caller should make sure `new_node` have the same value to the search key to `entry()`,
    ///
    /// ignoring `K` captured in VacantEntry
    #[inline]
    pub unsafe fn insert_node_unchecked(self, new_node: P) -> OccupiedEntry<'a, P, Tag> {
        let new_ptr = new_node.into_raw() as *mut P::Target;
        self.tree._insert(new_ptr, self.node, self.direction);
        OccupiedEntry { node: unsafe { NonNull::new_unchecked(new_ptr) }, tree: self.tree }
    }
}

impl<'a, P, Tag, K> VacantEntry<'a, P, Tag, K>
where
    P: Pointer,
    P::Target: AvlItem<Tag> + Borrow<K>,
    K: PartialEq,
{
    /// insert new_node at the position, for entry search by borrowed key
    ///
    /// # Safety
    ///
    /// We have debug_assert_eq the key an the node
    #[inline]
    pub fn insert_node(self, new_node: P) -> OccupiedEntry<'a, P, Tag> {
        debug_assert!(&self.key == new_node.as_ref().borrow());
        let new_ptr = new_node.into_raw() as *mut P::Target;
        self.tree._insert(new_ptr, self.node, self.direction);
        OccupiedEntry { node: unsafe { NonNull::new_unchecked(new_ptr) }, tree: self.tree }
    }
}