1use super::{Felt, MerkleError, RpoDigest, StarkField};
  2
  3// NODE INDEX
  4// ================================================================================================
  5
  6/// Address to an arbitrary node in a binary tree using level order form.
  7///
  8/// The position is represented by the pair `(depth, pos)`, where for a given depth `d` elements
  9/// are numbered from $0..(2^d)-1$. Example:
 10///
 11/// ```ignore
 12/// depth
 13/// 0             0
 14/// 1         0        1
 15/// 2      0    1    2    3
 16/// 3     0 1  2 3  4 5  6 7
 17/// ```
 18///
 19/// The root is represented by the pair $(0, 0)$, its left child is $(1, 0)$ and its right child
 20/// $(1, 1)$.
 21#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
 22pub struct NodeIndex {
 23    depth: u8,
 24    value: u64,
 25}
 26
 27impl NodeIndex {
 28    // CONSTRUCTORS
 29    // --------------------------------------------------------------------------------------------
 30
 31    /// Creates a new node index.
 32    ///
 33    /// # Errors
 34    /// Returns an error if the `value` is greater than or equal to 2^{depth}.
 35    pub const fn new(depth: u8, value: u64) -> Result<Self, MerkleError> {
 36        if (64 - value.leading_zeros()) > depth as u32 {
 37            Err(MerkleError::InvalidIndex { depth, value })
 38        } else {
 39            Ok(Self { depth, value })
 40        }
 41    }
 42
 43    /// Creates a new node index for testing purposes.
 44    ///
 45    /// # Panics
 46    /// Panics if the `value` is greater than or equal to 2^{depth}.
 47    #[cfg(test)]
 48    pub fn make(depth: u8, value: u64) -> Self {
 49        Self::new(depth, value).unwrap()
 50    }
 51
 52    /// Creates a node index from a pair of field elements representing the depth and value.
 53    ///
 54    /// # Errors
 55    /// Returns an error if:
 56    /// - `depth` doesn't fit in a `u8`.
 57    /// - `value` is greater than or equal to 2^{depth}.
 58    pub fn from_elements(depth: &Felt, value: &Felt) -> Result<Self, MerkleError> {
 59        let depth = depth.as_int();
 60        let depth = u8::try_from(depth).map_err(|_| MerkleError::DepthTooBig(depth))?;
 61        let value = value.as_int();
 62        Self::new(depth, value)
 63    }
 64
 65    /// Creates a new node index pointing to the root of the tree.
 66    pub const fn root() -> Self {
 67        Self { depth: 0, value: 0 }
 68    }
 69
 70    /// Computes the value of the sibling of the current node.
 71    pub fn sibling(mut self) -> Self {
 72        self.value ^= 1;
 73        self
 74    }
 75
 76    // PROVIDERS
 77    // --------------------------------------------------------------------------------------------
 78
 79    /// Builds a node to be used as input of a hash function when computing a Merkle path.
 80    ///
 81    /// Will evaluate the parity of the current instance to define the result.
 82    pub const fn build_node(&self, slf: RpoDigest, sibling: RpoDigest) -> [RpoDigest; 2] {
 83        if self.is_value_odd() {
 84            [sibling, slf]
 85        } else {
 86            [slf, sibling]
 87        }
 88    }
 89
 90    /// Returns the scalar representation of the depth/value pair.
 91    ///
 92    /// It is computed as `2^depth + value`.
 93    pub const fn to_scalar_index(&self) -> u64 {
 94        (1 << self.depth as u64) + self.value
 95    }
 96
 97    /// Returns the depth of the current instance.
 98    pub const fn depth(&self) -> u8 {
 99        self.depth
100    }
101
102    /// Returns the value of this index.
103    pub const fn value(&self) -> u64 {
104        self.value
105    }
106
107    /// Returns true if the current instance points to a right sibling node.
108    pub const fn is_value_odd(&self) -> bool {
109        (self.value & 1) == 1
110    }
111
112    /// Returns `true` if the depth is `0`.
113    pub const fn is_root(&self) -> bool {
114        self.depth == 0
115    }
116
117    // STATE MUTATORS
118    // --------------------------------------------------------------------------------------------
119
120    /// Traverse one level towards the root, decrementing the depth by `1`.
121    pub fn move_up(&mut self) -> &mut Self {
122        self.depth = self.depth.saturating_sub(1);
123        self.value >>= 1;
124        self
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use proptest::prelude::*;
132
133    #[test]
134    fn test_node_index_value_too_high() {
135        assert_eq!(
136            NodeIndex::new(0, 0).unwrap(),
137            NodeIndex { depth: 0, value: 0 }
138        );
139        match NodeIndex::new(0, 1) {
140            Err(MerkleError::InvalidIndex { depth, value }) => {
141                assert_eq!(depth, 0);
142                assert_eq!(value, 1);
143            }
144            _ => unreachable!(),
145        }
146    }
147
148    #[test]
149    fn test_node_index_can_represent_depth_64() {
150        assert!(NodeIndex::new(64, u64::MAX).is_ok());
151    }
152
153    prop_compose! {
154        fn node_index()(value in 0..2u64.pow(u64::BITS - 1)) -> NodeIndex {
155            // unwrap never panics because the range of depth is 0..u64::BITS
156            let mut depth = value.ilog2() as u8;
157            if value > (1 << depth) { // round up
158                depth += 1;
159            }
160            NodeIndex::new(depth, value.into()).unwrap()
161        }
162    }
163
164    proptest! {
165        #[test]
166        fn arbitrary_index_wont_panic_on_move_up(
167            mut index in node_index(),
168            count in prop::num::u8::ANY,
169        ) {
170            for _ in 0..count {
171                index.move_up();
172            }
173        }
174    }
175}