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