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