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