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