[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

eth/tracers/api_blocktrace.go

@@ -81,6 +81,7 @@ 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) {
+	legacyTrace := false
 	if config == nil {
 		config = &TraceConfig{
 			LogConfig: &vm.LogConfig{
@@ -91,6 +92,10 @@ func (api *API) createTraceEnvAndGetBlockTrace(ctx context.Context, config *Trac
 			},
 		}
 	} else if config.Tracer != nil {
+		if *config.Tracer == "legacy" {
+			legacyTrace = true
+		}
+
 		config.Tracer = nil
 		log.Warn("Tracer params is unsupported")
 	}
@@ -109,5 +114,13 @@ 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
+	}
+	l2Trace.StorageTrace.ApplyFilter(legacyTrace)
+	for _, st := range l2Trace.TxStorageTraces {
+		st.ApplyFilter(legacyTrace)
+	}
+	return l2Trace, nil
 }

rollup/tracing/tracing.go

@@ -26,6 +26,7 @@ import (
 	"github.com/scroll-tech/go-ethereum/rollup/fees"
 	"github.com/scroll-tech/go-ethereum/rollup/rcfg"
 	"github.com/scroll-tech/go-ethereum/rollup/withdrawtrie"
+	zktrie "github.com/scroll-tech/zktrie/types"
 )
 
 var (
@@ -40,6 +41,9 @@ var (
 // TracerWrapper implements ScrollTracerWrapper interface
 type TracerWrapper struct{}
 
+// alias for proof list
+type proofList = state.FullProofList
+
 // NewTracerWrapper TracerWrapper creates a new TracerWrapper
 func NewTracerWrapper() *TracerWrapper {
 	return &TracerWrapper{}
@@ -69,7 +73,7 @@ type TraceEnv struct {
 	// The following Mutexes are used to protect against parallel read/write,
 	// since txs are executed in parallel.
 	pMu sync.Mutex // for `TraceEnv.StorageTrace.Proofs`
-	sMu sync.Mutex // for `TraceEnv.state``
+	sMu sync.Mutex // for `TraceEnv.state`
 	cMu sync.Mutex // for `TraceEnv.Codes`
 
 	*types.StorageTrace
@@ -109,10 +113,13 @@ func CreateTraceEnvHelper(chainConfig *params.ChainConfig, logConfig *vm.LogConf
 		state:            statedb,
 		blockCtx:         blockCtx,
 		StorageTrace: &types.StorageTrace{
-			RootBefore:    rootBefore,
-			RootAfter:     block.Root(),
-			Proofs:        make(map[string][]hexutil.Bytes),
-			StorageProofs: make(map[string]map[string][]hexutil.Bytes),
+			RootBefore:     rootBefore,
+			RootAfter:      block.Root(),
+			Proofs:         make(map[string][]hexutil.Bytes),
+			StorageProofs:  make(map[string]map[string][]hexutil.Bytes),
+			FlattenProofs:  make(map[common.Hash]hexutil.Bytes),
+			AddressHashes:  make(map[common.Address]common.Hash),
+			StoreKeyHashes: make(map[common.Hash]common.Hash),
 		},
 		Codes:             make(map[common.Hash]vm.CodeInfo),
 		ZkTrieTracer:      make(map[string]state.ZktrieProofTracer),
@@ -170,12 +177,15 @@ func CreateTraceEnv(chainConfig *params.ChainConfig, chainContext core.ChainCont
 
 	key := coinbase.String()
 	if _, exist := env.Proofs[key]; !exist {
-		proof, err := env.state.GetProof(coinbase)
+		proof, addrHash, err := env.state.GetFullProof(coinbase)
 		if err != nil {
 			log.Error("Proof for coinbase not available", "coinbase", coinbase, "error", err)
 			// but we still mark the proofs map with nil array
 		}
-		env.Proofs[key] = types.WrapProof(proof)
+		// TODO:
+		env.AddressHashes[coinbase] = addrHash
+		env.fillFlattenStorageProof(nil, proof)
+		env.Proofs[key] = types.WrapProof(proof.GetData())
 	}
 
 	return env, nil
@@ -249,21 +259,32 @@ func (env *TraceEnv) GetBlockTrace(block *types.Block) (*types.BlockTrace, error
 	pend.Wait()
 
 	// after all tx has been traced, collect "deletion proof" for zktrie
+	deleteionProofs := make(map[common.Hash]hexutil.Bytes)
+
 	for _, tracer := range env.ZkTrieTracer {
 		delProofs, err := tracer.GetDeletionProofs()
 		if err != nil {
 			log.Error("deletion proof failure", "error", err)
 		} else {
-			for _, proof := range delProofs {
+			for key, proof := range delProofs {
+				deleteionProofs[common.BytesToHash(key.Bytes())] = proof
 				env.DeletionProofs = append(env.DeletionProofs, proof)
 			}
 		}
 	}
 
+	//TODO: merge deletion proof
+	for k, v := range deleteionProofs {
+		env.FlattenProofs[k] = v
+	}
+
 	// build dummy per-tx deletion proof
 	for _, txStorageTrace := range env.TxStorageTraces {
 		if txStorageTrace != nil {
 			txStorageTrace.DeletionProofs = env.DeletionProofs
+			for k, v := range deleteionProofs {
+				txStorageTrace.FlattenProofs[k] = v
+			}
 		}
 	}
 
@@ -371,8 +392,11 @@ func (env *TraceEnv) getTxResult(state *state.StateDB, index int, block *types.B
 	}
 
 	txStorageTrace := &types.StorageTrace{
-		Proofs:        make(map[string][]hexutil.Bytes),
-		StorageProofs: make(map[string]map[string][]hexutil.Bytes),
+		Proofs:         make(map[string][]hexutil.Bytes),
+		StorageProofs:  make(map[string]map[string][]hexutil.Bytes),
+		FlattenProofs:  make(map[common.Hash]hexutil.Bytes),
+		AddressHashes:  make(map[common.Address]common.Hash),
+		StoreKeyHashes: make(map[common.Hash]common.Hash),
 	}
 	// still we have no state root for per tx, only set the head and tail
 	if index == 0 {
@@ -406,13 +430,17 @@ func (env *TraceEnv) getTxResult(state *state.StateDB, index int, block *types.B
 		if existed {
 			continue
 		}
-		proof, err := state.GetProof(addr)
+		proof, addrHash, err := state.GetFullProof(addr)
 		if err != nil {
 			log.Error("Proof not available", "address", addrStr, "error", err)
 			// but we still mark the proofs map with nil array
 		}
-		wrappedProof := types.WrapProof(proof)
+		wrappedProof := types.WrapProof(proof.GetData())
 		env.pMu.Lock()
+		// TODO:
+		env.fillFlattenStorageProof(txStorageTrace, proof)
+		txStorageTrace.AddressHashes[addr] = addrHash
+		env.AddressHashes[addr] = addrHash
 		env.Proofs[addrStr] = wrappedProof
 		txStorageTrace.Proofs[addrStr] = wrappedProof
 		env.pMu.Unlock()
@@ -464,18 +492,27 @@ func (env *TraceEnv) getTxResult(state *state.StateDB, index int, block *types.B
 			}
 			env.sMu.Unlock()
 
-			var proof [][]byte
+			var proof proofList
+			var keyHash common.Hash
 			var err error
 			if zktrieTracer.Available() {
-				proof, err = state.GetSecureTrieProof(zktrieTracer, key)
+				proof, keyHash, err = state.GetSecureTrieProof(zktrieTracer, key)
 			} else {
-				proof, err = state.GetSecureTrieProof(trie, key)
+				proof, keyHash, err = state.GetSecureTrieProof(trie, key)
 			}
 			if err != nil {
 				log.Error("Storage proof not available", "error", err, "address", addrStr, "key", keyStr)
 				// but we still mark the proofs map with nil array
 			}
-			wrappedProof := types.WrapProof(proof)
+
+			env.pMu.Lock()
+			// TODO:
+			env.fillFlattenStorageProof(txStorageTrace, proof)
+			txStorageTrace.StoreKeyHashes[key] = keyHash
+			env.StoreKeyHashes[key] = keyHash
+			env.pMu.Unlock()
+
+			wrappedProof := types.WrapProof(proof.GetData())
 			env.sMu.Lock()
 			txm[keyStr] = wrappedProof
 			m[keyStr] = wrappedProof
@@ -514,6 +551,18 @@ func (env *TraceEnv) getTxResult(state *state.StateDB, index int, block *types.B
 	return nil
 }
 
+func (env *TraceEnv) fillFlattenStorageProof(trace *types.StorageTrace, proof proofList) {
+	for _, i := range proof {
+		// the "raw key" is in fact a zktrie.Hash (bytes stored with little-endian)
+		// we need to convert it into big-endian
+		hash := common.BytesToHash(zktrie.NewHashFromBytes(i.Key)[:])
+		env.FlattenProofs[hash] = i.Value
+		if trace != nil {
+			trace.FlattenProofs[hash] = i.Value
+		}
+	}
+}
+
 // fillBlockTrace content after all the txs are finished running.
 func (env *TraceEnv) fillBlockTrace(block *types.Block) (*types.BlockTrace, error) {
 	defer func(t time.Time) {
@@ -542,10 +591,13 @@ func (env *TraceEnv) fillBlockTrace(block *types.Block) (*types.BlockTrace, erro
 
 	for addr, storages := range intrinsicStorageProofs {
 		if _, existed := env.Proofs[addr.String()]; !existed {
-			if proof, err := statedb.GetProof(addr); err != nil {
+			if proof, addrHash, err := statedb.GetFullProof(addr); err != nil {
 				log.Error("Proof for intrinstic address not available", "error", err, "address", addr)
 			} else {
-				env.Proofs[addr.String()] = types.WrapProof(proof)
+				// TODO:
+				env.fillFlattenStorageProof(nil, proof)
+				env.AddressHashes[addr] = addrHash
+				env.Proofs[addr.String()] = types.WrapProof(proof.GetData())
 			}
 		}
 
@@ -557,10 +609,13 @@ func (env *TraceEnv) fillBlockTrace(block *types.Block) (*types.BlockTrace, erro
 			if _, existed := env.StorageProofs[addr.String()][slot.String()]; !existed {
 				if trie, err := statedb.GetStorageTrieForProof(addr); err != nil {
 					log.Error("Storage proof for intrinstic address not available", "error", err, "address", addr)
-				} else if proof, err := statedb.GetSecureTrieProof(trie, slot); err != nil {
+				} else if proof, keyHash, err := statedb.GetSecureTrieProof(trie, slot); err != nil {
 					log.Error("Get storage proof for intrinstic address failed", "error", err, "address", addr, "slot", slot)
 				} else {
-					env.StorageProofs[addr.String()][slot.String()] = types.WrapProof(proof)
+					// TODO:
+					env.fillFlattenStorageProof(nil, proof)
+					env.StoreKeyHashes[slot] = keyHash
+					env.StorageProofs[addr.String()][slot.String()] = types.WrapProof(proof.GetData())
 				}
 			}
 		}