intrex 0.2.0

//! Mutable red-black tree accessor.
use core::{mem::swap, ops};

use crate::{
    bintree::{NodesCmp, NodesCmpKey, Slot},
    node_data::{NodesDataLendMut, NodesDataLendMutGat},
};

use super::*;

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

/// Accessor to a red-black tree.
#[derive(Debug)]
pub struct TreeAccessorMut<'head, Nodes, Index> {
    raw: crate::bintree::accessor_mut::TreeAccessorMut<'head, Nodes, Index>,
}

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

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

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

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

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

    /// 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.raw.front_mut()
    }

    /// 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.raw.back_mut()
    }

    /// Call the specified closure for every element, passing a mutable
    /// reference for each of them.
    #[inline]
    pub fn for_each_mut<B>(
        &mut self,
        f: impl FnMut(Index, <Nodes as NodesDataLendMutGat<'_, Index>>::Data) -> ops::ControlFlow<B>,
    ) -> ops::ControlFlow<B>
    where
        Nodes: NodesDataLendMut<Index>,
    {
        self.raw.for_each_mut(f)
    }

    /// Call the specified closure for every element in a given index 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>>,
        f: impl FnMut(Index, <Nodes as NodesDataLendMutGat<'_, Index>>::Data) -> ops::ControlFlow<B>,
    ) -> ops::ControlFlow<B>
    where
        Nodes: NodesDataLendMut<Index>,
    {
        self.raw.for_each_in_index_range_mut(range, f)
    }
}

/// # Red-Black Tree Operations
impl<'tree, Nodes, Index> TreeAccessorMut<'tree, Nodes, Index>
where
    Nodes: NodesRbMut<Index>,
    Index: PartialEq + Clone,
{
    /// Insert an existing node in the pool into the red-black tree at a
    /// specified position and rebalance it.
    ///
    /// `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 insert_node(&mut self, new_node: Index, slot: Slot<Index>) -> Option<Index> {
        let old_node = self.raw.tree().read(self.nodes()).slot(slot.clone());

        if let Some(old_node) = &old_node {
            // Replace an existing node while preserving all fields except
            // application data
            self.raw.replace_node(old_node.clone(), new_node.clone());
            let color = self.nodes().node_color(old_node.clone());
            self.nodes_mut().set_node_color(new_node.clone(), color);

            self.nodes_mut()
                .post_replace_node(old_node.clone(), new_node);
        } else {
            self.nodes_mut().set_node_left(new_node.clone(), None);
            self.nodes_mut().set_node_right(new_node.clone(), None);

            // Insert `new_node` to an empty place
            self.raw.set_slot(slot.clone(), Some(new_node.clone()));
            self.nodes_mut()
                .set_node_parent(new_node.clone(), slot.clone().parent());

            // Initially paint the new node red. This maintains the black height
            // invariant.
            self.nodes_mut()
                .set_node_color(new_node.clone(), Color::Red);

            self.nodes_mut().post_set_slot(new_node.clone());

            // Rebalance the tree
            self.rebalance_after_insert(new_node);
        }

        old_node
    }

    /// Rebalance the tree after a single node insertion.
    ///
    /// The node must be initially painted [`Color::Red`] by the caller.
    ///
    /// # Panics
    ///
    /// This method may panic if the tree invariants were already violated
    /// before the insertion.
    pub fn rebalance_after_insert(&mut self, new_node: Index) {
        let mut cur = new_node;

        loop {
            // `cur` is a red node potentially violating the red node
            // invariant.
            debug_assert_eq!(self.nodes().node_color(cur.clone()), Color::Red);

            let cur_slot = self.raw.tree().read(self.nodes()).parent_slot(cur.clone());

            let Some(mut parent) = cur_slot.clone().parent() else {
                break;
            };

            if self.nodes().node_color(parent.clone()) == Color::Black {
                // If `cur.parent` is black, then the red node invariant
                // is already maintained.
                break;
            }

            // `cur` is an illegal red child of `parent`.

            let parent_slot = self
                .raw
                .tree()
                .read(self.nodes())
                .parent_slot(parent.clone());

            let (uncle, grandparent, parent_side) = match parent_slot {
                Slot::Root => {
                    // If `cur.parent` is a red root, then we can simply
                    // repaint it black to restore the invariants.
                    self.nodes_mut().set_node_color(parent, Color::Black);
                    break;
                }
                Slot::Left(node) => (self.nodes().node_right(node.clone()), node, false),
                Slot::Right(node) => (self.nodes().node_left(node.clone()), node, true),
            };

            if let Some(uncle) = uncle.clone()
                && self.nodes().node_color(uncle.clone()) == Color::Red
            {
                //     gp(B)
                //     /   \
                //   u(R)   p(R)
                //         /   \
                //        x   c(R)
                //
                // Swap the colors of `grandparent`, `uncle`, and `parent`,
                // resolving the violation at `cur`. This, however, may create
                // a new violation at `grandparent`.
                self.nodes_mut().set_node_color(parent, Color::Black);
                self.nodes_mut().set_node_color(uncle, Color::Black);
                self.nodes_mut()
                    .set_node_color(grandparent.clone(), Color::Red);
                cur = grandparent;
                continue;
            }

            let cur_side = matches!(cur_slot, Slot::Right(_));
            if cur_side != parent_side {
                //     gp(B)
                //     /   \
                //   u(B)   p(R)
                //          /
                //        c(R)
                //
                // Apply tree rotation on `parent` to reduce this to the
                // case below.
                if parent_side {
                    self.nodes_mut().pre_rotate_right(parent.clone());
                    self.raw.rotate_right(parent.clone());
                } else {
                    self.nodes_mut().pre_rotate_left(parent.clone());
                    self.raw.rotate_left(parent.clone());
                }

                // `parent` is now the illegal red child of `cur`.
                swap(&mut cur, &mut parent);
                drop(cur_slot);
            }

            //     gp(B)
            //     /   \
            //   u(B)   p(R)
            //         /   \
            //        x   c(R)
            //
            // Apply tree rotation on `grandparent` to move `parent` into
            // its position.
            if parent_side {
                self.nodes_mut().pre_rotate_left(grandparent.clone());
                self.raw.rotate_left(grandparent.clone());
            } else {
                self.nodes_mut().pre_rotate_right(grandparent.clone());
                self.raw.rotate_right(grandparent.clone());
            }

            //       p(R)
            //       /   \
            //    gp(B)   c(R)
            //    /   \
            //  u(B)   x
            //
            // Now `cur` has one less black ancestor, violating the black
            // count invariant. Swapping the colors of `parent` and
            // `grandparent` increases the black height of `cur` while
            // maintaining those of `grandparent`'s descendants, restoring
            // the invariant. This also resolves the red node violation.
            self.nodes_mut().set_node_color(parent, Color::Black);
            self.nodes_mut().set_node_color(grandparent, Color::Red);
            break;
        } // loop
    } // fn

    /// Remove a node from the tree and rebalance it.
    pub fn remove_node(&mut self, node: Index) {
        let mut left = self.nodes().node_left(node.clone());
        let mut right = self.nodes().node_right(node.clone());

        if left.is_some() && right.is_some() {
            // Swap `node` with its successor and remove the successor instead.
            let next = self
                .raw
                .tree()
                .read(self.nodes())
                .next_index(node.clone())
                .unwrap();
            let node_color = self.nodes().node_color(node.clone());
            let next_color = self.nodes().node_color(next.clone());
            self.nodes_mut().pre_swap_nodes(node.clone(), next.clone());
            self.raw.swap_nodes(node.clone(), next.clone());
            self.nodes_mut().set_node_color(node.clone(), next_color);
            self.nodes_mut().set_node_color(next, node_color);

            left = None;
            debug_assert!(self.nodes().node_left(node.clone()).is_none());
            right = self.nodes().node_right(node.clone());
        }

        let node_slot = self.raw.tree().read(self.nodes()).parent_slot(node.clone());

        let left_some = left.is_some();

        match (left, right) {
            (None, None) => {
                self.nodes_mut()
                    .pre_clear_slot(node.clone(), node_slot.clone());
                self.raw.set_slot(node_slot.clone(), None);
                if !matches!(node_slot, Slot::Root)
                    && self.nodes().node_color(node.clone()) == Color::Black
                {
                    // This created an imbalance, which must be fixed.
                    self.rebalance_after_remove_black(node_slot);
                }
            }
            (None, Some(child)) | (Some(child), None) => {
                // Replace `node` with `child` and repaint it as black.
                //
                // Using a tree rotation is a bit inefficient but reduces the
                // number of reshaping hooks that need to be implemented by
                // users.
                debug_assert_eq!(self.nodes().node_color(node.clone()), Color::Black);
                debug_assert_eq!(self.nodes().node_color(child.clone()), Color::Red);
                self.nodes_mut().set_node_color(child.clone(), Color::Black);
                if left_some {
                    self.nodes_mut().pre_rotate_right(node.clone());
                    self.raw.rotate_right(node.clone());
                    self.nodes_mut()
                        .pre_clear_slot(node, Slot::Right(child.clone()));
                    self.nodes_mut().set_node_right(child, None);
                } else {
                    self.nodes_mut().pre_rotate_left(node.clone());
                    self.raw.rotate_left(node.clone());
                    self.nodes_mut()
                        .pre_clear_slot(node, Slot::Left(child.clone()));
                    self.nodes_mut().set_node_left(child, None);
                }
            }
            (Some(_), Some(_)) => {
                unreachable!()
            }
        }
    }

    /// Rebalance the tree after a single black height reduction.
    ///
    /// `cur_slot` specifies the subtree that has its black node counts reduced
    /// by one, violating the invariant.
    ///
    /// # Panics
    ///
    /// This method may panic if the tree invariants were already violated
    /// before the black height reduction.
    pub fn rebalance_after_remove_black(&mut self, mut cur_slot: Slot<Index>) {
        macro_rules! color {
            ($i:expr) => {
                Option::map_or($i, Color::Black, |i| self.nodes().node_color(i))
            };
        }

        loop {
            let (mut close_nibling, mut distant_nibling, parent, mut sibling, cur_side) =
                match cur_slot {
                    Slot::Root => {
                        // Decreasing the whole tree's black height does not
                        // violate the invariant.
                        break;
                    }
                    Slot::Left(parent) => {
                        let sibling = self.nodes().node_right(parent.clone()).expect(
                            "sibling is null - tree invariants were \
                            already violated before black height reduction",
                        );
                        (
                            self.nodes().node_left(sibling.clone()),
                            self.nodes().node_right(sibling.clone()),
                            parent,
                            sibling,
                            false,
                        )
                    }
                    Slot::Right(parent) => {
                        //         p(_)
                        //         /   \
                        //      s(_)   c(*)
                        //      /   \
                        //   dn(_)  cn(_)
                        let sibling = self.nodes().node_left(parent.clone()).expect(
                            "sibling is null - tree invariants were \
                            already violated before black height reduction",
                        );
                        (
                            self.nodes().node_right(sibling.clone()),
                            self.nodes().node_left(sibling.clone()),
                            parent,
                            sibling,
                            true,
                        )
                    }
                };

            if color!(Some(sibling.clone())) == Color::Red {
                //         p(B)
                //         /   \
                //      s(R)   c(*)
                //      /   \
                //   dn(B)  cn(B)
                debug_assert_eq!(color!(Some(parent.clone())), Color::Black);
                debug_assert_eq!(color!(close_nibling.clone()), Color::Black);
                debug_assert_eq!(color!(distant_nibling.clone()), Color::Black);

                // Apply tree rotation on `parent`, moving `sibling` into its
                // position.
                if cur_side {
                    self.nodes_mut().pre_rotate_right(parent.clone());
                    self.raw.rotate_right(parent.clone());
                } else {
                    self.nodes_mut().pre_rotate_left(parent.clone());
                    self.raw.rotate_left(parent.clone());
                }

                // Flip the colors of `parent` and `sibling`.
                self.nodes_mut().set_node_color(parent.clone(), Color::Red);
                self.nodes_mut().set_node_color(sibling, Color::Black);

                //          s(B)
                //         /   \
                //     dn(B)   p(R)
                //            /    \
                //          cn(B)  c(*)
                //
                // `cur` is still one black node short, but it now has a red
                // parent and a black sibling, which can be handled by the
                // remaining cases.
                sibling = close_nibling.expect(
                    "`close_nibling` is null in red sibling case - tree \
                    invariants were already violated before black height \
                    reduction",
                );
                if cur_side {
                    close_nibling = self.nodes().node_right(sibling.clone());
                    distant_nibling = self.nodes().node_left(sibling.clone());
                } else {
                    close_nibling = self.nodes().node_left(sibling.clone());
                    distant_nibling = self.nodes().node_right(sibling.clone());
                }
            }

            debug_assert_eq!(color!(Some(sibling.clone())), Color::Black);

            let distant_nibling = if let Some(distant_nibling) = distant_nibling
                && color!(Some(distant_nibling.clone())) == Color::Red
            {
                distant_nibling
            } else if let Some(close_nibling) = close_nibling.clone()
                && color!(Some(close_nibling.clone())) == Color::Red
            {
                //         p(_)
                //         /   \
                //      s(B)   c(*)
                //      /   \
                //   dn(B)  cn(R)
                //
                // Apply tree rotation on `sibling`, moving `close_nibling`
                // into its position.
                if cur_side {
                    self.nodes_mut().pre_rotate_left(sibling.clone());
                    self.raw.rotate_left(sibling.clone());
                } else {
                    self.nodes_mut().pre_rotate_right(sibling.clone());
                    self.raw.rotate_right(sibling.clone());
                }

                //            p(_)
                //            /   \
                //        cn(R)   c(*)
                //         /
                //      s(B)
                //      /
                //   dn(B)
                //
                // Flip the colors of `sibling` and `close_nibling` to
                // preserve invariants.
                self.nodes_mut()
                    .set_node_color(close_nibling.clone(), Color::Black);
                self.nodes_mut().set_node_color(sibling.clone(), Color::Red);

                let distant_nibling = sibling;
                sibling = close_nibling;

                distant_nibling
            } else if color!(Some(parent.clone())) == Color::Red {
                //         p(R)
                //         /   \
                //      s(B)   c(*)
                //      /   \
                //   dn(B)  cn(B)
                //
                // Flip the colors of `sibling` and `parent`. This increases
                // the black height of `cur`, restoring the invariant.
                self.nodes_mut().set_node_color(parent, Color::Black);
                self.nodes_mut().set_node_color(sibling, Color::Red);
                break;
            } else {
                //         p(B)
                //         /   \
                //      s(B)   c(*)
                //      /   \
                //   dn(B)  cn(B)
                //
                // Repaint `sibling` red, removing the imbalance between
                // `sibling` and `close_nibling`.
                self.nodes_mut().set_node_color(sibling, Color::Red);

                // `parent` is now one black short compared to the rest of
                // the tree, so rebalancing must continue here.
                cur_slot = self.raw.tree().read(self.nodes()).parent_slot(parent);
                continue;
            };

            //         p(_)
            //         /   \
            //      s(B)   c(*)
            //      /   \
            //   dn(R)  cn(_)
            //
            // Apply tree rotation on `parent`, moving `sibling`
            // into its position.
            if cur_side {
                self.nodes_mut().pre_rotate_right(parent.clone());
                self.raw.rotate_right(parent.clone());
            } else {
                self.nodes_mut().pre_rotate_left(parent.clone());
                self.raw.rotate_left(parent.clone());
            }

            // Swap the colors of `sibling` and `parent`.
            let parent_color = color!(Some(parent.clone()));
            self.nodes_mut()
                .set_node_color(sibling.clone(), parent_color);
            self.nodes_mut()
                .set_node_color(parent.clone(), Color::Black);

            // Recolor `distant_nibling` black.
            self.nodes_mut()
                .set_node_color(distant_nibling, Color::Black);

            //         s(_)
            //         /   \
            //     dn(B)   p(B)
            //            /    \
            //         cn(_)   c(*)
            //
            // The tree is now balanced.
            break;
        } // loop
    } // fn
}

/// # Red-Black Search Tree Operations
impl<'tree, Nodes, Index> TreeAccessorMut<'tree, Nodes, Index>
where
    Nodes: NodesRbMut<Index>,
    Index: PartialEq + Clone,
{
    /// Insert an existing node in the pool into the red-black search tree and
    /// rebalance it.
    ///
    /// `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 tree already contains a node with the same key, 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>
    where
        Nodes: NodesCmp<Index>,
    {
        // Insert `new_node` without rebalancing
        if let Some(old_node) = self.raw.bst_insert_node(new_node.clone()) {
            // `new_node` replaced an existing node - no rebalancing required
            let color = self.nodes().node_color(old_node.clone());
            self.nodes_mut().set_node_color(new_node.clone(), color);

            self.nodes_mut()
                .post_replace_node(old_node.clone(), new_node);

            Some(old_node)
        } else {
            self.nodes_mut().post_set_slot(new_node.clone());

            // Initially paint the new node red. This maintains the black height
            // invariant.
            self.nodes_mut()
                .set_node_color(new_node.clone(), Color::Red);

            // Rebalance the tree
            self.rebalance_after_insert(new_node);
            None
        }
    }

    /// Remove the node matching `key` from the red-black search tree and
    /// rebalance it.
    ///
    /// If the tree contains a node with the given key, its index is returned.
    /// If no matching node is found, `None` is returned, and the tree is left
    /// unmodified.
    pub fn bst_remove_node<Key>(&mut self, key: &Key) -> Option<Index>
    where
        Nodes: NodesCmpKey<Index, Key>,
    {
        let node = self.raw.tree().read(self.nodes()).bst_find_index(key)?;
        self.remove_node(node.clone());
        Some(node)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::rbtree::tests::{Node, any_rbbst};
    use proptest::prelude::*;
    use proptest_state_machine::{ReferenceStateMachine, StateMachineTest, prop_state_machine};
    use std::{collections::BTreeSet, prelude::rust_2024::*, vec};

    #[proptest::property_test]
    fn pt_insert_node(
        #[strategy = any_rbbst(true)] (mut nodes, mut root): (Vec<Node<u8>>, Tree),
        slot: prop::sample::Index,
    ) {
        // Pick a slot to insert a new node
        let free_slots = if root.is_empty() {
            vec![Slot::Root]
        } else {
            root.read(&nodes[..])
                .indices()
                .flat_map(|i| {
                    let node = &nodes[i];
                    [
                        node.left.is_none().then_some(Slot::Left(i)),
                        node.right.is_none().then_some(Slot::Right(i)),
                    ]
                })
                .flatten()
                .collect()
        };

        let slot = *slot.get(&free_slots);

        // Calculate the expected resulting sequence
        let mut expected_values: Vec<u8> = root.read(&nodes[..]).values().copied().collect();

        let value = 0;

        let i = match slot {
            Slot::Root => 0,
            Slot::Left(parent) => expected_values
                .iter()
                .position(|x| *x == nodes[parent].value)
                .unwrap(),
            Slot::Right(parent) => {
                expected_values
                    .iter()
                    .position(|x| *x == nodes[parent].value)
                    .unwrap()
                    + 1
            }
        };
        expected_values.insert(i, value);

        // Insert a node
        let new_node_i = nodes.len();
        nodes.push(Node {
            value,
            ..<_>::default()
        });
        assert!(
            root.write(&mut nodes[..])
                .insert_node(new_node_i, slot)
                .is_none()
        );
        root.read(&nodes[..]).validate().unwrap_or_else(|e| {
            panic!("validation failed: {e}\n\nroot = {root:?}\nnodes = {nodes:?}")
        });

        // Check the resulting sequence
        let got_values: Vec<u8> = root.read(&nodes[..]).values().copied().collect();

        assert_eq!(got_values, expected_values);
    }

    #[proptest::property_test]
    fn pt_remove_node(
        #[strategy = any_rbbst(false)] (mut nodes, mut root): (Vec<Node<u8>>, Tree),
        node: prop::sample::Index,
    ) {
        // Pick a node to remove
        let node_i = node.index(nodes.len());

        // Calculate the expected resulting sequence
        let mut expected_values: Vec<u8> = root.read(&nodes[..]).values().copied().collect();
        let node_value = nodes[node_i].value;
        expected_values.retain(|x| *x != node_value);

        // Remove the node
        root.write(&mut nodes[..]).remove_node(node_i);
        root.read(&nodes[..]).validate().unwrap_or_else(|e| {
            panic!("validation failed: {e}\n\nroot = {root:?}\nnodes = {nodes:?}")
        });
        root.read(&nodes[..])
            .bst_validate(|nodes, i| nodes[i].value)
            .unwrap_or_else(|e| {
                panic!("validation failed: {e}\n\nroot = {root:?}\nnodes = {nodes:?}")
            });

        // Check the resulting sequence
        let got_values: Vec<u8> = root.read(&nodes[..]).values().copied().collect();

        assert_eq!(got_values, expected_values);
    }

    /// A reference state machine representing a set.
    ///
    /// Used for state machine testing with [`proptest_state_machine`].
    struct SetStateMachine;

    #[derive(Debug, Clone)]
    enum SetTransition {
        Insert(u8),
        Remove(u8),
    }

    impl ReferenceStateMachine for SetStateMachine {
        type State = BTreeSet<u8>;
        type Transition = SetTransition;

        fn init_state() -> BoxedStrategy<Self::State> {
            Just(BTreeSet::new()).boxed()
        }

        fn transitions(state: &Self::State) -> BoxedStrategy<Self::Transition> {
            let insert = any::<u8>().prop_map(SetTransition::Insert);
            let remove_bad = any::<u8>().prop_map(SetTransition::Remove);
            if state.is_empty() {
                prop_oneof![insert, remove_bad].boxed()
            } else {
                let remove = prop::sample::select(state.iter().copied().collect::<Vec<_>>())
                    .prop_map(SetTransition::Remove);
                prop_oneof![insert, remove, remove_bad].boxed()
            }
        }

        fn apply(mut state: Self::State, transition: &Self::Transition) -> Self::State {
            match *transition {
                SetTransition::Insert(x) => {
                    state.insert(x);
                }
                SetTransition::Remove(x) => {
                    state.remove(&x);
                }
            }

            state
        }
    }

    /// Defines the system-under-test used by the test case
    /// [`pt_bst_insert_and_remove_nodes`].
    struct PtBstInsertAndRemoveNodes;

    struct PtBstInsertAndRemoveNodesSystem {
        nodes: Vec<Node<u8>>,
        tree: Tree,
        free_nodes: Vec<usize>,
    }

    impl StateMachineTest for PtBstInsertAndRemoveNodes {
        type SystemUnderTest = PtBstInsertAndRemoveNodesSystem;
        type Reference = SetStateMachine;

        fn init_test(
            ref_state: &<Self::Reference as ReferenceStateMachine>::State,
        ) -> Self::SystemUnderTest {
            assert!(ref_state.is_empty());
            PtBstInsertAndRemoveNodesSystem {
                nodes: Vec::default(),
                tree: Tree::default(),
                free_nodes: Vec::new(),
            }
        }

        fn apply(
            mut state: Self::SystemUnderTest,
            _ref_state: &<Self::Reference as ReferenceStateMachine>::State,
            transition: <Self::Reference as ReferenceStateMachine>::Transition,
        ) -> Self::SystemUnderTest {
            match transition {
                SetTransition::Insert(value) => {
                    let node = Node {
                        value,
                        ..<_>::default()
                    };
                    let node_i = if let Some(i) = state.free_nodes.pop() {
                        state.nodes[i] = node;
                        i
                    } else {
                        let i = state.nodes.len();
                        state.nodes.push(node);
                        i
                    };
                    if let Some(old_node_i) = state
                        .tree
                        .write(&mut state.nodes[..])
                        .bst_insert_node(node_i)
                    {
                        state.free_nodes.push(old_node_i);
                    }
                }
                SetTransition::Remove(value) => {
                    if let Some(old_node_i) = state
                        .tree
                        .write(&mut state.nodes[..])
                        .bst_remove_node(&value)
                    {
                        state.free_nodes.push(old_node_i);
                    }
                }
            }
            state
        }

        fn check_invariants(
            state: &Self::SystemUnderTest,
            ref_state: &<Self::Reference as ReferenceStateMachine>::State,
        ) {
            state
                .tree
                .read(&state.nodes[..])
                .validate()
                .unwrap_or_else(|e| panic!("validation failed: {e}"));
            state
                .tree
                .read(&state.nodes[..])
                .bst_validate(|nodes, i| nodes[i].value)
                .unwrap_or_else(|e| panic!("validation failed: {e}"));

            // The sets must be equal
            let got_values = state.tree.read(&state.nodes[..]).into_values().copied();
            let ref_values = ref_state.iter().copied();
            assert!(got_values.eq(ref_values));
        }
    }

    prop_state_machine! {
        #[test]
        fn pt_bst_insert_and_remove_nodes(sequential 1..20
            => PtBstInsertAndRemoveNodes);
    }
}