1//! Index for nodes of a binary tree based on an in-order tree walk.
  2//!
  3//! In-order walks have the parent node index split its left and right subtrees. All the left
  4//! children have indexes lower than the parent, meanwhile all the right subtree higher indexes.
  5//! This property makes it is easy to compute changes to the index by adding or subtracting the
  6//! leaves count.
  7use core::num::NonZeroUsize;
  8
  9/// Index of nodes in a perfectly balanced binary tree based on an in-order tree walk.
 10#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
 11pub struct InOrderIndex {
 12    idx: usize,
 13}
 14
 15impl InOrderIndex {
 16    /// Constructor for a new [InOrderIndex].
 17    pub fn new(idx: NonZeroUsize) -> InOrderIndex {
 18        InOrderIndex { idx: idx.get() }
 19    }
 20
 21    /// Constructs an index from a leaf position.
 22    ///
 23    /// Panics:
 24    ///
 25    /// If `leaf` is higher than or equal to `usize::MAX / 2`.
 26    pub fn from_leaf_pos(leaf: usize) -> InOrderIndex {
 27        // Convert the position from 0-indexed to 1-indexed, since the bit manipulation in this
 28        // implementation only works 1-indexed counting.
 29        let pos = leaf + 1;
 30        InOrderIndex { idx: pos * 2 - 1 }
 31    }
 32
 33    /// True if the index is pointing at a leaf.
 34    ///
 35    /// Every odd number represents a leaf.
 36    pub fn is_leaf(&self) -> bool {
 37        self.idx & 1 == 1
 38    }
 39
 40    /// Returns the level of the index.
 41    ///
 42    /// Starts at level zero for leaves and increases by one for each parent.
 43    pub fn level(&self) -> u32 {
 44        self.idx.trailing_zeros()
 45    }
 46
 47    /// Returns the index of the left child.
 48    ///
 49    /// Panics:
 50    ///
 51    /// If the index corresponds to a leaf.
 52    pub fn left_child(&self) -> InOrderIndex {
 53        // The left child is itself a parent, with an index that splits its left/right subtrees. To
 54        // go from the parent index to its left child, it is only necessary to subtract the count
 55        // of elements on the child's right subtree + 1.
 56        let els = 1 << (self.level() - 1);
 57        InOrderIndex { idx: self.idx - els }
 58    }
 59
 60    /// Returns the index of the right child.
 61    ///
 62    /// Panics:
 63    ///
 64    /// If the index corresponds to a leaf.
 65    pub fn right_child(&self) -> InOrderIndex {
 66        // To compute the index of the parent of the right subtree it is sufficient to add the size
 67        // of its left subtree + 1.
 68        let els = 1 << (self.level() - 1);
 69        InOrderIndex { idx: self.idx + els }
 70    }
 71
 72    /// Returns the index of the parent node.
 73    pub fn parent(&self) -> InOrderIndex {
 74        // If the current index corresponds to a node in a left tree, to go up a level it is
 75        // required to add the number of nodes of the right sibling, analogously if the node is a
 76        // right child, going up requires subtracting the number of nodes in its left subtree.
 77        //
 78        // Both of the above operations can be performed by bitwise manipulation. Below the mask
 79        // sets the number of trailing zeros to be equal the new level of the index, and the bit
 80        // marks the parent.
 81        let target = self.level() + 1;
 82        let bit = 1 << target;
 83        let mask = bit - 1;
 84        let idx = self.idx ^ (self.idx & mask);
 85        InOrderIndex { idx: idx | bit }
 86    }
 87
 88    /// Returns the index of the sibling node.
 89    pub fn sibling(&self) -> InOrderIndex {
 90        let parent = self.parent();
 91        if *self > parent {
 92            parent.left_child()
 93        } else {
 94            parent.right_child()
 95        }
 96    }
 97}
 98
 99#[cfg(test)]
100mod test {
101    use super::InOrderIndex;
102    use proptest::prelude::*;
103
104    proptest! {
105        #[test]
106        fn proptest_inorder_index_random(count in 1..1000usize) {
107            let left_pos = count * 2;
108            let right_pos = count * 2 + 1;
109
110            let left = InOrderIndex::from_leaf_pos(left_pos);
111            let right = InOrderIndex::from_leaf_pos(right_pos);
112
113            assert!(left.is_leaf());
114            assert!(right.is_leaf());
115            assert_eq!(left.parent(), right.parent());
116            assert_eq!(left.parent().right_child(), right);
117            assert_eq!(left, right.parent().left_child());
118            assert_eq!(left.sibling(), right);
119            assert_eq!(left, right.sibling());
120        }
121    }
122
123    #[test]
124    fn test_inorder_index_basic() {
125        let left = InOrderIndex::from_leaf_pos(0);
126        let right = InOrderIndex::from_leaf_pos(1);
127
128        assert!(left.is_leaf());
129        assert!(right.is_leaf());
130        assert_eq!(left.parent(), right.parent());
131        assert_eq!(left.parent().right_child(), right);
132        assert_eq!(left, right.parent().left_child());
133        assert_eq!(left.sibling(), right);
134        assert_eq!(left, right.sibling());
135    }
136}