getStorageProof rpc method

blockchain/blockchain.go

@@ -36,6 +36,7 @@ type Reader interface {
 	StateUpdateByHash(hash *felt.Felt) (update *core.StateUpdate, err error)
 
 	HeadState() (core.StateReader, StateCloser, error)
+	HeadTrie() (core.TrieReader, StateCloser, error)
 	StateAtBlockHash(blockHash *felt.Felt) (core.StateReader, StateCloser, error)
 	StateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error)
 	PendingState() (core.StateReader, StateCloser, error)
@@ -769,6 +770,17 @@ func (b *Blockchain) HeadState() (core.StateReader, StateCloser, error) {
 	return core.NewState(txn), txn.Discard, nil
 }
 
+func (b *Blockchain) HeadTrie() (core.TrieReader, StateCloser, error) {
+	// Note: I'm not sure I should open a new db txn since the TrieReader is a State
+	// so the same instance of the state we create in HeadState will do job.
+	txn, err := b.database.NewTransaction(false)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return core.NewState(txn), txn.Discard, nil
+}
+
 // StateAtBlockNumber returns a StateReader that provides a stable view to the state at the given block number
 func (b *Blockchain) StateAtBlockNumber(blockNumber uint64) (core.StateReader, StateCloser, error) {
 	b.listener.OnRead("StateAtBlockNumber")

core/state.go

@@ -44,6 +44,17 @@ type StateReader interface {
 	Class(classHash *felt.Felt) (*DeclaredClass, error)
 }
 
+// TrieReader used for storage proofs, can only be supported by current state implementation (for now, we plan to add db snapshots)
+var _ TrieReader = (*State)(nil)
+
+//go:generate mockgen -destination=../mocks/mock_trie.go -package=mocks github.com/NethermindEth/juno/core TrieReader
+type TrieReader interface {
+	ClassTrie() (*trie.Trie, func() error, error)
+	StorageTrie() (*trie.Trie, func() error, error)
+	StorageTrieForAddr(addr *felt.Felt) (*trie.Trie, error)
+	StateAndClassRoot() (*felt.Felt, *felt.Felt, error)
+}
+
 type State struct {
 	*history
 	txn db.Transaction
@@ -129,6 +140,18 @@ func (s *State) storage() (*trie.Trie, func() error, error) {
 	return s.globalTrie(db.StateTrie, trie.NewTriePedersen)
 }
 
+func (s *State) StorageTrie() (*trie.Trie, func() error, error) {
+	return s.storage()
+}
+
+func (s *State) ClassTrie() (*trie.Trie, func() error, error) {
+	return s.classesTrie()
+}
+
+func (s *State) StorageTrieForAddr(addr *felt.Felt) (*trie.Trie, error) {
+	return storage(addr, s.txn)
+}
+
 func (s *State) classesTrie() (*trie.Trie, func() error, error) {
 	return s.globalTrie(db.ClassesTrie, trie.NewTriePoseidon)
 }
@@ -547,7 +570,7 @@ func (s *State) Revert(blockNumber uint64, update *StateUpdate) error {
 
 	err = s.performStateDeletions(blockNumber, update.StateDiff)
 	if err != nil {
-		return fmt.Errorf("error performing state deletions: %v", err)
+		return fmt.Errorf("build reverse diff: %v", err)
 	}
 
 	stateTrie, storageCloser, err := s.storage()
@@ -581,6 +604,7 @@ func (s *State) purgeNoClassContracts() error {
 	// As noClassContracts are not in StateDiff.DeployedContracts we can only purge them if their storage no longer exists.
 	// Updating contracts with reverse diff will eventually lead to the deletion of noClassContract's storage key from db. Thus,
 	// we can use the lack of key's existence as reason for purging noClassContracts.
+
 	for addr := range noClassContracts {
 		noClassC, err := NewContractUpdater(&addr, s.txn)
 		if err != nil {
@@ -743,3 +767,35 @@ func (s *State) performStateDeletions(blockNumber uint64, diff *StateDiff) error
 
 	return nil
 }
+
+func (s *State) StateAndClassRoot() (*felt.Felt, *felt.Felt, error) {
+	var storageRoot, classesRoot *felt.Felt
+
+	sStorage, closer, err := s.storage()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if storageRoot, err = sStorage.Root(); err != nil {
+		return nil, nil, err
+	}
+
+	if err = closer(); err != nil {
+		return nil, nil, err
+	}
+
+	classes, closer, err := s.classesTrie()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if classesRoot, err = classes.Root(); err != nil {
+		return nil, nil, err
+	}
+
+	if err = closer(); err != nil {
+		return nil, nil, err
+	}
+
+	return storageRoot, classesRoot, nil
+}

core/trie/key.go

@@ -28,6 +28,9 @@ func (k *Key) SubKey(n uint8) (*Key, error) {
 	if n > k.len {
 		return nil, errors.New(fmt.Sprint("cannot subtract key of length %i from key of length %i", n, k.len))
 	}
+	if n == k.len {
+		return &Key{}, nil
+	}
 
 	newKey := &Key{len: n}
 	copy(newKey.bitset[:], k.bitset[len(k.bitset)-int((k.len+7)/8):]) //nolint:mnd

core/trie/node.go

@@ -18,7 +18,7 @@ type Node struct {
 }
 
 // Hash calculates the hash of a [Node]
-func (n *Node) Hash(path *Key, hashFunc hashFunc) *felt.Felt {
+func (n *Node) Hash(path *Key, hashFunc HashFunc) *felt.Felt {
 	if path.Len() == 0 {
 		// we have to deference the Value, since the Node can released back
 		// to the NodePool and be reused anytime
@@ -33,7 +33,7 @@ func (n *Node) Hash(path *Key, hashFunc hashFunc) *felt.Felt {
 }
 
 // Hash calculates the hash of a [Node]
-func (n *Node) HashFromParent(parnetKey, nodeKey *Key, hashFunc hashFunc) *felt.Felt {
+func (n *Node) HashFromParent(parnetKey, nodeKey *Key, hashFunc HashFunc) *felt.Felt {
 	path := path(nodeKey, parnetKey)
 	return n.Hash(&path, hashFunc)
 }

core/trie/proof.go

@@ -13,7 +13,7 @@ var (
 )
 
 type ProofNode interface {
-	Hash(hash hashFunc) *felt.Felt
+	Hash(hash HashFunc) *felt.Felt
 	Len() uint8
 	PrettyPrint()
 }
@@ -23,7 +23,7 @@ type Binary struct {
 	RightHash *felt.Felt
 }
 
-func (b *Binary) Hash(hash hashFunc) *felt.Felt {
+func (b *Binary) Hash(hash HashFunc) *felt.Felt {
 	return hash(b.LeftHash, b.RightHash)
 }
 
@@ -42,7 +42,7 @@ type Edge struct {
 	Path  *Key       // path from parent to child
 }
 
-func (e *Edge) Hash(hash hashFunc) *felt.Felt {
+func (e *Edge) Hash(hash HashFunc) *felt.Felt {
 	length := make([]byte, len(e.Path.bitset))
 	length[len(e.Path.bitset)-1] = e.Path.len
 	pathFelt := e.Path.Felt()
@@ -199,7 +199,7 @@ func traverseNodes(currNode ProofNode, path *[]ProofNode, nodeHashes map[felt.Fe
 // merges paths in the specified order [commonNodes..., leftNodes..., rightNodes...]
 // ordering of the merged path is not important
 // since SplitProofPath can discover the left and right paths using the merged path and the rootHash
-func MergeProofPaths(leftPath, rightPath []ProofNode, hash hashFunc) ([]ProofNode, *felt.Felt, error) {
+func MergeProofPaths(leftPath, rightPath []ProofNode, hash HashFunc) ([]ProofNode, *felt.Felt, error) {
 	merged := []ProofNode{}
 	minLen := min(len(leftPath), len(rightPath))
 
@@ -236,7 +236,7 @@ func MergeProofPaths(leftPath, rightPath []ProofNode, hash hashFunc) ([]ProofNod
 // SplitProofPath splits the merged proof path into two paths (left and right), which were merged before
 // it first validates that the merged path is not circular, the split happens at most once and rootHash exists
 // then calls traverseNodes to split the path to left and right paths
-func SplitProofPath(mergedPath []ProofNode, rootHash *felt.Felt, hash hashFunc) ([]ProofNode, []ProofNode, error) {
+func SplitProofPath(mergedPath []ProofNode, rootHash *felt.Felt, hash HashFunc) ([]ProofNode, []ProofNode, error) {
 	commonPath := []ProofNode{}
 	leftPath := []ProofNode{}
 	rightPath := []ProofNode{}
@@ -316,7 +316,7 @@ func GetProof(key *Key, tri *Trie) ([]ProofNode, error) {
 
 // verifyProof checks if `leafPath` leads from `root` to `leafHash` along the `proofNodes`
 // https://github.com/eqlabs/pathfinder/blob/main/crates/merkle-tree/src/tree.rs#L2006
-func VerifyProof(root *felt.Felt, key *Key, value *felt.Felt, proofs []ProofNode, hash hashFunc) bool {
+func VerifyProof(root *felt.Felt, key *Key, value *felt.Felt, proofs []ProofNode, hash HashFunc) bool {
 	expectedHash := root
 	remainingPath := NewKey(key.len, key.bitset[:])
 	for i, proofNode := range proofs {
@@ -345,7 +345,7 @@ func VerifyProof(root *felt.Felt, key *Key, value *felt.Felt, proofs []ProofNode
 				return true
 			}
 
-			if !proofNode.Path.Equal(subKey) {
+			if !proofNode.Path.Equal(subKey) && !subKey.Equal(&Key{}) {
 				return false
 			}
 			expectedHash = proofNode.Child
@@ -363,7 +363,7 @@ func VerifyProof(root *felt.Felt, key *Key, value *felt.Felt, proofs []ProofNode
 // and therefore it's hash won't match the expected root.
 // ref: https://github.com/ethereum/go-ethereum/blob/v1.14.3/trie/proof.go#L484
 func VerifyRangeProof(root *felt.Felt, keys, values []*felt.Felt, proofKeys [2]*Key, proofValues [2]*felt.Felt,
-	proofs [2][]ProofNode, hash hashFunc,
+	proofs [2][]ProofNode, hash HashFunc,
 ) (bool, error) {
 	// Step 0: checks
 	if len(keys) != len(values) {
@@ -440,7 +440,7 @@ func ensureMonotonicIncreasing(proofKeys [2]*Key, keys []*felt.Felt) error {
 }
 
 // compressNode determines if the node needs compressed, and if so, the len needed to arrive at the next key
-func compressNode(idx int, proofNodes []ProofNode, hashF hashFunc) (int, uint8, error) {
+func compressNode(idx int, proofNodes []ProofNode, hashF HashFunc) (int, uint8, error) {
 	parent := proofNodes[idx]
 
 	if idx == len(proofNodes)-1 {
@@ -474,7 +474,7 @@ func compressNode(idx int, proofNodes []ProofNode, hashF hashFunc) (int, uint8,
 }
 
 func assignChild(i, compressedParent int, parentNode *Node,
-	nilKey, leafKey, parentKey *Key, proofNodes []ProofNode, hashF hashFunc,
+	nilKey, leafKey, parentKey *Key, proofNodes []ProofNode, hashF HashFunc,
 ) (*Key, error) {
 	childInd := i + compressedParent + 1
 	childKey, err := getChildKey(childInd, parentKey, leafKey, nilKey, proofNodes, hashF)
@@ -494,7 +494,7 @@ func assignChild(i, compressedParent int, parentNode *Node,
 // ProofToPath returns a set of storage nodes from the root to the end of the proof path.
 // The storage nodes will have the hashes of the children, but only the key of the child
 // along the path outlined by the proof.
-func ProofToPath(proofNodes []ProofNode, leafKey *Key, hashF hashFunc) ([]StorageNode, error) {
+func ProofToPath(proofNodes []ProofNode, leafKey *Key, hashF HashFunc) ([]StorageNode, error) {
 	pathNodes := []StorageNode{}
 
 	// Child keys that can't be derived are set to nilKey, so that we can store the node
@@ -552,7 +552,7 @@ func ProofToPath(proofNodes []ProofNode, leafKey *Key, hashF hashFunc) ([]Storag
 	return pathNodes, nil
 }
 
-func skipNode(pNode ProofNode, pathNodes []StorageNode, hashF hashFunc) bool {
+func skipNode(pNode ProofNode, pathNodes []StorageNode, hashF HashFunc) bool {
 	lastNode := pathNodes[len(pathNodes)-1].node
 	noLeftMatch, noRightMatch := false, false
 	if lastNode.LeftHash != nil && !pNode.Hash(hashF).Equal(lastNode.LeftHash) {
@@ -607,7 +607,7 @@ func getParentKey(idx int, compressedParentOffset uint8, leafKey *Key,
 	return crntKey, err
 }
 
-func getChildKey(childIdx int, crntKey, leafKey, nilKey *Key, proofNodes []ProofNode, hashF hashFunc) (*Key, error) {
+func getChildKey(childIdx int, crntKey, leafKey, nilKey *Key, proofNodes []ProofNode, hashF HashFunc) (*Key, error) {
 	if childIdx > len(proofNodes)-1 {
 		return nilKey, nil
 	}

core/trie/trie.go

@@ -13,7 +13,7 @@ import (
 	"github.com/NethermindEth/juno/db"
 )
 
-type hashFunc func(*felt.Felt, *felt.Felt) *felt.Felt
+type HashFunc func(*felt.Felt, *felt.Felt) *felt.Felt
 
 // Trie is a dense Merkle Patricia Trie (i.e., all internal nodes have two children).
 //
@@ -37,7 +37,7 @@ type Trie struct {
 	rootKey *Key
 	maxKey  *felt.Felt
 	storage *Storage
-	hash    hashFunc
+	hash    HashFunc
 
 	dirtyNodes     []*Key
 	rootKeyIsDirty bool
@@ -53,7 +53,7 @@ func NewTriePoseidon(storage *Storage, height uint8) (*Trie, error) {
 	return newTrie(storage, height, crypto.Poseidon)
 }
 
-func newTrie(storage *Storage, height uint8, hash hashFunc) (*Trie, error) {
+func newTrie(storage *Storage, height uint8, hash HashFunc) (*Trie, error) {
 	if height > felt.Bits {
 		return nil, fmt.Errorf("max trie height is %d, got: %d", felt.Bits, height)
 	}
@@ -668,6 +668,10 @@ func (t *Trie) RootKey() *Key {
 	return t.rootKey
 }
 
+func (t *Trie) HashFunc() HashFunc {
+	return t.hash
+}
+
 func (t *Trie) Dump() {
 	t.dump(0, nil)
 }