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 sibling index of the current node.
 78    pub const fn sibling(mut self) -> Self {
 79        self.value ^= 1;
 80        self
 81    }
 82
 83    /// Returns left child index of the current node.
 84    pub const fn left_child(mut self) -> Self {
 85        self.depth += 1;
 86        self.value <<= 1;
 87        self
 88    }
 89
 90    /// Returns right child index of the current node.
 91    pub const fn right_child(mut self) -> Self {
 92        self.depth += 1;
 93        self.value = (self.value << 1) + 1;
 94        self
 95    }
 96
 97    // PROVIDERS
 98    // --------------------------------------------------------------------------------------------
 99
100    /// Builds a node to be used as input of a hash function when computing a Merkle path.
101    ///
102    /// Will evaluate the parity of the current instance to define the result.
103    pub const fn build_node(&self, slf: RpoDigest, sibling: RpoDigest) -> [RpoDigest; 2] {
104        if self.is_value_odd() {
105            [sibling, slf]
106        } else {
107            [slf, sibling]
108        }
109    }
110
111    /// Returns the scalar representation of the depth/value pair.
112    ///
113    /// It is computed as `2^depth + value`.
114    pub const fn to_scalar_index(&self) -> u64 {
115        (1 << self.depth as u64) + self.value
116    }
117
118    /// Returns the depth of the current instance.
119    pub const fn depth(&self) -> u8 {
120        self.depth
121    }
122
123    /// Returns the value of this index.
124    pub const fn value(&self) -> u64 {
125        self.value
126    }
127
128    /// Returns true if the current instance points to a right sibling node.
129    pub const fn is_value_odd(&self) -> bool {
130        (self.value & 1) == 1
131    }
132
133    /// Returns `true` if the depth is `0`.
134    pub const fn is_root(&self) -> bool {
135        self.depth == 0
136    }
137
138    // STATE MUTATORS
139    // --------------------------------------------------------------------------------------------
140
141    /// Traverses one level towards the root, decrementing the depth by `1`.
142    pub fn move_up(&mut self) {
143        self.depth = self.depth.saturating_sub(1);
144        self.value >>= 1;
145    }
146
147    /// Traverses towards the root until the specified depth is reached.
148    ///
149    /// Assumes that the specified depth is smaller than the current depth.
150    pub fn move_up_to(&mut self, depth: u8) {
151        debug_assert!(depth < self.depth);
152        let delta = self.depth.saturating_sub(depth);
153        self.depth = self.depth.saturating_sub(delta);
154        self.value >>= delta as u32;
155    }
156}
157
158impl Display for NodeIndex {
159    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
160        write!(f, "depth={}, value={}", self.depth, self.value)
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use proptest::prelude::*;
168
169    #[test]
170    fn test_node_index_value_too_high() {
171        assert_eq!(NodeIndex::new(0, 0).unwrap(), NodeIndex { depth: 0, value: 0 });
172        match NodeIndex::new(0, 1) {
173            Err(MerkleError::InvalidIndex { depth, value }) => {
174                assert_eq!(depth, 0);
175                assert_eq!(value, 1);
176            }
177            _ => unreachable!(),
178        }
179    }
180
181    #[test]
182    fn test_node_index_can_represent_depth_64() {
183        assert!(NodeIndex::new(64, u64::MAX).is_ok());
184    }
185
186    prop_compose! {
187        fn node_index()(value in 0..2u64.pow(u64::BITS - 1)) -> NodeIndex {
188            // unwrap never panics because the range of depth is 0..u64::BITS
189            let mut depth = value.ilog2() as u8;
190            if value > (1 << depth) { // round up
191                depth += 1;
192            }
193            NodeIndex::new(depth, value).unwrap()
194        }
195    }
196
197    proptest! {
198        #[test]
199        fn arbitrary_index_wont_panic_on_move_up(
200            mut index in node_index(),
201            count in prop::num::u8::ANY,
202        ) {
203            for _ in 0..count {
204                index.move_up();
205            }
206        }
207    }
208}