From 2f7a13d9d5f51d8350b7e9a283dcebc92b9489f8 Mon Sep 17 00:00:00 2001 From: weiihann Date: Tue, 21 Jan 2025 12:53:21 +0800 Subject: [PATCH] remove dups from storage keys --- rpc/storage.go | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/rpc/storage.go b/rpc/storage.go index 9f850b32cf..c2856cb521 100644 --- a/rpc/storage.go +++ b/rpc/storage.go @@ -165,6 +165,21 @@ func (h *Handler) StorageProof(id BlockID, classes = utils.Unique(classes) contracts = utils.Unique(contracts) + // Remove duplicates from the storage keys + mergedStorageKeys := make(map[felt.Felt][]felt.Felt) + for _, storageKey := range storageKeys { + if existing, ok := mergedStorageKeys[storageKey.Contract]; ok { + mergedStorageKeys[storageKey.Contract] = append(existing, storageKey.Keys...) + } else { + mergedStorageKeys[storageKey.Contract] = storageKey.Keys + } + } + + uniqueStorageKeys := make([]StorageKeys, 0, len(mergedStorageKeys)) + for contract, keys := range mergedStorageKeys { + uniqueStorageKeys = append(uniqueStorageKeys, StorageKeys{Contract: contract, Keys: utils.Unique(keys)}) + } + classProof, err := getClassProof(classTrie, classes) if err != nil { return nil, ErrInternal.CloneWithData(err) @@ -175,7 +190,7 @@ func (h *Handler) StorageProof(id BlockID, return nil, ErrInternal.CloneWithData(err) } - contractStorageProof, err := getContractStorageProof(state, storageKeys) + contractStorageProof, err := getContractStorageProof(state, uniqueStorageKeys) if err != nil { return nil, ErrInternal.CloneWithData(err) } @@ -261,8 +276,7 @@ func getContractStorageProof(state core.StateReader, storageKeys []StorageKeys) } contractStorageProof := trie.NewProofNodeSet() - keys := utils.Unique(storageKey.Keys) - for _, key := range keys { + for _, key := range storageKey.Keys { if err := contractStorageTrie.Prove(&key, contractStorageProof); err != nil { return nil, err }