feat : added pathfinder_getClassProof method

.gitignore

@@ -24,4 +24,4 @@ pathfinder-var.env
 .vscode/
 
 # mdbook compilation
-**/book/
+**/book/

crates/merkle-tree/src/class.rs

@@ -1,16 +1,25 @@
+use std::ops::ControlFlow;
+
 use anyhow::Context;
-use pathfinder_common::hash::PoseidonHash;
+use bitvec::order::Msb0;
+use bitvec::prelude::BitSlice;
+use pathfinder_common::hash::{PedersenHash, PoseidonHash};
+use pathfinder_common::trie::TrieNode;
 use pathfinder_common::{
     BlockNumber,
+    CasmHash,
     ClassCommitment,
     ClassCommitmentLeafHash,
     ClassHash,
     SierraHash,
+    StorageAddress,
+    StorageValue,
 };
 use pathfinder_crypto::Felt;
 use pathfinder_storage::{Transaction, TrieUpdate};
 
-use crate::tree::MerkleTree;
+use crate::merkle_node::InternalNode;
+use crate::tree::{MerkleTree, Visit};
 
 /// A [Patricia Merkle tree](MerkleTree) used to calculate commitments to
 /// Starknet's Sierra classes.
@@ -71,6 +80,162 @@ impl<'tx> ClassCommitmentTree<'tx> {
         let commitment = ClassCommitment(update.root_commitment);
         Ok((commitment, update))
     }
+
+    /// Generates a proof for a given `key`
+    pub fn get_proof(
+        tx: &'tx Transaction<'tx>,
+        block: BlockNumber,
+        class_hash: ClassHash,
+    ) -> anyhow::Result<Option<Vec<TrieNode>>> {
+        let root = tx
+            .class_root_index(block)
+            .context("Querying class root index")?;
+
+        let Some(root) = root else {
+            return Ok(None);
+        };
+
+        let storage = ClassTrieStorage {
+            tx,
+            block: Some(block),
+        };
+
+        let casm = tx
+            .casm_hash_at(block.into(), class_hash)
+            .context("Querying CASM hash")?;
+
+        let Some(casm) = casm else {
+            return Ok(None);
+        };
+
+        MerkleTree::<PoseidonHash, 251>::get_proof(root, &storage, casm.view_bits())
+    }
+}
+
+/// A [Patricia Merkle tree](MerkleTree) used to calculate commitments to
+/// Starknet's Sierra classes.
+///
+/// It maps a class's [SierraHash] to its [ClassCommitmentLeafHash]
+///
+/// Tree data is persisted by a sqlite table 'tree_class'.
+
+pub struct ClassStorageTree<'tx> {
+    tree: MerkleTree<PoseidonHash, 251>,
+    storage: ClassStorage<'tx>,
+}
+
+impl<'tx> ClassStorageTree<'tx> {
+    pub fn empty(tx: &'tx Transaction<'tx>) -> Self {
+        let storage = ClassStorage { tx, block: None };
+        let tree = MerkleTree::empty();
+
+        Self { tree, storage }
+    }
+
+    pub fn load(tx: &'tx Transaction<'tx>, block: BlockNumber) -> anyhow::Result<Self> {
+        let root = tx
+            .class_root_index(block)
+            .context("Querying class root index")?;
+
+        let Some(root) = root else {
+            return Ok(Self::empty(tx));
+        };
+
+        let storage = ClassStorage {
+            tx,
+            block: Some(block),
+        };
+
+        let tree = MerkleTree::new(root);
+
+        Ok(Self { tree, storage })
+    }
+
+    pub fn with_verify_hashes(mut self, verify_hashes: bool) -> Self {
+        self.tree = self.tree.with_verify_hashes(verify_hashes);
+        self
+    }
+
+    /// Generates a proof for `key`. See [`MerkleTree::get_proof`].
+    pub fn get_proof(
+        tx: &'tx Transaction<'tx>,
+        block: BlockNumber,
+        key: &BitSlice<u8, Msb0>,
+    ) -> anyhow::Result<Option<Vec<TrieNode>>> {
+        let root = tx
+            .class_root_index(block)
+            .context("Querying class root index")?;
+
+        let Some(root) = root else {
+            return Ok(None);
+        };
+
+        let storage = ClassStorage {
+            tx,
+            block: Some(block),
+        };
+
+        MerkleTree::<PedersenHash, 251>::get_proof(root, &storage, key)
+    }
+
+    pub fn set(&mut self, address: StorageAddress, value: StorageValue) -> anyhow::Result<()> {
+        let key = address.view_bits().to_owned();
+        self.tree.set(&self.storage, key, value.0)
+    }
+
+    /// Commits the changes and calculates the new node hashes. Returns the new
+    /// commitment and any potentially newly created nodes.
+    pub fn commit(self) -> anyhow::Result<(CasmHash, TrieUpdate)> {
+        let update = self.tree.commit(&self.storage)?;
+        let commitment = CasmHash(update.root_commitment);
+        Ok((commitment, update))
+    }
+
+    /// See [`MerkleTree::dfs`]
+    pub fn dfs<B, F: FnMut(&InternalNode, &BitSlice<u8, Msb0>) -> ControlFlow<B, Visit>>(
+        &mut self,
+        f: &mut F,
+    ) -> anyhow::Result<Option<B>> {
+        self.tree.dfs(&self.storage, f)
+    }
+}
+
+struct ClassTrieStorage<'tx> {
+    tx: &'tx Transaction<'tx>,
+    block: Option<BlockNumber>,
+}
+
+impl crate::storage::Storage for ClassTrieStorage<'_> {
+    fn get(&self, index: u64) -> anyhow::Result<Option<pathfinder_storage::StoredNode>> {
+        self.tx.storage_trie_node(index)
+    }
+
+    fn hash(&self, index: u64) -> anyhow::Result<Option<Felt>> {
+        self.tx.storage_trie_node_hash(index)
+    }
+
+    fn leaf(&self, path: &BitSlice<u8, Msb0>) -> anyhow::Result<Option<Felt>> {
+        assert!(path.len() == 251);
+
+        let Some(block) = self.block else {
+            return Ok(None);
+        };
+
+        let sierra =
+            ClassHash(Felt::from_bits(path).context("Mapping leaf path to contract address")?);
+
+        let casm = self
+            .tx
+            .casm_hash_at(block.into(), sierra)
+            .context("Querying CASM hash")?;
+        let Some(casm) = casm else {
+            return Ok(None);
+        };
+
+        let value = self.tx.class_commitment_leaf(block, &casm)?.map(|x| x.0);
+
+        Ok(value)
+    }
 }
 
 struct ClassStorage<'tx> {

crates/merkle-tree/src/lib.rs

@@ -3,7 +3,7 @@ pub mod merkle_node;
 pub mod storage;
 pub mod tree;
 
-mod class;
+pub mod class;
 mod contract;
 mod transaction;
 

crates/rpc/src/pathfinder.rs

@@ -8,4 +8,5 @@ pub fn register_routes() -> RpcRouterBuilder {
         .register("pathfinder_version",              || { pathfinder_common::consts::VERGEN_GIT_DESCRIBE })
         .register("pathfinder_getProof",             methods::get_proof)
         .register("pathfinder_getTransactionStatus", methods::get_transaction_status)
+        .register("pathfinder_getClassProof",        methods::get_proof_class)
 }

crates/rpc/src/pathfinder/methods.rs

@@ -1,5 +1,5 @@
 mod get_proof;
 mod get_transaction_status;
 
-pub(crate) use get_proof::get_proof;
+pub(crate) use get_proof::{get_proof, get_proof_class};
 pub(crate) use get_transaction_status::get_transaction_status;

crates/rpc/src/pathfinder/methods/get_proof.rs

@@ -3,7 +3,8 @@ use pathfinder_common::prelude::*;
 use pathfinder_common::trie::TrieNode;
 use pathfinder_common::BlockId;
 use pathfinder_crypto::Felt;
-use pathfinder_merkle_tree::{ContractsStorageTree, StorageCommitmentTree};
+use pathfinder_merkle_tree::class::ClassStorageTree;
+use pathfinder_merkle_tree::{ClassCommitmentTree, ContractsStorageTree, StorageCommitmentTree};
 use serde::{Deserialize, Serialize};
 use serde_with::skip_serializing_none;
 
@@ -16,6 +17,13 @@ pub struct GetProofInput {
     pub keys: Vec<StorageAddress>,
 }
 
+#[derive(Deserialize, Debug, PartialEq, Eq)]
+pub struct GetProofInputClass {
+    pub block_id: BlockId,
+    pub class_hash: ClassHash,
+    pub keys: Vec<StorageAddress>,
+}
+
 // FIXME: allow `generate_rpc_error_subset!` to work with enum struct variants.
 // This may not actually be possible though.
 #[derive(Debug)]
@@ -151,6 +159,18 @@ pub struct GetProofOutput {
     contract_data: Option<ContractData>,
 }
 
+#[derive(Debug, Serialize)]
+#[skip_serializing_none]
+pub struct GetProofOutputClass {
+    /// Required to verify that the hash of the class commitment and the root of
+    /// the [contract_proof](GetProofOutput::contract_proof) matches the
+    /// [state_commitment](Self#state_commitment). Present only for Starknet
+    /// blocks 0.11.0 onwards.
+    class_commitment: Option<ClassCommitment>,
+    /// Membership / Non-membership proof for the queried contract classes
+    class_proof: ProofNodes,
+}
+
 /// Returns all the necessary data to trustlessly verify storage slots for a
 /// particular contract.
 pub async fn get_proof(
@@ -278,6 +298,95 @@ pub async fn get_proof(
     jh.await.context("Database read panic or shutting down")?
 }
 
+/// Returns all the necessary data to trustlessly verify class changes for a
+/// particular contract.
+pub async fn get_proof_class(
+    context: RpcContext,
+    input: GetProofInputClass,
+) -> Result<GetProofOutputClass, GetProofError> {
+    const MAX_KEYS: usize = 100;
+    if input.keys.len() > MAX_KEYS {
+        return Err(GetProofError::ProofLimitExceeded {
+            limit: MAX_KEYS as u32,
+            requested: input.keys.len() as u32,
+        });
+    }
+
+    let block_id = match input.block_id {
+        BlockId::Pending => {
+            return Err(GetProofError::Internal(anyhow!(
+                "'pending' is not currently supported by this method!"
+            )))
+        }
+        other => other.try_into().expect("Only pending cast should fail"),
+    };
+
+    let storage = context.storage.clone();
+    let span = tracing::Span::current();
+
+    let jh = tokio::task::spawn_blocking(move || {
+        let _g = span.enter();
+        let mut db = storage
+            .connection()
+            .context("Opening database connection")?;
+
+        let tx = db.transaction().context("Creating database transaction")?;
+
+        // Use internal error to indicate that the process of querying for a particular
+        // block failed, which is not the same as being sure that the block is
+        // not in the db.
+        let header = tx
+            .block_header(block_id)
+            .context("Fetching block header")?
+            .ok_or(GetProofError::BlockNotFound)?;
+
+        let class_commitment = match header.class_commitment {
+            ClassCommitment::ZERO => None,
+            other => Some(other),
+        };
+
+        // Generate a proof for this class. If the class does not exist, this will
+        // be a "non membership" proof.
+        let class_proof = ClassCommitmentTree::get_proof(&tx, header.number, input.class_hash)
+            .context("Creating contract proof")?
+            .ok_or(GetProofError::ProofMissing)?;
+        let class_proof = ProofNodes(class_proof);
+
+        let class_root_exists = tx
+            .class_root_exists(header.number)
+            .context("Fetching class root existence")?;
+
+        if !class_root_exists {
+            return Ok(GetProofOutputClass {
+                class_commitment,
+                class_proof,
+            });
+        };
+
+        let mut class_proofs = Vec::new();
+        for k in &input.keys {
+            let proof = ClassStorageTree::get_proof(&tx, header.number, k.view_bits())
+                .context("Get proof from class tree")?
+                .ok_or_else(|| {
+                    let e = anyhow!(
+                        "Storage proof missing for key {:?}, but should be present",
+                        k
+                    );
+                    tracing::warn!("{e}");
+                    e
+                })?;
+            class_proofs.push(ProofNodes(proof));
+        }
+
+        Ok(GetProofOutputClass {
+            class_commitment,
+            class_proof,
+        })
+    });
+
+    jh.await.context("Database read panic or shutting down")?
+}
+
 #[cfg(test)]
 mod tests {
     use pathfinder_common::macro_prelude::*;