[Feat] Optimize for zktrie state

core/state/state_prove.go

@@ -74,15 +74,18 @@ func (s *StateDB) GetStorageTrieForProof(addr common.Address) (Trie, error) {
 
 // GetSecureTrieProof handle any interface with Prove (should be a Trie in most case) and
 // deliver the proof in bytes
-func (s *StateDB) GetSecureTrieProof(trieProve TrieProve, key common.Hash) ([][]byte, error) {
+func (s *StateDB) GetSecureTrieProof(trieProve TrieProve, key common.Hash) (FullProofList, common.Hash, error) {
 
-	var proof proofList
+	var proof FullProofList
+	var hash common.Hash
 	var err error
 	if s.IsZktrie() {
 		key_s, _ := zkt.ToSecureKeyBytes(key.Bytes())
-		err = trieProve.Prove(key_s.Bytes(), 0, &proof)
+		hash = common.BytesToHash(key_s.Bytes())
+		err = trieProve.Prove(hash.Bytes(), 0, &proof)
 	} else {
-		err = trieProve.Prove(crypto.Keccak256(key.Bytes()), 0, &proof)
+		hash = common.BytesToHash(crypto.Keccak256(key.Bytes()))
+		err = trieProve.Prove(hash.Bytes(), 0, &proof)
 	}
-	return proof, err
+	return proof, hash, err
 }

core/state/statedb.go

@@ -54,6 +54,33 @@ func (n *proofList) Delete(key []byte) error {
 	panic("not supported")
 }
 
+type fullProof struct {
+	Key   []byte
+	Value []byte
+}
+
+type FullProofList []fullProof
+
+func (n *FullProofList) Put(key []byte, value []byte) error {
+	*n = append(*n, fullProof{
+		Key:   key,
+		Value: value,
+	})
+	return nil
+}
+
+func (n *FullProofList) Delete(key []byte) error {
+	panic("not supported")
+}
+
+func (n FullProofList) GetData() (out [][]byte) {
+	out = make([][]byte, 0, len(n))
+	for _, i := range n {
+		out = append(out, i.Value)
+	}
+	return
+}
+
 // StateDB structs within the ethereum protocol are used to store anything
 // within the merkle trie. StateDBs take care of caching and storing
 // nested states. It's the general query interface to retrieve:
@@ -343,6 +370,21 @@ func (s *StateDB) GetProofByHash(addrHash common.Hash) ([][]byte, error) {
 	return proof, err
 }
 
+// GetFullProof returns the Merkle proof for a given account, with both node data and key
+// also the key for address is provided
+func (s *StateDB) GetFullProof(addr common.Address) (FullProofList, common.Hash, error) {
+	var hash common.Hash
+	if s.IsZktrie() {
+		addr_s, _ := zkt.ToSecureKeyBytes(addr.Bytes())
+		hash = common.BytesToHash(addr_s.Bytes())
+	} else {
+		hash = crypto.Keccak256Hash(addr.Bytes())
+	}
+	var proof FullProofList
+	err := s.trie.Prove(hash[:], 0, &proof)
+	return proof, hash, err
+}
+
 func (s *StateDB) GetLiveStateAccount(addr common.Address) *types.StateAccount {
 	obj, ok := s.stateObjects[addr]
 	if !ok {
@@ -361,7 +403,11 @@ func (s *StateDB) GetStorageProof(a common.Address, key common.Hash) ([][]byte,
 	if trie == nil {
 		return nil, errors.New("storage trie for requested address does not exist")
 	}
-	return s.GetSecureTrieProof(trie, key)
+	proof, _, err := s.GetSecureTrieProof(trie, key)
+	if err != nil {
+		return nil, err
+	}
+	return proof.GetData(), nil
 }
 
 // GetCommittedState retrieves a value from the given account's committed storage trie.

core/types/l2trace.go

@@ -45,11 +45,37 @@ type StorageTrace struct {
 	// All storage proofs BEFORE execution
 	StorageProofs map[string]map[string][]hexutil.Bytes `json:"storageProofs,omitempty"`
 
+	// The "flatten" db nodes
+	FlattenProofs map[common.Hash]hexutil.Bytes `json:"flattenProofs,omitempty"`
+
+	// The hash of secured addresses
+	AddressHashes map[common.Address]common.Hash `json:"addressHashes,omitempty"`
+	// The hash of secured store key
+	StoreKeyHashes map[common.Hash]common.Hash `json:"storeKeyHashes,omitempty"`
+
 	// Node entries for deletion, no need to distinguish what it is from, just read them
 	// into the partial db
 	DeletionProofs []hexutil.Bytes `json:"deletionProofs,omitempty"`
 }
 
+func (tr *StorageTrace) ApplyFilter(legacy bool) {
+	if legacy {
+		tr.FlattenProofs = nil
+		tr.AddressHashes = nil
+		tr.StoreKeyHashes = nil
+	} else {
+		for k := range tr.Proofs {
+			tr.Proofs[k] = []hexutil.Bytes{}
+		}
+		for _, st := range tr.StorageProofs {
+			for k := range st {
+				st[k] = []hexutil.Bytes{}
+			}
+		}
+		tr.DeletionProofs = []hexutil.Bytes{}
+	}
+}
+
 // ExecutionResult groups all structured logs emitted by the EVM
 // while replaying a transaction in debug mode as well as transaction
 // execution status, the amount of gas used and the return value

core/vm/logger.go

@@ -50,12 +50,15 @@ func (s Storage) Copy() Storage {
 
 // LogConfig are the configuration options for structured logger the EVM
 type LogConfig struct {
-	EnableMemory     bool // enable memory capture
-	DisableStack     bool // disable stack capture
-	DisableStorage   bool // disable storage capture
-	EnableReturnData bool // enable return data capture
-	Debug            bool // print output during capture end
-	Limit            int  // maximum length of output, but zero means unlimited
+	EnableMemory       bool    // enable memory capture
+	DisableStack       bool    // disable stack capture
+	DisableStorage     bool    // disable storage capture
+	EnableReturnData   bool    // enable return data capture
+	Debug              bool    // print output during capture end
+	Limit              int     // maximum length of output, but zero means unlimited
+	StorageProofFormat *string // format of storage proofs, can be
+	// "legacy" (use the legacy proof format) or
+	// "union" (output both flatten and legacy proof)
 	// Chain overrides, can be used to execute a trace using future fork rules
 	Overrides *params.ChainConfig `json:"overrides,omitempty"`
 }

eth/tracers/api_blocktrace.go

@@ -81,6 +81,8 @@ func (api *API) GetTxBlockTraceOnTopOfBlock(ctx context.Context, tx *types.Trans
 
 // Make trace environment for current block, and then get the trace for the block.
 func (api *API) createTraceEnvAndGetBlockTrace(ctx context.Context, config *TraceConfig, block *types.Block) (*types.BlockTrace, error) {
+	legacyStorageTrace := true
+	unionStorageTrace := false
 	if config == nil {
 		config = &TraceConfig{
 			LogConfig: &vm.LogConfig{
@@ -95,6 +97,14 @@ func (api *API) createTraceEnvAndGetBlockTrace(ctx context.Context, config *Trac
 		log.Warn("Tracer params is unsupported")
 	}
 
+	if config.LogConfig != nil && config.StorageProofFormat != nil {
+		if *config.StorageProofFormat == "flatten" {
+			legacyStorageTrace = false
+		} else if *config.StorageProofFormat == "union" {
+			unionStorageTrace = true
+		}
+	}
+
 	parent, err := api.blockByNumberAndHash(ctx, rpc.BlockNumber(block.NumberU64()-1), block.ParentHash())
 	if err != nil {
 		return nil, err
@@ -109,5 +119,15 @@ func (api *API) createTraceEnvAndGetBlockTrace(ctx context.Context, config *Trac
 	}
 
 	chaindb := api.backend.ChainDb()
-	return api.scrollTracerWrapper.CreateTraceEnvAndGetBlockTrace(api.backend.ChainConfig(), api.chainContext(ctx), api.backend.Engine(), chaindb, statedb, parent, block, true)
+	l2Trace, err := api.scrollTracerWrapper.CreateTraceEnvAndGetBlockTrace(api.backend.ChainConfig(), api.chainContext(ctx), api.backend.Engine(), chaindb, statedb, parent, block, true)
+	if err != nil {
+		return nil, err
+	}
+	if !unionStorageTrace {
+		l2Trace.StorageTrace.ApplyFilter(legacyStorageTrace)
+		for _, st := range l2Trace.TxStorageTraces {
+			st.ApplyFilter(legacyStorageTrace)
+		}
+	}
+	return l2Trace, nil
 }