intrex 0.2.0

//! Mutable binary tree accessor.
use crate::node_data::{NodesDataLendMut, NodesDataLendMutGat};

use super::*;

impl<Index> Tree<Index> {
    /// Create a mutable accessor to the tree, using `Nodes: `[`NodesLinkMut`]
    /// to access nodes.
    #[inline]
    pub fn write<Nodes>(&mut self, nodes: Nodes) -> TreeAccessorMut<'_, Nodes, Index>
    where
        Nodes: NodesLinkMut<Index>,
    {
        TreeAccessorMut { tree: self, nodes }
    }
}

/// Accessor to a binary tree.
#[derive(Debug)]
pub struct TreeAccessorMut<'head, Nodes, Index> {
    tree: &'head mut Tree<Index>,
    nodes: Nodes,
}

impl<'tree, Nodes, Index> TreeAccessorMut<'tree, Nodes, Index>
where
    Nodes: NodesLinkMut<Index>,
    Index: PartialEq + Clone,
{
    /// Borrow the node accessor.
    #[inline]
    pub fn nodes(&self) -> &Nodes {
        &self.nodes
    }

    /// Mutably borrow the node accessor.
    #[inline]
    pub fn nodes_mut(&mut self) -> &mut Nodes {
        &mut self.nodes
    }

    /// Borrow the tree root.
    #[inline]
    pub fn tree(&self) -> &Tree<Index> {
        self.tree
    }
    /// Mutably borrow the tree root.
    #[inline]
    pub fn tree_mut(&mut self) -> &mut Tree<Index> {
        self.tree
    }

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

    /// Get the first (leftmost) element's index.
    #[inline]
    pub fn front_index(&self) -> Option<Index> {
        self.front_index_in(self.tree.root.clone())
    }

    /// Get the last (rightmost) element's index.
    #[inline]
    pub fn back_index(&self) -> Option<Index> {
        self.back_index_in(self.tree.root.clone())
    }

    /// Get the first (smallest) child's index in `node`'s subtree.
    #[inline]
    fn front_index_in(&self, node: Option<Index>) -> Option<Index>
    where
        Index: Clone,
    {
        core::iter::successors(node, |i| self.nodes.node_left(i.clone())).last()
    }

    /// Get the last (largest) child's index in `node`'s subtree.
    #[inline]
    fn back_index_in(&self, node: Option<Index>) -> Option<Index>
    where
        Index: Clone,
    {
        core::iter::successors(node, |i| self.nodes.node_right(i.clone())).last()
    }

    /// Borrow the first element.
    #[doc(alias = "first_mut")]
    #[inline]
    pub fn front_mut(&mut self) -> Option<<Nodes as NodesDataLendMutGat<'_, Index>>::Data>
    where
        Nodes: NodesDataLendMut<Index>,
    {
        self.front_index().map(|p| self.nodes.node_data_lend_mut(p))
    }

    /// Borrow the last element.
    #[doc(alias = "last_mut")]
    #[inline]
    pub fn back_mut(&mut self) -> Option<<Nodes as NodesDataLendMutGat<'_, Index>>::Data>
    where
        Nodes: NodesDataLendMut<Index>,
    {
        self.back_index().map(|p| self.nodes.node_data_lend_mut(p))
    }

    /// Update a [`Slot`].
    ///
    /// Warning: This method updates one direction of a link but does not the
    /// other. The other direction must be updated manually to maintain
    /// validity.
    #[inline]
    pub fn set_slot(&mut self, slot: Slot<Index>, new_target: Option<Index>) {
        match slot {
            Slot::Root => self.tree.root = new_target,
            Slot::Left(node) => self.nodes.set_node_left(node, new_target),
            Slot::Right(node) => self.nodes.set_node_right(node, new_target),
        }
    }

    /// Call the specified closure for every element, passing a mutable
    /// reference for each of them.
    #[inline]
    pub fn for_each_mut<B>(
        &mut self,
        mut f: impl FnMut(Index, <Nodes as NodesDataLendMutGat<'_, Index>>::Data) -> ops::ControlFlow<B>,
    ) -> ops::ControlFlow<B>
    where
        Nodes: NodesDataLendMut<Index>,
    {
        let Some(first) = self.front_index() else {
            return ops::ControlFlow::Continue(());
        };
        let mut current = first.clone();
        loop {
            let next = next_index(&self.nodes, current.clone());
            f(current.clone(), self.nodes.node_data_lend_mut(current))?;
            let Some(next) = next else { break };
            current = next;
        }
        ops::ControlFlow::Continue(())
    }

    /// Call the specified closure for every element in a given in-order range,
    /// passing a mutable reference for each of them.
    #[inline]
    pub fn for_each_in_index_range_mut<B>(
        &mut self,
        range: impl ops::RangeBounds<Option<Index>>,
        mut f: impl FnMut(Index, <Nodes as NodesDataLendMutGat<'_, Index>>::Data) -> ops::ControlFlow<B>,
    ) -> ops::ControlFlow<B>
    where
        Nodes: NodesDataLendMut<Index>,
    {
        let (Some([first, last]), _) = self.tree().read(&self.nodes).resolve_index_range(range)
        else {
            return ops::ControlFlow::Continue(());
        };
        let mut current = first;
        loop {
            let next = (current != last).then(|| next_index(&self.nodes, current.clone()));
            f(current.clone(), self.nodes.node_data_lend_mut(current))?;
            let Some(Some(next)) = next else { break };
            current = next;
        }
        ops::ControlFlow::Continue(())
    }
}

/// # Tree Mutation
impl<'tree, Nodes, Index> TreeAccessorMut<'tree, Nodes, Index>
where
    Nodes: NodesLinkMut<Index>,
    Index: PartialEq + Clone,
{
    /// Replace the in-tree node `node` with the out-of-tree node `new_node`.
    ///
    /// This method transfers all tree links accessible via [`NodesLinkMut`]
    /// (`parent`, `left`, and `right`) from `node` to `new_node` and removes
    /// `node` from the tree.
    ///
    /// This method does not modify `node`'s link fields.
    pub fn replace_node(&mut self, node: Index, new_node: Index) {
        // `node` is in-tree and `new_node` is out-of-tree, so they must not
        // be the same node
        assert!(node != new_node);

        // Update outbound links
        let parent = self.nodes.node_parent(node.clone());
        let left = self.nodes.node_left(node.clone());
        let right = self.nodes.node_right(node.clone());
        self.nodes.set_node_parent(new_node.clone(), parent.clone());
        self.nodes.set_node_left(new_node.clone(), left.clone());
        self.nodes.set_node_right(new_node.clone(), right.clone());

        // Update inbound links
        if let Some(parent) = parent {
            if self.nodes.node_left(parent.clone()).as_ref() == Some(&node) {
                self.nodes
                    .set_node_left(parent.clone(), Some(new_node.clone()));
            } else if self.nodes.node_right(parent.clone()).as_ref() == Some(&node) {
                self.nodes
                    .set_node_right(parent.clone(), Some(new_node.clone()));
            } else {
                unreachable!("`node.parent` does not contain `node`")
            }
        } else {
            self.tree.root = Some(new_node.clone());
        }

        if let Some(left) = left {
            self.nodes.set_node_parent(left, Some(new_node.clone()));
        }

        if let Some(right) = right {
            self.nodes.set_node_parent(right, Some(new_node.clone()));
        }
    }

    /// Swap the positions of two in-tree nodes `a` and `b`.
    pub fn swap_nodes(&mut self, a: Index, b: Index) {
        if a == b {
            return;
        }

        let slot_a = self.tree().read(&self.nodes).parent_slot(a.clone());
        let parent_a = slot_a.clone().parent();
        let left_a = self.nodes.node_left(a.clone());
        let right_a = self.nodes.node_right(a.clone());

        let slot_b = self.tree().read(&self.nodes).parent_slot(b.clone());
        let parent_b = slot_b.clone().parent();
        let left_b = self.nodes.node_left(b.clone());
        let right_b = self.nodes.node_right(b.clone());

        // Update links from `a`
        self.nodes.set_node_left(
            a.clone(),
            if left_b.as_ref() == Some(&a) {
                Some(b.clone())
            } else {
                left_b.clone()
            },
        );
        self.nodes.set_node_right(
            a.clone(),
            if right_b.as_ref() == Some(&a) {
                Some(b.clone())
            } else {
                right_b.clone()
            },
        );

        if let Some(left_a) = &left_a
            && left_a != &b
        {
            self.nodes.set_node_parent(left_a.clone(), Some(b.clone()));
        }
        if let Some(right_a) = &right_a
            && right_a != &b
        {
            self.nodes.set_node_parent(right_a.clone(), Some(b.clone()));
        }

        // Update links from `b`
        self.nodes.set_node_left(
            b.clone(),
            if left_a.as_ref() == Some(&b) {
                Some(a.clone())
            } else {
                left_a.clone()
            },
        );
        self.nodes.set_node_right(
            b.clone(),
            if right_a.as_ref() == Some(&b) {
                Some(a.clone())
            } else {
                right_a.clone()
            },
        );

        if let Some(left_b) = &left_b
            && left_b != &a
        {
            self.nodes.set_node_parent(left_b.clone(), Some(a.clone()));
        }
        if let Some(right_b) = &right_b
            && right_b != &a
        {
            self.nodes.set_node_parent(right_b.clone(), Some(a.clone()));
        }

        // Update links to `a`
        if parent_a.as_ref() != Some(&b) {
            self.set_slot(slot_a, Some(b.clone()));
        }

        self.nodes.set_node_parent(
            a.clone(),
            if parent_b.as_ref() == Some(&a) {
                Some(b.clone())
            } else {
                parent_b.clone()
            },
        );

        // Update links to `b`
        if parent_b.as_ref() != Some(&a) {
            self.set_slot(slot_b, Some(a.clone()));
        }

        self.nodes.set_node_parent(
            b.clone(),
            if parent_a.as_ref() == Some(&b) {
                Some(a)
            } else {
                parent_a
            },
        );
    }

    /// Perform a left [rotation][1] at the given node.
    ///
    /// The node must have a right child.
    ///
    /// ```text
    ///       x              y
    ///      / \            / \
    ///     a   y    →    x   c
    ///        / \       / \
    ///       b   c     a   b
    /// ```
    ///
    /// [1]: https://en.wikipedia.org/wiki/Tree_rotation
    pub fn rotate_left(&mut self, x: Index) {
        let y = self
            .nodes
            .node_right(x.clone())
            .expect("`rotate_left` requires a right child");
        let b = self.nodes.node_left(y.clone());
        let parent_slot = self.tree.read(&self.nodes).parent_slot(x.clone());

        // Reattach `b`
        self.nodes.set_node_right(x.clone(), b.clone());
        if let Some(b) = b {
            self.nodes.set_node_parent(b, Some(x.clone()));
        }

        // Reattach `x`
        self.nodes
            .set_node_parent(y.clone(), parent_slot.clone().parent());
        self.set_slot(parent_slot, Some(y.clone()));

        // Reattach `y`
        self.nodes.set_node_left(y.clone(), Some(x.clone()));
        self.nodes.set_node_parent(x, Some(y));
    }

    /// Perform a right [rotation][1] at the given node.
    ///
    /// The node must have a left child.
    ///
    /// ```text
    ///       y            x
    ///      / \          / \
    ///     x   c   →    a   y
    ///    / \              / \
    ///   a   b            b   c
    /// ```
    ///
    /// [1]: https://en.wikipedia.org/wiki/Tree_rotation
    pub fn rotate_right(&mut self, y: Index) {
        let x = self
            .nodes
            .node_left(y.clone())
            .expect("`rotate_right` requires a left child");
        let b = self.nodes.node_right(x.clone());
        let parent_slot = self.tree.read(&self.nodes).parent_slot(y.clone());

        // Reattach `b`
        self.nodes.set_node_left(y.clone(), b.clone());
        if let Some(b) = b {
            self.nodes.set_node_parent(b, Some(y.clone()));
        }

        // Reattach `x`
        self.nodes
            .set_node_parent(x.clone(), parent_slot.clone().parent());
        self.set_slot(parent_slot, Some(x.clone()));

        // Reattach `y`
        self.nodes.set_node_right(x.clone(), Some(y.clone()));
        self.nodes.set_node_parent(y, Some(x));
    }
}

/// # BST Update Operations
///
/// The following methods maintain the binary search tree invariants in `self`.
impl<'tree, Nodes, Index> TreeAccessorMut<'tree, Nodes, Index>
where
    Nodes: NodesLinkMut<Index> + NodesCmp<Index>,
    Index: PartialEq + Clone,
{
    /// Insert an existing node in the pool into the position inside the BST
    /// determined based on its key.
    ///
    /// `new_node`'s link fields are overwritten.
    /// This means that, if `new_node` is already a member of another tree,
    /// the tree will become corrupted.
    ///
    /// If the slot is already filled by another node, it is replaced by
    /// `new_node` while preserving the tree structure.
    /// The old node is returned as `Some(existing_node)`.
    /// The contents of `existing_node`'s link fields are unspecified and must
    /// not be relied upon.
    pub fn bst_insert_node(&mut self, new_node: Index) -> Option<Index> {
        match bst_slot_to_insert(
            &self.nodes,
            |nodes, node| nodes.cmp_nodes(node, new_node.clone()),
            self.tree.root.clone(),
        ) {
            Ok(existing_node) => {
                self.replace_node(existing_node.clone(), new_node);
                Some(existing_node)
            }
            Err(slot) => {
                self.nodes.set_node_left(new_node.clone(), None);
                self.nodes.set_node_right(new_node.clone(), None);
                self.set_slot(slot.clone(), Some(new_node.clone()));
                self.nodes.set_node_parent(
                    new_node,
                    match slot {
                        Slot::Root => None,
                        Slot::Left(node) | Slot::Right(node) => Some(node),
                    },
                );
                None
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bintree::tests::{Node, any_bst, any_tree, tree_from_str};

    use proptest::prelude::*;
    use std::prelude::rust_2024::*;

    #[test]
    fn for_each_mut() {
        let (mut nodes, mut tree) = tree_from_str("(3 (1 _ (2 _ _)) (4 _ _))");
        _ = tree.write(&mut nodes[..]).for_each_mut(|_, value| {
            *value *= 10;
            core::ops::ControlFlow::Continue::<()>(())
        });
        let (nodes2, tree2) = tree_from_str("(30 (10 _ (20 _ _)) (40 _ _))");
        assert_eq!(nodes, nodes2);
        assert_eq!(tree, tree2);
    }

    #[proptest::property_test]
    fn pt_rotate_preserves_bst(
        #[strategy = any_bst()] (mut nodes, mut root): (Vec<Node<u8>>, Tree),
        ops: Vec<(prop::sample::Index, bool)>,
    ) {
        let expected_values: Vec<u8> = root.read(&nodes[..]).values().copied().collect();
        if expected_values.is_empty() {
            return Ok(());
        }

        for (index, right) in ops {
            let node_i = index.index(nodes.len());
            let node = &nodes[node_i];
            let right = match (&node.left, &node.right) {
                (None, None) => continue,
                (None, Some(_)) => false,
                (Some(_), None) => true,
                (Some(_), Some(_)) => right,
            };
            if right {
                root.write(&mut nodes[..]).rotate_right(node_i);
            } else {
                root.write(&mut nodes[..]).rotate_left(node_i);
            }
        }

        root.read(&nodes[..]).validate().unwrap();
        root.read(&nodes[..])
            .bst_validate(|nodes, i| nodes[i].value)
            .unwrap();

        let got_values: Vec<u8> = root.read(&nodes[..]).values().copied().collect();
        assert_eq!(got_values, expected_values);
    }

    #[proptest::property_test]
    fn pt_swap_nodes(
        #[strategy = any_tree()] (mut nodes, mut root): (Vec<Node<u8>>, Tree),
        target_nodes: [prop::sample::Index; 2],
    ) {
        if nodes.is_empty() {
            return Ok(());
        }

        let indices: Vec<usize> = root.read(&nodes[..]).indices().collect();
        let target_nodes = target_nodes.map(|i| i.index(indices.len()));

        // Calculate the expected set of final elements
        let mut expected_values: Vec<u8> = root.read(&nodes[..]).values().copied().collect();
        expected_values.swap(target_nodes[0], target_nodes[1]);

        root.write(&mut nodes[..])
            .swap_nodes(indices[target_nodes[0]], indices[target_nodes[1]]);

        // Check the result
        let got_values: Vec<u8> = root.read(&nodes[..]).values().copied().collect();
        assert_eq!(got_values, expected_values);
    }

    #[proptest::property_test]
    fn pt_bst_insert_node(
        #[strategy = any_bst()] (mut nodes, mut root): (Vec<Node<u8>>, Tree),
        new_value: u8,
    ) {
        // Calculate the expected set of final elements
        let mut expected_values: Vec<u8> = root.read(&nodes[..]).values().copied().collect();
        if let Err(i) = expected_values.binary_search(&new_value) {
            expected_values.insert(i, new_value);
        }

        // Insert a new node
        let i = nodes.len();
        nodes.push(Node {
            value: new_value,
            ..<_>::default()
        });
        let old_node = root.write(&mut nodes[..]).bst_insert_node(i);

        // Check the result
        let got_values: Vec<u8> = root.read(&nodes[..]).values().copied().collect();
        assert_eq!(got_values, expected_values);

        if let Some(i) = old_node {
            assert_eq!(nodes[i].value, new_value);
        }
    }
}