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