1use super::{Felt, MerkleError, RpoDigest, StarkField};
  2
  3// NODE INDEX
  4// ================================================================================================
  5
  6/// A Merkle tree address to an arbitrary node.
  7#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
  8pub struct NodeIndex {
  9    depth: u8,
 10    value: u64,
 11}
 12
 13impl NodeIndex {
 14    // CONSTRUCTORS
 15    // --------------------------------------------------------------------------------------------
 16
 17    /// Creates a new node index.
 18    pub const fn new(depth: u8, value: u64) -> Self {
 19        Self { depth, value }
 20    }
 21
 22    /// Creates a node index from a pair of field elements representing the depth and value.
 23    ///
 24    /// # Errors
 25    ///
 26    /// Will error if the `u64` representation of the depth doesn't fit a `u8`.
 27    pub fn from_elements(depth: &Felt, value: &Felt) -> Result<Self, MerkleError> {
 28        let depth = depth.as_int();
 29        let depth = u8::try_from(depth).map_err(|_| MerkleError::DepthTooBig(depth))?;
 30        let value = value.as_int();
 31        Ok(Self::new(depth, value))
 32    }
 33
 34    /// Creates a new node index pointing to the root of the tree.
 35    pub const fn root() -> Self {
 36        Self { depth: 0, value: 0 }
 37    }
 38
 39    /// Mutates the instance and returns it, replacing the depth.
 40    pub const fn with_depth(mut self, depth: u8) -> Self {
 41        self.depth = depth;
 42        self
 43    }
 44
 45    /// Computes the value of the sibling of the current node.
 46    pub fn sibling(mut self) -> Self {
 47        self.value ^= 1;
 48        self
 49    }
 50
 51    // PROVIDERS
 52    // --------------------------------------------------------------------------------------------
 53
 54    /// Builds a node to be used as input of a hash function when computing a Merkle path.
 55    ///
 56    /// Will evaluate the parity of the current instance to define the result.
 57    pub const fn build_node(&self, slf: RpoDigest, sibling: RpoDigest) -> [RpoDigest; 2] {
 58        if self.is_value_odd() {
 59            [sibling, slf]
 60        } else {
 61            [slf, sibling]
 62        }
 63    }
 64
 65    /// Returns the scalar representation of the depth/value pair.
 66    ///
 67    /// It is computed as `2^depth + value`.
 68    pub const fn to_scalar_index(&self) -> u64 {
 69        (1 << self.depth as u64) + self.value
 70    }
 71
 72    /// Returns the depth of the current instance.
 73    pub const fn depth(&self) -> u8 {
 74        self.depth
 75    }
 76
 77    /// Returns the value of the current depth.
 78    pub const fn value(&self) -> u64 {
 79        self.value
 80    }
 81
 82    /// Returns true if the current value fits the current depth for a binary tree.
 83    pub const fn is_valid(&self) -> bool {
 84        self.value < (1 << self.depth as u64)
 85    }
 86
 87    /// Returns true if the current instance points to a right sibling node.
 88    pub const fn is_value_odd(&self) -> bool {
 89        (self.value & 1) == 1
 90    }
 91
 92    /// Returns `true` if the depth is `0`.
 93    pub const fn is_root(&self) -> bool {
 94        self.depth == 0
 95    }
 96
 97    // STATE MUTATORS
 98    // --------------------------------------------------------------------------------------------
 99
100    /// Traverse one level towards the root, decrementing the depth by `1`.
101    pub fn move_up(&mut self) -> &mut Self {
102        self.depth = self.depth.saturating_sub(1);
103        self.value >>= 1;
104        self
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use proptest::prelude::*;
112
113    proptest! {
114        #[test]
115        fn arbitrary_index_wont_panic_on_move_up(
116            depth in prop::num::u8::ANY,
117            value in prop::num::u64::ANY,
118            count in prop::num::u8::ANY,
119        ) {
120            let mut index = NodeIndex::new(depth, value);
121            for _ in 0..count {
122                index.move_up();
123            }
124        }
125    }
126}