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,15 +1,7 @@
-use std::collections::HashSet;
-
 use anyhow::Context;
 use pathfinder_common::trie::TrieNode;
 use pathfinder_common::{
-    BlockHash,
-    BlockId,
-    BlockNumber,
-    ClassHash,
-    ContractAddress,
-    ContractNonce,
-    ContractRoot,
+    BlockHash, BlockId, BlockNumber, ClassHash, ContractAddress, ContractNonce, ContractRoot,
     StorageAddress,
 };
 use pathfinder_crypto::Felt;
@@ -378,7 +370,7 @@ fn get_class_proofs(
                 node_hash,
                 node: ProofNode(node),
             })
-            .collect::<HashSet<_>>()
+            .collect::<Vec<_>>()
             .into_iter()
             .collect();
     let classes_proof = NodeHashToNodeMappings(nodes);
@@ -423,7 +415,7 @@ fn get_contract_proofs(
                 node_hash,
                 node: ProofNode(node),
             })
-            .collect::<HashSet<_>>()
+            .collect::<Vec<_>>()
             .into_iter()
             .collect();
 
@@ -476,24 +468,24 @@ fn get_contract_storage_proofs(
                     .context("Querying contract root index")?;
 
                 if let Some(root) = root {
-                    let nodes: Vec<NodeHashToNodeMapping> = ContractsStorageTree::get_proofs(
+                    ContractsStorageTree::get_proofs(
                         &tx,
                         csk.contract_address,
                         block_number,
                         &csk.storage_keys,
                         root,
                     )?
                     .into_iter()
-                    .flatten()
-                    .map(|(node, node_hash)| NodeHashToNodeMapping {
-                        node_hash,
-                        node: ProofNode(node),
-                    })
-                    .collect::<HashSet<_>>()
-                    .into_iter()
-                    .collect();
-
-                    proofs.push(NodeHashToNodeMappings(nodes));
+                    .for_each(|vec_nodes| {
+                        let nodes = vec_nodes
+                            .into_iter()
+                            .map(|(node, node_hash)| NodeHashToNodeMapping {
+                                node_hash,
+                                node: ProofNode(node),
+                            })
+                            .collect::<Vec<_>>();
+                        proofs.push(NodeHashToNodeMappings(nodes));
+                    });
                 } else {
                     proofs.push(NodeHashToNodeMappings(vec![]));
                 }
@@ -872,6 +864,7 @@ mod tests {
         Class,
         ContractNonce,
         ContractStorage,
+        ContractStorageEmptyKeys,
     }
 
     impl PartialRequest {
@@ -933,6 +926,25 @@ mod tests {
                         }]),
                     }
                 }
+                Self::ContractStorageEmptyKeys => {
+                    let (contract_address, _update) = block
+                        .state_update
+                        .as_ref()
+                        .unwrap()
+                        .contract_updates
+                        .iter()
+                        .next()
+                        .unwrap();
+                    Input {
+                        block_id: BlockId::Number(BlockNumber::new_or_panic(1)),
+                        class_hashes: None,
+                        contract_addresses: None,
+                        contracts_storage_keys: Some(vec![ContractStorageKeys {
+                            contract_address: *contract_address,
+                            storage_keys: vec![],
+                        }]),
+                    }
+                }
             }
         }
     }
@@ -941,6 +953,7 @@ mod tests {
     #[case::class_request(PartialRequest::Class)]
     #[case::contract_request(PartialRequest::ContractNonce)]
     #[case::contract_storage_request(PartialRequest::ContractStorage)]
+    #[case::contract_storage_request(PartialRequest::ContractStorageEmptyKeys)]
     #[tokio::test]
     async fn partial_query(#[case] req: PartialRequest) {
         use pathfinder_storage::fake::{fill, generate};
@@ -999,6 +1012,23 @@ mod tests {
                 assert!(output.contracts_proof.nodes.0.is_empty());
                 assert!(output.contracts_proof.contract_leaves_data.is_empty());
             }
+            PartialRequest::ContractStorageEmptyKeys => {
+                // Some contract storage proof should be present
+                assert!(!output.contracts_storage_proofs.is_empty());
+                // It should have only one element and the nodes should be default empty
+                assert!(output.contracts_storage_proofs.len() == 1);
+                assert_eq!(
+                    output.contracts_storage_proofs[0].0[0].node.0,
+                    TrieNode::Binary {
+                        left: Felt::default(),
+                        right: Felt::default(),
+                    }
+                );
+                // The rest should be empty
+                assert!(output.classes_proof.0.is_empty());
+                assert!(output.contracts_proof.nodes.0.is_empty());
+                assert!(output.contracts_proof.contract_leaves_data.is_empty());
+            }
         }
     }
 }