feat: Add StorageIndex newtype

crates/merkle-tree/src/merkle_node.rs

@@ -12,6 +12,7 @@ use bitvec::prelude::BitVec;
 use bitvec::slice::BitSlice;
 use pathfinder_common::hash::FeltHash;
 use pathfinder_crypto::Felt;
+use pathfinder_storage::Storage;
 
 /// A node in a Binary Merkle-Patricia Tree graph.
 #[derive(Clone, Debug, PartialEq)]
@@ -28,11 +29,27 @@ pub enum InternalNode {
     Leaf,
 }
 
+/// A newtype for the storage index of a trie node.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct StorageIndex(u64);
+
+impl StorageIndex {
+    /// Create a new StorageIndex.
+    pub fn new(index: u64) -> Self {
+        Self(index)
+    }
+
+    /// Get the inner u64 value.
+    pub fn value(&self) -> u64 {
+        self.0
+    }
+}
+
 /// Describes the [InternalNode::Binary] variant.
 #[derive(Clone, Debug, PartialEq)]
 pub struct BinaryNode {
     /// The storage index of this node (if it was loaded from storage).
-    pub storage_index: Option<u64>,
+    pub storage_index: Option<StorageIndex>,
     /// The height of this node in the tree.
     pub height: usize,
     /// [Left](Direction::Left) child.
@@ -44,7 +61,7 @@ pub struct BinaryNode {
 #[derive(Clone, Debug, PartialEq)]
 pub struct EdgeNode {
     /// The storage index of this node (if it was loaded from storage).
-    pub storage_index: Option<u64>,
+    pub storage_index: Option<StorageIndex>,
     /// The starting height of this node in the tree.
     pub height: usize,
     /// The path this edge takes.
@@ -144,9 +161,9 @@ impl InternalNode {
         matches!(self, InternalNode::Leaf)
     }
 
-    pub fn storage_index(&self) -> Option<u64> {
+    pub fn storage_index(&self) -> Option<StorageIndex> {
         match self {
-            InternalNode::Unresolved(storage_index) => Some(*storage_index),
+            InternalNode::Unresolved(storage_index) => Some(StorageIndex::new(*storage_index)),
             InternalNode::Binary(binary) => binary.storage_index,
             InternalNode::Edge(edge) => edge.storage_index,
             InternalNode::Leaf => None,

crates/merkle-tree/src/tree.rs

@@ -63,7 +63,7 @@ use pathfinder_common::trie::TrieNode;
 use pathfinder_crypto::Felt;
 use pathfinder_storage::{Node, NodeRef, StoredNode, TrieUpdate};
 
-use crate::merkle_node::{BinaryNode, Direction, EdgeNode, InternalNode};
+use crate::merkle_node::{BinaryNode, Direction, EdgeNode, InternalNode, StorageIndex};
 use crate::storage::Storage;
 
 /// A Starknet binary Merkle-Patricia tree.
@@ -214,7 +214,7 @@ impl<H: FeltHash, const HEIGHT: usize> MerkleTree<H, HEIGHT> {
                 };
 
                 if let Some(storage_index) = binary.storage_index {
-                    removed.push(storage_index);
+                    removed.push(storage_index.value());
                 };
 
                 let node_index = added.len();
@@ -247,7 +247,7 @@ impl<H: FeltHash, const HEIGHT: usize> MerkleTree<H, HEIGHT> {
                 let node_index = added.len();
                 added.push((hash, persisted_node));
                 if let Some(storage_index) = edge.storage_index {
-                    removed.push(storage_index);
+                    removed.push(storage_index.value());
                 };
 
                 (hash, Some(NodeRef::Index(node_index)))
@@ -371,7 +371,7 @@ impl<H: FeltHash, const HEIGHT: usize> MerkleTree<H, HEIGHT> {
 
                 let old_node = node.replace(updated);
                 if let Some(index) = old_node.storage_index() {
-                    self.nodes_removed.push(index);
+                    self.nodes_removed.push(index.value());
                 };
             }
             None => {
@@ -438,7 +438,7 @@ impl<H: FeltHash, const HEIGHT: usize> MerkleTree<H, HEIGHT> {
                 InternalNode::Binary(_) => false,
                 _ => {
                     if let Some(index) = node.storage_index() {
-                        indexes_removed.push(index);
+                        indexes_removed.push(index.value());
                     };
                     true
                 }
@@ -471,7 +471,7 @@ impl<H: FeltHash, const HEIGHT: usize> MerkleTree<H, HEIGHT> {
                 // Replace the old binary node with the new edge node.
                 let old_node = node.replace(InternalNode::Edge(new_edge));
                 if let Some(index) = old_node.storage_index() {
-                    self.nodes_removed.push(index);
+                    self.nodes_removed.push(index.value());
                 };
             }
             None => {
@@ -747,25 +747,25 @@ impl<H: FeltHash, const HEIGHT: usize> MerkleTree<H, HEIGHT> {
 
         let node = match node {
             StoredNode::Binary { left, right } => InternalNode::Binary(BinaryNode {
-                storage_index: Some(index),
+                storage_index: Some(StorageIndex::new(index)),
                 height,
                 left: Rc::new(RefCell::new(InternalNode::Unresolved(left))),
                 right: Rc::new(RefCell::new(InternalNode::Unresolved(right))),
             }),
             StoredNode::Edge { child, path } => InternalNode::Edge(EdgeNode {
-                storage_index: Some(index),
+                storage_index: Some(StorageIndex::new(index)),
                 height,
                 path,
                 child: Rc::new(RefCell::new(InternalNode::Unresolved(child))),
             }),
             StoredNode::LeafBinary => InternalNode::Binary(BinaryNode {
-                storage_index: Some(index),
+                storage_index: Some(StorageIndex::new(index)),
                 height,
                 left: Rc::new(RefCell::new(InternalNode::Leaf)),
                 right: Rc::new(RefCell::new(InternalNode::Leaf)),
             }),
             StoredNode::LeafEdge { path } => InternalNode::Edge(EdgeNode {
-                storage_index: Some(index),
+                storage_index: Some(StorageIndex::new(index)),
                 height,
                 path,
                 child: Rc::new(RefCell::new(InternalNode::Leaf)),
@@ -794,7 +794,7 @@ impl<H: FeltHash, const HEIGHT: usize> MerkleTree<H, HEIGHT> {
         if let Some(child_edge) = resolved_child.as_edge().cloned() {
             parent.path.extend_from_bitslice(&child_edge.path);
             if let Some(storage_index) = child_edge.storage_index {
-                self.nodes_removed.push(storage_index);
+                self.nodes_removed.push(storage_index.value());
             }
             parent.child = child_edge.child;
         }
@@ -1583,11 +1583,7 @@ mod tests {
 
     mod real_world {
         use pathfinder_common::{
-            class_commitment,
-            class_commitment_leaf_hash,
-            felt,
-            sierra_hash,
-            BlockNumber,
+            class_commitment, class_commitment_leaf_hash, felt, sierra_hash, BlockNumber,
             ClassCommitmentLeafHash,
         };
         use pathfinder_storage::RootIndexUpdate;