fix: fix get_storage_proof

The storage proofs were being returned out of order,
and certain use cases of the endpoint where there are no
storage keys but the root was still needed were not implemented.

crates/rpc/src/method/get_storage_proof.rs

@@ -1,5 +1,3 @@
-use std::collections::HashSet;
-
 use anyhow::Context;
 use pathfinder_common::trie::TrieNode;
 use pathfinder_common::{
@@ -9,6 +7,7 @@ use pathfinder_common::{
     ClassHash,
     ContractAddress,
     ContractNonce,
+    ContractRoot,
     StorageAddress,
 };
 use pathfinder_crypto::Felt;
@@ -368,18 +367,19 @@ fn get_class_proofs(
         return Ok((class_root_hash, NodeHashToNodeMappings(vec![])));
     };
 
-    let nodes: Vec<NodeHashToNodeMapping> =
-        ClassCommitmentTree::get_proofs(tx, block_number, class_hashes, class_root_idx)?
-            .into_iter()
-            .flatten()
-            .map(|(node, node_hash)| NodeHashToNodeMapping {
+    let nodes = ClassCommitmentTree::get_proofs(tx, block_number, class_hashes, class_root_idx)?;
+
+    let mut proof: Vec<NodeHashToNodeMapping> = Vec::new();
+    for node in nodes {
+        for (merkle, node_hash) in node {
+            let node_map = NodeHashToNodeMapping {
                 node_hash,
-                node: ProofNode(node),
-            })
-            .collect::<HashSet<_>>()
-            .into_iter()
-            .collect();
-    let classes_proof = NodeHashToNodeMappings(nodes);
+                node: ProofNode(merkle),
+            };
+            proof.push(node_map);
+        }
+    }
+    let classes_proof = NodeHashToNodeMappings(proof);
 
     Ok((class_root_hash, classes_proof))
 }
@@ -414,18 +414,20 @@ fn get_contract_proofs(
     };
 
     let nodes =
-        StorageCommitmentTree::get_proofs(tx, block_number, contract_addresses, storage_root_idx)?
-            .into_iter()
-            .flatten()
-            .map(|(node, node_hash)| NodeHashToNodeMapping {
+        StorageCommitmentTree::get_proofs(tx, block_number, contract_addresses, storage_root_idx)?;
+
+    let mut proof: Vec<NodeHashToNodeMapping> = Vec::new();
+    for node in nodes {
+        for (merkle, node_hash) in node {
+            let node_map = NodeHashToNodeMapping {
                 node_hash,
-                node: ProofNode(node),
-            })
-            .collect::<HashSet<_>>()
-            .into_iter()
-            .collect();
+                node: ProofNode(merkle),
+            };
+            proof.push(node_map);
+        }
+    }
 
-    let contract_proof_nodes = NodeHashToNodeMappings(nodes);
+    let contract_proof_nodes = NodeHashToNodeMappings(proof);
 
     let contract_leaves_data = contract_addresses
         .iter()
@@ -460,31 +462,55 @@ fn get_contract_storage_proofs(
         Some(contracts_storage_keys) => {
             let mut proofs = vec![];
             for csk in contracts_storage_keys {
-                let root = tx
-                    .contract_root_index(block_number, csk.contract_address)
-                    .context("Querying contract root index")?;
-
-                if let Some(root) = root {
-                    let nodes: Vec<NodeHashToNodeMapping> = ContractsStorageTree::get_proofs(
-                        &tx,
-                        csk.contract_address,
-                        block_number,
-                        &csk.storage_keys,
-                        root,
-                    )?
-                    .into_iter()
-                    .flatten()
-                    .map(|(node, node_hash)| NodeHashToNodeMapping {
+                if csk.storage_keys.is_empty() {
+                    let contract_root = tx
+                        .contract_root(block_number, csk.contract_address)
+                        .context("Querying contract's root")?
+                        .unwrap_or_default();
+
+                    let node_hash = match contract_root {
+                        ContractRoot(value) => value,
+                    };
+
+                    let merkle = TrieNode::Binary {
+                        left: Felt::default(),
+                        right: Felt::default(),
+                    };
+
+                    let node_map = NodeHashToNodeMapping {
                         node_hash,
-                        node: ProofNode(node),
-                    })
-                    .collect::<HashSet<_>>()
-                    .into_iter()
-                    .collect();
+                        node: ProofNode(merkle),
+                    };
 
-                    proofs.push(NodeHashToNodeMappings(nodes));
+                    proofs.push(NodeHashToNodeMappings(vec![node_map]));
                 } else {
-                    proofs.push(NodeHashToNodeMappings(vec![]));
+                    let root = tx
+                        .contract_root_index(block_number, csk.contract_address)
+                        .context("Querying contract root index")?;
+
+                    if let Some(root) = root {
+                        let nodes = ContractsStorageTree::get_proofs(
+                            &tx,
+                            csk.contract_address,
+                            block_number,
+                            &csk.storage_keys,
+                            root,
+                        )?;
+
+                        for node in nodes {
+                            let mut proof: Vec<NodeHashToNodeMapping> = Vec::new();
+                            for (merkle, node_hash) in node {
+                                let node_map = NodeHashToNodeMapping {
+                                    node_hash,
+                                    node: ProofNode(merkle),
+                                };
+                                proof.push(node_map);
+                            }
+                            proofs.push(NodeHashToNodeMappings(proof));
+                        }
+                    } else {
+                        proofs.push(NodeHashToNodeMappings(vec![]));
+                    }
                 }
             }