intrex 0.2.0

//! Intrusive red-black tree with items addressed by indices.
//!
//! # Node Access Traits
//!
//! The tree accessors manipulate nodes through application-provided types
//! (supplied through a `Nodes` generic parameter) implementing *node access
//! traits*.
//!
//! ## Color
//!
//! The following traits provide access to the color field associated with
//! nodes:
//!
//! | Trait           | Receiver         | Read | Write |
//! | --------------- | ---------------- | ---- | ----- |
//! | [`NodesRb`]     | `&self`          | x    |       |
//! | [`NodesRbMut`]  | `&mut self`      |      | x     |
//!
//! They have straightforward forwarding implementations for reference types
//! (`&T`, `&mut T`).
//!
//! ## Links
//!
//! The left/right child fields are accessed via the [`crate::bintree`]
//! module's traits, which this module builds upon.
//!
//! ## Key Comparison
//!
//! This module uses the key comparison traits defined in [`crate::bintree`]
//! to optionally support binary search trees.
//!
//! ## Values
//!
//! This module uses the node value access traits defined in
//! [`crate::node_data`].
//!
//! # Examples
//!
//! ```rust
//! use std::ops::Bound;
//! use intrex::rbtree::{Tree, Link, map_link_cmp_mut};
//!
//! type Node = (i32, Link);
//!
//! let mut nodes: Vec<Node> = Vec::new();
//! let mut tree = Tree::default();
//!
//! let mut callbacks = map_link_cmp_mut(
//!     &mut nodes,
//!     |x: &Node| &x.1,
//!     |x: &mut Node| &mut x.1,
//!     |x: &Node, y: &Node| Ord::cmp(&x.0, &y.0),
//!     |x: &Node, y: &i32| Ord::cmp(&x.0, y),
//! );
//!
//! // Add nodes
//! for value in [3, 4, 5, 2, 1] {
//!     let mut i = callbacks.pool.len();
//!     callbacks.pool.push((value, Link::default()));
//!     tree.write(&mut callbacks).bst_insert_node(i);
//! }
//!
//! // Remove node
//! let i = tree.read(&callbacks).bst_find_index(&2).unwrap();
//! tree.write(&mut callbacks).remove_node(i);
//!
//! // Find nodes
//! assert!(tree.read(&callbacks).bst_find_index(&3).is_some());
//! assert!(tree.read(&callbacks).bst_find_index(&2).is_none());
//!
//! // Enumerate the nodes in order
//! let values: Vec<i32> = tree
//!     .read(&callbacks)
//!     .values()
//!     .map(|node| node.0)
//!     .collect();
//! assert_eq!(values, vec![1, 3, 4, 5]);
//! ```
#![warn(missing_docs)]

pub mod accessor;
pub mod accessor_mut;
mod nodes;

pub use nodes::*;

/// Red-black tree header.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Tree<Index = usize> {
    /// The binary tree header.
    pub raw: crate::bintree::Tree<Index>,
}

impl<Index> Default for Tree<Index> {
    #[inline]
    fn default() -> Self {
        Self {
            raw: <_>::default(),
        }
    }
}

impl<Index> Tree<Index> {
    /// Construct an empty tree header.
    #[inline]
    pub fn new() -> Self {
        Self::default()
    }

    /// Check if the tree is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.raw.is_empty()
    }
}

/// Node color in a red-black tree.
#[derive(Default, Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum Color {
    /// Black node.
    ///
    /// The red-black tree invariants require that every path from the root node
    /// to a leaf node contain the same number of black nodes (black height).
    ///
    /// All `None` nodes are considered black.
    #[default]
    Black,
    /// Red node.
    ///
    /// The red-black tree invariants require that a red node not have
    /// a red node child.
    Red,
}

/// Per-node link data stored in a red-black tree element.
///
/// This struct contains all fields used by
/// [`NodesWithMapLink`][]\[[`Mut`][2]\] to implement
/// [`NodesRb`][]\[[`Mut`][1]\].
///
/// [1]: NodesRbMut
/// [2]: NodesWithMapLinkMut
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Link<Index = usize> {
    /// The index of the parent.
    pub parent: Option<Index>,
    /// The index of the left child.
    pub left: Option<Index>,
    /// The index of the right child.
    pub right: Option<Index>,
    /// The color of this node.
    pub color: Color,
}

impl<Index> Link<Index> {
    /// The default value of `Self`.
    pub const DEFAULT: Self = Self {
        parent: None,
        left: None,
        right: None,
        color: Color::Black,
    };
}

impl<Index> Default for Link<Index> {
    #[inline]
    fn default() -> Self {
        Self::DEFAULT
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        bintree::{NodesCmp, NodesCmpKey, NodesLink, NodesLinkMut},
        node_data::{NodesDataLend, NodesDataLendGat, NodesDataLendMut, NodesDataLendMutGat},
    };

    use super::*;
    use core::cmp::Ordering;
    use proptest::prelude::*;
    use std::{prelude::rust_2024::*, rc::Rc};

    #[derive(Default, Debug, Clone, Copy, PartialEq)]
    pub(crate) struct Node<T> {
        pub value: T,
        pub color: Color,
        pub left: Option<usize>,
        pub right: Option<usize>,
        pub parent: Option<usize>,
    }

    impl<T> NodesLink<usize> for [Node<T>] {
        fn node_parent(&self, node: usize) -> Option<usize> {
            self[node].parent
        }

        fn node_left(&self, node: usize) -> Option<usize> {
            self[node].left
        }

        fn node_right(&self, node: usize) -> Option<usize> {
            self[node].right
        }
    }

    impl<T> NodesLinkMut<usize> for [Node<T>] {
        fn set_node_parent(&mut self, node: usize, parent: Option<usize>) {
            self[node].parent = parent;
        }

        fn set_node_left(&mut self, node: usize, left: Option<usize>) {
            self[node].left = left;
        }

        fn set_node_right(&mut self, node: usize, right: Option<usize>) {
            self[node].right = right;
        }
    }

    impl<T> NodesRb<usize> for [Node<T>]
    where
        T: Ord,
    {
        fn node_color(&self, node: usize) -> Color {
            self[node].color
        }
    }

    impl<T> NodesRbMut<usize> for [Node<T>]
    where
        T: Ord,
    {
        fn set_node_color(&mut self, node: usize, color: Color) {
            self[node].color = color;
        }
    }

    impl<T> NodesCmp<usize> for [Node<T>]
    where
        T: Ord,
    {
        fn cmp_nodes(&self, lhs: usize, rhs: usize) -> core::cmp::Ordering {
            Ord::cmp(&self[lhs].value, &self[rhs].value)
        }
    }

    impl<T: Ord> NodesCmpKey<usize, T> for [Node<T>] {
        fn cmp_node_key(&self, node: usize, key: &T) -> Ordering {
            Ord::cmp(&self[node].value, key)
        }
    }

    impl<'this, T> NodesDataLendGat<'this, usize> for [Node<T>] {
        type Data = &'this T;
    }

    impl<T> NodesDataLend<usize> for [Node<T>] {
        fn node_data_lend(&self, node: usize) -> <Self as NodesDataLendGat<'_, usize>>::Data {
            &self[node].value
        }
    }

    impl<'this, T> NodesDataLendMutGat<'this, usize> for [Node<T>] {
        type Data = &'this mut T;
    }

    impl<T> NodesDataLendMut<usize> for [Node<T>] {
        fn node_data_lend_mut(
            &mut self,
            node: usize,
        ) -> <Self as NodesDataLendMutGat<'_, usize>>::Data {
            &mut self[node].value
        }
    }

    /// Construct a proptest strategy producing random red-black BSTs.
    pub(crate) fn any_rbbst(allow_empty: bool) -> impl Strategy<Value = (Vec<Node<u8>>, Tree)> {
        #[derive(Debug, Default)]
        struct INode {
            color: Color,
            left: Option<Box<INode>>,
            right: Option<Box<INode>>,
        }

        ((!allow_empty) as _..=3)
            .prop_flat_map(|depth| {
                let mut cur = Rc::new(((|| None) as fn() -> Option<INode>).boxed());

                for _ in 0..depth {
                    let child = || cur.clone().prop_map(|c| c.map(Box::new));
                    let root = prop_oneof![
                        prop::array::uniform(child()).prop_map(|[left, right]| Some(INode {
                            color: Color::Black,
                            left,
                            right
                        })),
                        prop::array::uniform(child()).prop_map(|[left1, left2, right]| Some(
                            INode {
                                color: Color::Black,
                                left: Some(Box::new(INode {
                                    color: Color::Red,
                                    left: left1,
                                    right: left2
                                })),
                                right
                            }
                        )),
                        prop::array::uniform(child()).prop_map(|[left, right1, right2]| Some(
                            INode {
                                color: Color::Black,
                                left,
                                right: Some(Box::new(INode {
                                    color: Color::Red,
                                    left: right1,
                                    right: right2
                                })),
                            }
                        )),
                        prop::array::uniform(child()).prop_map(|[left1, left2, right1, right2]| {
                            Some(INode {
                                color: Color::Black,
                                left: Some(Box::new(INode {
                                    color: Color::Red,
                                    left: left1,
                                    right: left2,
                                })),
                                right: Some(Box::new(INode {
                                    color: Color::Red,
                                    left: right1,
                                    right: right2,
                                })),
                            })
                        }),
                    ];
                    cur = Rc::new(root.boxed());
                }

                (cur, any::<bool>())
            })
            .prop_map(|(mut root, make_root_red)| {
                // Paint the root red if `make_root_red == true` and it does
                // not violate the tree invariants
                if let Some(root) = &mut root
                    && make_root_red
                    && root.left.as_ref().is_none_or(|n| n.color == Color::Black)
                    && root.right.as_ref().is_none_or(|n| n.color == Color::Black)
                {
                    root.color = Color::Red;
                }

                struct State {
                    next_value: u8,
                    nodes: Vec<Node<u8>>,
                }

                fn convert_node(state: &mut State, node: Option<&INode>) -> Option<usize> {
                    let node = node?;

                    let left = convert_node(state, node.left.as_deref());

                    let value = state.next_value;
                    state.next_value += 1;

                    let right = convert_node(state, node.right.as_deref());

                    let i = state.nodes.len();
                    state.nodes.push(Node {
                        value,
                        color: node.color,
                        left,
                        right,
                        parent: None,
                    });

                    for k in [left, right].into_iter().flatten() {
                        state.nodes[k].parent = Some(i);
                    }

                    Some(i)
                }

                let mut state = State {
                    next_value: 1,
                    nodes: Vec::new(),
                };
                let tree = Tree {
                    raw: crate::bintree::Tree {
                        root: convert_node(&mut state, root.as_ref()),
                    },
                };

                (state.nodes, tree)
            })
    }

    /// Check that [`any_rbbst`] produces valid trees.
    #[proptest::property_test]
    fn pt_any_rbbst_validate(#[strategy = any_rbbst(true)] (nodes, root): (Vec<Node<u8>>, Tree)) {
        root.read(&nodes[..]).validate().unwrap();
        root.read(&nodes[..])
            .bst_validate(|nodes, i| nodes[i])
            .unwrap();
        assert!(root.read(&nodes[..]).values().is_sorted());
    }
}