fix: add validation to NodeIndex constructor and remove BitIterator

src/merkle/index.rs

@@ -1,13 +1,23 @@
 use super::{Felt, MerkleError, RpoDigest, StarkField};
-use crate::bit::BitIterator;
 
 // NODE INDEX
 // ================================================================================================
 
-/// A Merkle tree address to an arbitrary node.
+/// Address to an arbitrary node in a binary tree using level order form.
 ///
-/// The position is relative to a tree in level order, where for a given depth `d` elements are
-/// numbered from $0..2^d$.
+/// The position is represented by the pair `(depth, pos)`, where for a given depth `d` elements
+/// are numbered from $0..(2^d)-1$. Example:
+///
+/// ```ignore
+/// depth
+/// 0             0
+/// 1         0        1
+/// 2      0    1    2    3
+/// 3     0 1  2 3  4 5  6 7
+/// ```
+///
+/// The root is represented by the pair $(0, 0)$, its left child is $(1, 0)$ and its right child
+/// $(1, 1)$.
 #[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
 pub struct NodeIndex {
     depth: u8,
@@ -19,20 +29,37 @@ impl NodeIndex {
     // --------------------------------------------------------------------------------------------
 
     /// Creates a new node index.
-    pub const fn new(depth: u8, value: u64) -> Self {
-        Self { depth, value }
+    ///
+    /// # Errors
+    /// Returns an error if the `value` is greater than or equal to 2^{depth}.
+    pub const fn new(depth: u8, value: u64) -> Result<Self, MerkleError> {
+        if (64 - value.leading_zeros()) > depth as u32 {
+            Err(MerkleError::InvalidIndex { depth, value })
+        } else {
+            Ok(Self { depth, value })
+        }
+    }
+
+    /// Creates a new node index for testing purposes.
+    ///
+    /// # Panics
+    /// Panics if the `value` is greater than or equal to 2^{depth}.
+    #[cfg(test)]
+    pub fn make(depth: u8, value: u64) -> Self {
+        Self::new(depth, value).unwrap()
     }
 
     /// Creates a node index from a pair of field elements representing the depth and value.
     ///
     /// # Errors
-    ///
-    /// Will error if the `u64` representation of the depth doesn't fit a `u8`.
+    /// Returns an error if:
+    /// - `depth` doesn't fit in a `u8`.
+    /// - `value` is greater than or equal to 2^{depth}.
     pub fn from_elements(depth: &Felt, value: &Felt) -> Result<Self, MerkleError> {
         let depth = depth.as_int();
         let depth = u8::try_from(depth).map_err(|_| MerkleError::DepthTooBig(depth))?;
         let value = value.as_int();
-        Ok(Self::new(depth, value))
+        Self::new(depth, value)
     }
 
     /// Creates a new node index pointing to the root of the tree.
@@ -40,12 +67,6 @@ impl NodeIndex {
         Self { depth: 0, value: 0 }
     }
 
-    /// Mutates the instance and returns it, replacing the depth.
-    pub const fn with_depth(mut self, depth: u8) -> Self {
-        self.depth = depth;
-        self
-    }
-
     /// Computes the value of the sibling of the current node.
     pub fn sibling(mut self) -> Self {
         self.value ^= 1;
@@ -83,11 +104,6 @@ impl NodeIndex {
         self.value
     }
 
-    /// Returns true if the current value fits the current depth for a binary tree.
-    pub const fn is_valid(&self) -> bool {
-        self.value < (1 << self.depth as u64)
-    }
-
     /// Returns true if the current instance points to a right sibling node.
     pub const fn is_value_odd(&self) -> bool {
         (self.value & 1) == 1
@@ -98,19 +114,6 @@ impl NodeIndex {
         self.depth == 0
     }
 
-    /// Returns a bit iterator for the `value`.
-    ///
-    /// Bits read from left-to-right represent which internal node's child should be visited to
-    /// arrive at the leaf. From the right-to-left the bit represent the position the hash of the
-    /// current element should go.
-    ///
-    /// Additionally, the value that is not visited are the sibling values necessary for a Merkle
-    /// opening.
-    pub fn bit_iterator(&self) -> BitIterator {
-        let depth: u32 = self.depth.into();
-        BitIterator::new(self.value).skip_back(u64::BITS - depth)
-    }
-
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 
@@ -127,14 +130,43 @@ mod tests {
     use super::*;
     use proptest::prelude::*;
 
+    #[test]
+    fn test_node_index_value_too_high() {
+        assert_eq!(
+            NodeIndex::new(0, 0).unwrap(),
+            NodeIndex { depth: 0, value: 0 }
+        );
+        match NodeIndex::new(0, 1) {
+            Err(MerkleError::InvalidIndex { depth, value }) => {
+                assert_eq!(depth, 0);
+                assert_eq!(value, 1);
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    #[test]
+    fn test_node_index_can_represent_depth_64() {
+        assert!(NodeIndex::new(64, u64::MAX).is_ok());
+    }
+
+    prop_compose! {
+        fn node_index()(value in 0..2u64.pow(u64::BITS - 1)) -> NodeIndex {
+            // unwrap never panics because the range of depth is 0..u64::BITS
+            let mut depth = value.ilog2() as u8;
+            if value > (1 << depth) { // round up
+                depth += 1;
+            }
+            NodeIndex::new(depth, value.into()).unwrap()
+        }
+    }
+
     proptest! {
         #[test]
         fn arbitrary_index_wont_panic_on_move_up(
-            depth in prop::num::u8::ANY,
-            value in prop::num::u64::ANY,
+            mut index in node_index(),
             count in prop::num::u8::ANY,
         ) {
-            let mut index = NodeIndex::new(depth, value);
             for _ in 0..count {
                 index.move_up();
             }