zktrie part3: include zktrie witness in block trace; add demo for generating witness data for mpt circuit (#123)

* demo for witness generator

* add simple binary for witness generating

* revert the witgen utility

* purge some code

* induce mptwitness argument

* handle case for failed tx, and special block in POA

* mptwitness: refactor for enabling output ops with different order

* induce witness gen with rwtable order

* self-documenting options for mptwitness

* fix trace flag according to review

* fix issue in handling failed tx

* refactor

* update according to reviews

Co-authored-by: HAOYUatHZ <haoyu@protonmail.com>

cmd/geth/config.go

@@ -35,6 +35,7 @@ import (
 	"github.com/scroll-tech/go-ethereum/cmd/utils"
 	"github.com/scroll-tech/go-ethereum/eth/catalyst"
 	"github.com/scroll-tech/go-ethereum/eth/ethconfig"
+	"github.com/scroll-tech/go-ethereum/internal/debug"
 	"github.com/scroll-tech/go-ethereum/internal/ethapi"
 	"github.com/scroll-tech/go-ethereum/log"
 	"github.com/scroll-tech/go-ethereum/metrics"
@@ -148,6 +149,7 @@ func makeConfigNode(ctx *cli.Context) (*node.Node, gethConfig) {
 	if ctx.GlobalIsSet(utils.EthStatsURLFlag.Name) {
 		cfg.Ethstats.URL = ctx.GlobalString(utils.EthStatsURLFlag.Name)
 	}
+	applyTraceConfig(ctx, &cfg.Eth)
 	applyMetricConfig(ctx, &cfg)
 
 	return stack, cfg
@@ -211,6 +213,12 @@ func dumpConfig(ctx *cli.Context) error {
 	return nil
 }
 
+func applyTraceConfig(ctx *cli.Context, cfg *ethconfig.Config) {
+	subCfg := debug.ConfigTrace(ctx)
+	cfg.TraceCacheLimit = subCfg.TraceCacheLimit
+	cfg.MPTWitness = subCfg.MPTWitness
+}
+
 func applyMetricConfig(ctx *cli.Context, cfg *gethConfig) {
 	if ctx.GlobalIsSet(utils.MetricsEnabledFlag.Name) {
 		cfg.Metrics.Enabled = ctx.GlobalBool(utils.MetricsEnabledFlag.Name)

cmd/geth/usage.go

@@ -57,12 +57,6 @@ var AppHelpFlagGroups = []flags.FlagGroup{
 			utils.WhitelistFlag,
 		},
 	},
-	{
-		Name: "NOTRACE CLIENT",
-		Flags: []cli.Flag{
-			utils.TraceCacheLimit,
-		},
-	},
 	{
 		Name: "LIGHT CLIENT",
 		Flags: []cli.Flag{

cmd/utils/flags.go

@@ -249,12 +249,6 @@ var (
 		Name:  "override.arrowglacier",
 		Usage: "Manually specify Arrow Glacier fork-block, overriding the bundled setting",
 	}
-	// NoTrace settings
-	TraceCacheLimit = cli.IntFlag{
-		Name:  "trace.limit",
-		Usage: "Handle the latest several blockResults",
-		Value: ethconfig.Defaults.TraceCacheLimit,
-	}
 	// Light server and client settings
 	LightServeFlag = cli.IntFlag{
 		Name:  "light.serve",
@@ -1100,13 +1094,6 @@ func MakeAddress(ks *keystore.KeyStore, account string) (accounts.Account, error
 	return accs[index], nil
 }
 
-func setTrace(ctx *cli.Context, cfg *ethconfig.Config) {
-	// NoTrace flag
-	if ctx.GlobalIsSet(TraceCacheLimit.Name) {
-		cfg.TraceCacheLimit = ctx.GlobalInt(TraceCacheLimit.Name)
-	}
-}
-
 // setEtherbase retrieves the etherbase either from the directly specified
 // command line flags or from the keystore if CLI indexed.
 func setEtherbase(ctx *cli.Context, ks *keystore.KeyStore, cfg *ethconfig.Config) {
@@ -1498,7 +1485,6 @@ func SetEthConfig(ctx *cli.Context, stack *node.Node, cfg *ethconfig.Config) {
 	if keystores := stack.AccountManager().Backends(keystore.KeyStoreType); len(keystores) > 0 {
 		ks = keystores[0].(*keystore.KeyStore)
 	}
-	setTrace(ctx, cfg)
 	setEtherbase(ctx, ks, cfg)
 	setGPO(ctx, &cfg.GPO, ctx.GlobalString(SyncModeFlag.Name) == "light")
 	setTxPool(ctx, &cfg.TxPool)

common/hexutil/hexutil.go

@@ -64,6 +64,11 @@ func Decode(input string) ([]byte, error) {
 	if !has0xPrefix(input) {
 		return nil, ErrMissingPrefix
 	}
+	// better compatible with odd size string
+	if len(input)%2 != 0 {
+		input = "0x0" + input[2:]
+	}
+
 	b, err := hex.DecodeString(input[2:])
 	if err != nil {
 		err = mapError(err)

core/blockchain.go

@@ -47,6 +47,7 @@ import (
 	"github.com/scroll-tech/go-ethereum/metrics"
 	"github.com/scroll-tech/go-ethereum/params"
 	"github.com/scroll-tech/go-ethereum/trie"
+	"github.com/scroll-tech/go-ethereum/trie/zkproof"
 )
 
 var (
@@ -134,6 +135,7 @@ type CacheConfig struct {
 	SnapshotLimit       int           // Memory allowance (MB) to use for caching snapshot entries in memory
 	Preimages           bool          // Whether to store preimage of trie key to the disk
 	TraceCacheLimit     int
+	MPTWitness          int // How to generate witness data for mpt circuit, 0: nothing, 1: natural
 
 	SnapshotWait bool // Wait for snapshot construction on startup. TODO(karalabe): This is a dirty hack for testing, nuke it
 }
@@ -147,6 +149,7 @@ var defaultCacheConfig = &CacheConfig{
 	SnapshotLimit:   256,
 	SnapshotWait:    true,
 	TraceCacheLimit: 32,
+	MPTWitness:      int(zkproof.MPTWitnessNothing),
 }
 
 // BlockChain represents the canonical chain given a database with a genesis
@@ -1370,7 +1373,6 @@ func (bc *BlockChain) writeBlockResult(state *state.StateDB, block *types.Block,
 	}
 
 	blockResult.BlockTrace = types.NewTraceBlock(bc.chainConfig, block, &coinbase)
-	blockResult.StorageTrace.RootAfter = state.GetRootHash()
 	for i, tx := range block.Transactions() {
 		evmTrace := blockResult.ExecutionResults[i]
 		// Contract is called
@@ -1383,6 +1385,11 @@ func (bc *BlockChain) writeBlockResult(state *state.StateDB, block *types.Block,
 			evmTrace.ByteCode = hexutil.Encode(tx.Data())
 		}
 	}
+
+	if err := zkproof.FillBlockResultForMPTWitness(zkproof.MPTWitnessType(bc.cacheConfig.MPTWitness), blockResult); err != nil {
+		log.Error("fill mpt witness fail", "error", err)
+	}
+
 	return blockResult
 }
 

core/types/l2trace.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"encoding/json"
 	"runtime"
 	"sync"
 
@@ -25,6 +26,7 @@ type BlockResult struct {
 	BlockTrace       *BlockTrace        `json:"blockTrace"`
 	StorageTrace     *StorageTrace      `json:"storageTrace"`
 	ExecutionResults []*ExecutionResult `json:"executionResults"`
+	MPTWitness       *json.RawMessage   `json:"mptwitness,omitempty"`
 }
 
 // StorageTrace stores proofs of storage needed by storage circuit

core/types/zktrie/byte32.go

@@ -18,6 +18,23 @@ func (b *Byte32) Hash() (*big.Int, error) {
 	}
 	return hash, nil
 }
+
+func (b *Byte32) Bytes() []byte { return b[:] }
+
+// same action as common.Hash (truncate bytes longer than 32 bytes FROM beginning,
+// and padding 0 at the beginning for shorter bytes)
+func NewByte32FromBytes(b []byte) *Byte32 {
+
+	byte32 := new(Byte32)
+
+	if len(b) > 32 {
+		b = b[len(b)-32:]
+	}
+
+	copy(byte32[32-len(b):], b)
+	return byte32
+}
+
 func NewByte32FromBytesPaddingZero(b []byte) *Byte32 {
 	if len(b) != 32 && len(b) != 20 {
 		panic(fmt.Errorf("do not support length except for 120bit and 256bit now. data: %v len: %v", b, len(b)))

core/vm/logger.go

@@ -318,7 +318,7 @@ func (l *StructLogger) CaptureEnter(typ OpCode, from common.Address, to common.A
 	theLog.getOrInitExtraData()
 	// handling additional updating for CALL/STATICCALL/CALLCODE/CREATE/CREATE2 only
 	// append extraData part for the log, capture the account status (the nonce / balance has been updated in capture enter)
-	wrappedStatus, _ := getWrappedAccountForAddr(l, to)
+	wrappedStatus := getWrappedAccountForAddr(l, to)
 	theLog.ExtraData.StateList = append(theLog.ExtraData.StateList, wrappedStatus)
 }
 
@@ -347,7 +347,7 @@ func (l *StructLogger) CaptureExit(output []byte, gasUsed uint64, err error) {
 		}
 
 		lastAccData := theLog.ExtraData.StateList[dataLen-1]
-		wrappedStatus, _ := getWrappedAccountForAddr(l, lastAccData.Address)
+		wrappedStatus := getWrappedAccountForAddr(l, lastAccData.Address)
 		theLog.ExtraData.StateList = append(theLog.ExtraData.StateList, wrappedStatus)
 	default:
 		//do nothing for other op code

core/vm/logger_trace.go

@@ -56,22 +56,20 @@ func traceStorage(l *StructLogger, scope *ScopeContext, extraData *types.ExtraDa
 		return nil
 	}
 	key := common.Hash(scope.Stack.peek().Bytes32())
-	storage, err := getWrappedAccountForStorage(l, scope.Contract.Address(), key)
+	storage := getWrappedAccountForStorage(l, scope.Contract.Address(), key)
-	if err == nil {
+	extraData.StateList = append(extraData.StateList, storage)
-		extraData.StateList = append(extraData.StateList, storage)
+
-	}
+	return nil
-	return err
 }
 
 // traceContractAccount gets the contract's account
 func traceContractAccount(l *StructLogger, scope *ScopeContext, extraData *types.ExtraData) error {
 	// Get account state.
-	state, err := getWrappedAccountForAddr(l, scope.Contract.Address())
+	state := getWrappedAccountForAddr(l, scope.Contract.Address())
-	if err == nil {
+	extraData.StateList = append(extraData.StateList, state)
-		extraData.StateList = append(extraData.StateList, state)
+	l.statesAffected[scope.Contract.Address()] = struct{}{}
-		l.statesAffected[scope.Contract.Address()] = struct{}{}
+
-	}
+	return nil
-	return err
 }
 
 // traceLastNAddressAccount returns func about the last N's address account.
@@ -83,37 +81,36 @@ func traceLastNAddressAccount(n int) traceFunc {
 		}
 
 		address := common.Address(stack.data[stack.len()-1-n].Bytes20())
-		state, err := getWrappedAccountForAddr(l, address)
+		state := getWrappedAccountForAddr(l, address)
-		if err == nil {
+		extraData.StateList = append(extraData.StateList, state)
-			extraData.StateList = append(extraData.StateList, state)
+		l.statesAffected[scope.Contract.Address()] = struct{}{}
-			l.statesAffected[scope.Contract.Address()] = struct{}{}
+
-		}
+		return nil
-		return err
 	}
 }
 
 // traceCaller gets caller address's account.
 func traceCaller(l *StructLogger, scope *ScopeContext, extraData *types.ExtraData) error {
 	address := scope.Contract.CallerAddress
-	state, err := getWrappedAccountForAddr(l, address)
+	state := getWrappedAccountForAddr(l, address)
-	if err == nil {
+
-		extraData.StateList = append(extraData.StateList, state)
+	extraData.StateList = append(extraData.StateList, state)
-		l.statesAffected[scope.Contract.Address()] = struct{}{}
+	l.statesAffected[scope.Contract.Address()] = struct{}{}
-	}
+
-	return err
+	return nil
 }
 
 // StorageWrapper will be empty
-func getWrappedAccountForAddr(l *StructLogger, address common.Address) (*types.AccountWrapper, error) {
+func getWrappedAccountForAddr(l *StructLogger, address common.Address) *types.AccountWrapper {
 	return &types.AccountWrapper{
 		Address:  address,
 		Nonce:    l.env.StateDB.GetNonce(address),
 		Balance:  (*hexutil.Big)(l.env.StateDB.GetBalance(address)),
 		CodeHash: l.env.StateDB.GetCodeHash(address),
-	}, nil
+	}
 }
 
-func getWrappedAccountForStorage(l *StructLogger, address common.Address, key common.Hash) (*types.AccountWrapper, error) {
+func getWrappedAccountForStorage(l *StructLogger, address common.Address, key common.Hash) *types.AccountWrapper {
 	return &types.AccountWrapper{
 		Address:  address,
 		Nonce:    l.env.StateDB.GetNonce(address),
@@ -123,5 +120,5 @@ func getWrappedAccountForStorage(l *StructLogger, address common.Address, key co
 			Key:   key.String(),
 			Value: l.env.StateDB.GetState(address, key).String(),
 		},
-	}, nil
+	}
 }

eth/backend.go

@@ -189,6 +189,7 @@ func New(stack *node.Node, config *ethconfig.Config) (*Ethereum, error) {
 			SnapshotLimit:       config.SnapshotCache,
 			Preimages:           config.Preimages,
 			TraceCacheLimit:     config.TraceCacheLimit,
+			MPTWitness:          config.MPTWitness,
 		}
 	)
 	eth.blockchain, err = core.NewBlockChain(chainDb, cacheConfig, chainConfig, eth.engine, vmConfig, eth.shouldPreserve, &config.TxLookupLimit)

eth/ethconfig/config.go

@@ -87,12 +87,11 @@ var Defaults = Config{
 		GasPrice: big.NewInt(params.GWei),
 		Recommit: 3 * time.Second,
 	},
-	TxPool:          core.DefaultTxPoolConfig,
+	TxPool:        core.DefaultTxPoolConfig,
-	RPCGasCap:       50000000,
+	RPCGasCap:     50000000,
-	RPCEVMTimeout:   5 * time.Second,
+	RPCEVMTimeout: 5 * time.Second,
-	GPO:             FullNodeGPO,
+	GPO:           FullNodeGPO,
-	RPCTxFeeCap:     1, // 1 ether
+	RPCTxFeeCap:   1, // 1 ether
-	TraceCacheLimit: 32,
 }
 
 func init() {
@@ -208,6 +207,7 @@ type Config struct {
 
 	// Trace option
 	TraceCacheLimit int
+	MPTWitness      int
 }
 
 // CreateConsensusEngine creates a consensus engine for the given chain configuration.

internal/debug/flags.go

@@ -91,6 +91,18 @@ var (
 		Name:  "trace",
 		Usage: "Write execution trace to the given file",
 	}
+	// NoTrace settings
+	traceCacheLimitFlag = cli.IntFlag{
+		Name:  "trace.limit",
+		Usage: "Handle the latest several blockResults",
+		Value: 32,
+	}
+	// mpt witness settings
+	mptWitnessFlag = cli.IntFlag{
+		Name:  "trace.mptwitness",
+		Usage: "Output witness for mpt circuit with Specified order (default = no output, 1 = by executing order",
+		Value: 0,
+	}
 )
 
 // Flags holds all command-line flags required for debugging.
@@ -107,6 +119,8 @@ var Flags = []cli.Flag{
 	blockprofilerateFlag,
 	cpuprofileFlag,
 	traceFlag,
+	traceCacheLimitFlag,
+	mptWitnessFlag,
 }
 
 var glogger *log.GlogHandler
@@ -117,6 +131,23 @@ func init() {
 	log.Root().SetHandler(glogger)
 }
 
+// TraceConfig export options about trace
+type TraceConfig struct {
+	TracePath string
+	// Trace option
+	TraceCacheLimit int
+	MPTWitness      int
+}
+
+func ConfigTrace(ctx *cli.Context) *TraceConfig {
+	cfg := new(TraceConfig)
+	cfg.TracePath = ctx.GlobalString(traceFlag.Name)
+	cfg.TraceCacheLimit = ctx.GlobalInt(traceCacheLimitFlag.Name)
+	cfg.MPTWitness = ctx.GlobalInt(mptWitnessFlag.Name)
+
+	return cfg
+}
+
 // Setup initializes profiling and logging based on the CLI flags.
 // It should be called as early as possible in the program.
 func Setup(ctx *cli.Context) error {

miner/worker.go

@@ -1163,6 +1163,7 @@ func (w *worker) commit(uncles []*types.Header, interval func(), update bool, st
 	// Deep copy receipts here to avoid interaction between different tasks.
 	receipts := copyReceipts(w.current.receipts)
 	s := w.current.state.Copy()
+	// when block is not mined by myself, we just omit
 	// complete storage before Finalize state (only RootAfter left unknown)
 	storage := &types.StorageTrace{
 		RootBefore:    s.GetRootHash(),
@@ -1174,6 +1175,8 @@ func (w *worker) commit(uncles []*types.Header, interval func(), update bool, st
 	if err != nil {
 		return err
 	}
+	storage.RootAfter = block.Header().Root
+
 	if w.isRunning() {
 		if interval != nil {
 			interval()

trie/zkproof/orderer.go

@@ -0,0 +1,229 @@
+package zkproof
+
+import (
+	"math/big"
+	"sort"
+
+	"github.com/scroll-tech/go-ethereum/common"
+	"github.com/scroll-tech/go-ethereum/common/hexutil"
+	"github.com/scroll-tech/go-ethereum/core/types"
+)
+
+type opIterator interface {
+	next() *types.AccountWrapper
+}
+
+type opOrderer interface {
+	absorb(*types.AccountWrapper)
+	end_absorb() opIterator
+}
+
+type iterateOp []*types.AccountWrapper
+
+func (ops *iterateOp) next() *types.AccountWrapper {
+
+	sl := *ops
+
+	if len(sl) == 0 {
+		return nil
+	}
+
+	*ops = sl[1:]
+	return sl[0]
+}
+
+type simpleOrderer struct {
+	savedOp []*types.AccountWrapper
+}
+
+func (od *simpleOrderer) absorb(st *types.AccountWrapper) {
+	od.savedOp = append(od.savedOp, st)
+}
+
+func (od *simpleOrderer) end_absorb() opIterator {
+	ret := iterateOp(od.savedOp)
+	return &ret
+}
+
+type multiOpIterator []opIterator
+
+func (opss *multiOpIterator) next() *types.AccountWrapper {
+
+	sl := *opss
+	if len(sl) == 0 {
+		return nil
+	}
+
+	op := sl[0].next()
+
+	for op == nil {
+
+		sl = sl[1:]
+		*opss = sl
+		if len(sl) == 0 {
+			return nil
+		}
+		op = sl[0].next()
+	}
+	return op
+}
+
+type rwTblOrderer struct {
+	traced map[string]*types.AccountWrapper
+
+	opAccNonce    map[string]*types.AccountWrapper
+	opAccBalance  map[string]*types.AccountWrapper
+	opAccCodeHash map[string]*types.AccountWrapper
+	opStorage     map[string]map[string]*types.StorageWrapper
+}
+
+func newRWTblOrderer(inited map[common.Address]*types.StateAccount) *rwTblOrderer {
+
+	traced := make(map[string]*types.AccountWrapper)
+
+	for addr, data := range inited {
+		if data == nil {
+			continue
+		}
+
+		bl := data.Balance
+		if bl == nil {
+			bl = big.NewInt(0)
+		}
+
+		traced[addr.String()] = &types.AccountWrapper{
+			Address:  addr,
+			Nonce:    data.Nonce,
+			Balance:  (*hexutil.Big)(bl),
+			CodeHash: common.BytesToHash(data.CodeHash),
+		}
+	}
+
+	return &rwTblOrderer{
+		traced:        traced,
+		opAccNonce:    make(map[string]*types.AccountWrapper),
+		opAccBalance:  make(map[string]*types.AccountWrapper),
+		opAccCodeHash: make(map[string]*types.AccountWrapper),
+		opStorage:     make(map[string]map[string]*types.StorageWrapper),
+	}
+}
+
+func (od *rwTblOrderer) absorb(st *types.AccountWrapper) {
+
+	addrStr := st.Address.String()
+
+	start, existed := od.traced[addrStr]
+	if !existed {
+		start = &types.AccountWrapper{
+			Balance: (*hexutil.Big)(big.NewInt(0)),
+		}
+	}
+	od.traced[addrStr] = st
+
+	// notice there would be at least one entry for all 3 fields when accessing an address
+	// this may caused extract "read" op in mpt circuit which has no corresponding one in rwtable
+	// we can avoid it unless obtaining more tips from the understanding of opcode
+	// but it would be ok if we have adopted the new lookup way (root_prev, root_cur) under discussion:
+	// https://github.com/privacy-scaling-explorations/zkevm-specs/issues/217
+
+	if traced, existed := od.opAccNonce[addrStr]; !existed {
+		traced = copyAccountState(st)
+		traced.Balance = start.Balance
+		traced.CodeHash = start.CodeHash
+		traced.Storage = nil
+		od.opAccNonce[addrStr] = traced
+	} else {
+		traced.Nonce = st.Nonce
+	}
+
+	if traced, existed := od.opAccBalance[addrStr]; !existed {
+		traced = copyAccountState(st)
+		traced.CodeHash = start.CodeHash
+		traced.Storage = nil
+		od.opAccBalance[addrStr] = traced
+	} else {
+		traced.Nonce = st.Nonce
+		traced.Balance = st.Balance
+	}
+
+	if traced, existed := od.opAccCodeHash[addrStr]; !existed {
+		traced = copyAccountState(st)
+		traced.Storage = nil
+		od.opAccCodeHash[addrStr] = traced
+	} else {
+		traced.Nonce = st.Nonce
+		traced.Balance = st.Balance
+		traced.CodeHash = st.CodeHash
+	}
+
+	if stg := st.Storage; stg != nil {
+		m, existed := od.opStorage[addrStr]
+		if !existed {
+			m = make(map[string]*types.StorageWrapper)
+			od.opStorage[addrStr] = m
+		}
+
+		// key must be unified into 32 bytes
+		keyBytes := hexutil.MustDecode(stg.Key)
+		m[common.BytesToHash(keyBytes).String()] = stg
+	}
+}
+
+func (od *rwTblOrderer) end_absorb() opIterator {
+	// now sort every map by address / key
+	// inited has collected all address, just sort address once
+	sortedAddrs := make([]string, 0, len(od.traced))
+	for addrs := range od.traced {
+		sortedAddrs = append(sortedAddrs, addrs)
+	}
+	sort.Strings(sortedAddrs)
+
+	var iterNonce []*types.AccountWrapper
+	var iterBalance []*types.AccountWrapper
+	var iterCodeHash []*types.AccountWrapper
+	var iterStorage []*types.AccountWrapper
+
+	for _, addrStr := range sortedAddrs {
+		if v, existed := od.opAccNonce[addrStr]; existed {
+			iterNonce = append(iterNonce, v)
+		}
+
+		if v, existed := od.opAccBalance[addrStr]; existed {
+			iterBalance = append(iterBalance, v)
+		}
+
+		if v, existed := od.opAccCodeHash[addrStr]; existed {
+			iterCodeHash = append(iterCodeHash, v)
+		}
+
+		if stgM, existed := od.opStorage[addrStr]; existed {
+
+			tracedStatus := od.traced[addrStr]
+			if tracedStatus == nil {
+				panic("missed traced status found in storage slot")
+			}
+
+			sortedKeys := make([]string, 0, len(stgM))
+			for key := range stgM {
+				sortedKeys = append(sortedKeys, key)
+			}
+			sort.Strings(sortedKeys)
+
+			for _, key := range sortedKeys {
+				st := copyAccountState(tracedStatus)
+				st.Storage = stgM[key]
+				iterStorage = append(iterStorage, st)
+			}
+		}
+
+	}
+
+	var finalRet []opIterator
+	for _, arr := range [][]*types.AccountWrapper{iterNonce, iterBalance, iterCodeHash, iterStorage} {
+		wrappedIter := iterateOp(arr)
+		finalRet = append(finalRet, &wrappedIter)
+	}
+
+	wrappedRet := multiOpIterator(finalRet)
+	return &wrappedRet
+}

trie/zkproof/types.go

@@ -0,0 +1,57 @@
+package zkproof
+
+import (
+	"github.com/scroll-tech/go-ethereum/common/hexutil"
+)
+
+type MPTWitnessType int
+
+const (
+	MPTWitnessNothing MPTWitnessType = iota
+	MPTWitnessNatural
+	MPTWitnessRWTbl
+)
+
+// SMTPathNode represent a node in the SMT Path, all hash is saved by the present of
+// zktype.Hash
+type SMTPathNode struct {
+	Value   hexutil.Bytes `json:"value"`
+	Sibling hexutil.Bytes `json:"sibling"`
+}
+
+// SMTPath is the whole path of SMT
+type SMTPath struct {
+	KeyPathPart *hexutil.Big  `json:"pathPart"` //the path part in key
+	Root        hexutil.Bytes `json:"root"`
+	Path        []SMTPathNode `json:"path,omitempty"` //path start from top
+	Leaf        *SMTPathNode  `json:"leaf,omitempty"` //would be omitted for empty leaf, the sibling indicate key
+}
+
+// StateAccount is the represent of StateAccount in L2 circuit
+// Notice in L2 we have different hash scheme against StateAccount.MarshalByte
+type StateAccount struct {
+	Nonce    int           `json:"nonce"`
+	Balance  *hexutil.Big  `json:"balance"` //just the common hex expression of integer (big-endian)
+	CodeHash hexutil.Bytes `json:"codeHash,omitempty"`
+}
+
+// StateStorage is the represent of a stored key-value pair for specified account
+type StateStorage struct {
+	Key   hexutil.Bytes `json:"key"` //notice this is the preimage of storage key
+	Value hexutil.Bytes `json:"value"`
+}
+
+// StorageTrace record the updating on state trie and (if changed) account trie
+// represent by the [before, after] updating of SMTPath amont tries and Account
+type StorageTrace struct {
+	// which log the trace is responded for, -1 indicate not caused
+	// by opcode (like gasRefund, coinbase, setNonce, etc)
+	Address         hexutil.Bytes    `json:"address"`
+	AccountKey      hexutil.Bytes    `json:"accountKey"`
+	AccountPath     [2]*SMTPath      `json:"accountPath"`
+	AccountUpdate   [2]*StateAccount `json:"accountUpdate"`
+	StateKey        hexutil.Bytes    `json:"stateKey,omitempty"`
+	CommonStateRoot hexutil.Bytes    `json:"commonStateRoot,omitempty"` //CommonStateRoot is used if there is no update on state storage
+	StatePath       [2]*SMTPath      `json:"statePath,omitempty"`
+	StateUpdate     [2]*StateStorage `json:"stateUpdate,omitempty"`
+}

trie/zkproof/writer.go

@@ -0,0 +1,740 @@
+package zkproof
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"math/big"
+
+	"github.com/scroll-tech/go-ethereum/common"
+	"github.com/scroll-tech/go-ethereum/common/hexutil"
+	"github.com/scroll-tech/go-ethereum/core/types"
+	zkt "github.com/scroll-tech/go-ethereum/core/types/zktrie"
+	"github.com/scroll-tech/go-ethereum/ethdb/memorydb"
+	"github.com/scroll-tech/go-ethereum/log"
+	"github.com/scroll-tech/go-ethereum/trie"
+)
+
+type proofList [][]byte
+
+func (n *proofList) Put(key []byte, value []byte) error {
+	*n = append(*n, value)
+	return nil
+}
+
+func (n *proofList) Delete(key []byte) error {
+	panic("not supported")
+}
+
+func addressToKey(addr common.Address) *zkt.Hash {
+	var preImage zkt.Byte32
+	copy(preImage[:], addr.Bytes())
+
+	h, err := preImage.Hash()
+	if err != nil {
+		log.Error("hash failure", "preImage", hexutil.Encode(preImage[:]))
+		return nil
+	}
+	return zkt.NewHashFromBigInt(h)
+}
+
+//resume the proof bytes into db and return the leaf node
+func resumeProofs(proof []hexutil.Bytes, db *memorydb.Database) *trie.Node {
+	for _, buf := range proof {
+
+		n, err := trie.DecodeSMTProof(buf)
+		if err != nil {
+			log.Warn("decode proof string fail", "error", err)
+		} else if n != nil {
+			k, err := n.Key()
+			if err != nil {
+				log.Warn("node has no valid key", "error", err)
+			} else {
+				//notice: must consistent with trie/merkletree.go
+				bt := k[:]
+				db.Put(bt, buf)
+				if n.Type == trie.NodeTypeLeaf || n.Type == trie.NodeTypeEmpty {
+					return n
+				}
+			}
+		}
+
+	}
+
+	return nil
+}
+
+// we have a trick here which suppose the proof array include all middle nodes along the
+// whole path in sequence, from root to leaf
+func decodeProofForMPTPath(proof proofList, path *SMTPath) {
+
+	var lastNode *trie.Node
+	keyPath := big.NewInt(0)
+	path.KeyPathPart = (*hexutil.Big)(keyPath)
+
+	keyCounter := big.NewInt(1)
+
+	for _, buf := range proof {
+		n, err := trie.DecodeSMTProof(buf)
+		if err != nil {
+			log.Warn("decode proof string fail", "error", err)
+		} else if n != nil {
+			k, err := n.Key()
+			if err != nil {
+				log.Warn("node has no valid key", "error", err)
+				return
+			}
+			if lastNode == nil {
+				//notice: use little-endian represent inside Hash ([:] or Bytes2())
+				path.Root = k[:]
+			} else {
+				if bytes.Equal(k[:], lastNode.ChildL[:]) {
+					path.Path = append(path.Path, SMTPathNode{
+						Value:   k[:],
+						Sibling: lastNode.ChildR[:],
+					})
+				} else if bytes.Equal(k[:], lastNode.ChildR[:]) {
+					path.Path = append(path.Path, SMTPathNode{
+						Value:   k[:],
+						Sibling: lastNode.ChildL[:],
+					})
+					keyPath.Add(keyPath, keyCounter)
+				} else {
+					panic("Unexpected proof form")
+				}
+				keyCounter.Mul(keyCounter, big.NewInt(2))
+			}
+			switch n.Type {
+			case trie.NodeTypeMiddle:
+				lastNode = n
+			case trie.NodeTypeLeaf:
+				vhash, _ := n.ValueKey()
+				path.Leaf = &SMTPathNode{
+					//here we just return the inner represent of hash (little endian, reversed byte order to common hash)
+					Value:   vhash[:],
+					Sibling: n.NodeKey[:],
+				}
+				//sanity check
+				keyPart := keyPath.Bytes()
+				for i, b := range keyPart {
+					ri := len(keyPart) - i
+					cb := path.Leaf.Sibling[ri-1] //notice the output is little-endian
+					if b&cb != b {
+						panic(fmt.Errorf("path key not match: part is %x but key is %x", keyPart, []byte(path.Leaf.Sibling[:])))
+					}
+				}
+
+				return
+			case trie.NodeTypeEmpty:
+				return
+			default:
+				panic(fmt.Errorf("unknown node type %d", n.Type))
+			}
+		}
+	}
+
+	panic("Unexpected finished here")
+}
+
+type zktrieProofWriter struct {
+	db                  *trie.ZktrieDatabase
+	tracingZktrie       *trie.ZkTrie
+	tracingStorageTries map[common.Address]*trie.ZkTrie
+	tracingAccounts     map[common.Address]*types.StateAccount
+}
+
+func NewZkTrieProofWriter(storage *types.StorageTrace) (*zktrieProofWriter, error) {
+
+	underlayerDb := memorydb.New()
+	zkDb := trie.NewZktrieDatabase(underlayerDb)
+
+	accounts := make(map[common.Address]*types.StateAccount)
+
+	// resuming proof bytes to underlayerDb
+	for addrs, proof := range storage.Proofs {
+		if n := resumeProofs(proof, underlayerDb); n != nil {
+			addr := common.HexToAddress(addrs)
+			if n.Type == trie.NodeTypeEmpty {
+				accounts[addr] = nil
+			} else if acc, err := types.UnmarshalStateAccount(n.Data()); err == nil {
+				if bytes.Equal(n.NodeKey[:], addressToKey(addr)[:]) {
+					accounts[addr] = acc
+				} else {
+					// should still mark the address as being trace (data not existed yet)
+					accounts[addr] = nil
+				}
+
+			} else {
+				return nil, fmt.Errorf("decode account bytes fail: %s, raw data [%x]", err, n.Data())
+			}
+
+		} else {
+			return nil, fmt.Errorf("can not resume proof for address %s", addrs)
+		}
+	}
+
+	storages := make(map[common.Address]*trie.ZkTrie)
+
+	for addrs, stgLists := range storage.StorageProofs {
+
+		addr := common.HexToAddress(addrs)
+		accState, existed := accounts[addr]
+		if !existed {
+			// trace is malformed but currently we just warn about that
+			log.Warn("no account state found for this addr, mal records", "address", addrs)
+			continue
+		} else if accState == nil {
+			// create an empty zktrie for uninit address
+			storages[addr], _ = trie.NewZkTrie(common.Hash{}, zkDb)
+			continue
+		}
+
+		for keys, proof := range stgLists {
+
+			if n := resumeProofs(proof, underlayerDb); n != nil {
+				var err error
+				storages[addr], err = trie.NewZkTrie(accState.Root, zkDb)
+				if err != nil {
+					return nil, fmt.Errorf("zktrie create failure for storage in addr <%s>: %s", err, addrs)
+				}
+
+			} else {
+				return nil, fmt.Errorf("can not resume proof for storage %s@%s", keys, addrs)
+			}
+
+		}
+	}
+
+	zktrie, err := trie.NewZkTrie(
+		storage.RootBefore,
+		trie.NewZktrieDatabase(underlayerDb),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("zktrie create failure: %s", err)
+	}
+
+	// sanity check
+	if !bytes.Equal(zktrie.Hash().Bytes(), storage.RootBefore.Bytes()) {
+		return nil, fmt.Errorf("unmatch init trie hash: expected %x but has %x", storage.RootBefore.Bytes(), zktrie.Hash().Bytes())
+	}
+
+	return &zktrieProofWriter{
+		db:                  zkDb,
+		tracingZktrie:       zktrie,
+		tracingAccounts:     accounts,
+		tracingStorageTries: storages,
+	}, nil
+}
+
+const (
+	posSSTOREBefore = 0
+	posCREATE       = 0
+	posCREATEAfter  = 1
+	posCALL         = 2
+	posSTATICCALL   = 0
+)
+
+func getAccountState(l *types.StructLogRes, pos int) *types.AccountWrapper {
+	if exData := l.ExtraData; exData == nil {
+		return nil
+	} else if len(exData.StateList) < pos {
+		return nil
+	} else {
+		return exData.StateList[pos]
+	}
+}
+
+func copyAccountState(st *types.AccountWrapper) *types.AccountWrapper {
+
+	var stg *types.StorageWrapper
+	if st.Storage != nil {
+		stg = &types.StorageWrapper{
+			Key:   st.Storage.Key,
+			Value: st.Storage.Value,
+		}
+	}
+
+	return &types.AccountWrapper{
+		Nonce:    st.Nonce,
+		Balance:  (*hexutil.Big)(big.NewInt(0).Set(st.Balance.ToInt())),
+		CodeHash: st.CodeHash,
+		Address:  st.Address,
+		Storage:  stg,
+	}
+}
+
+func getAccountDataFromLogState(state *types.AccountWrapper) *types.StateAccount {
+	return &types.StateAccount{
+		Nonce:    state.Nonce,
+		Balance:  (*big.Int)(state.Balance),
+		CodeHash: state.CodeHash.Bytes(),
+	}
+}
+
+// for sanity check
+func verifyAccount(addr common.Address, data *types.StateAccount, leaf *SMTPathNode) error {
+
+	if leaf == nil {
+		if data != nil {
+			return fmt.Errorf("path has no corresponding leaf for account")
+		} else {
+			return nil
+		}
+	}
+
+	addrKey := addressToKey(addr)
+	if !bytes.Equal(addrKey[:], leaf.Sibling) {
+		if data != nil {
+			return fmt.Errorf("unmatch leaf node in address: %s", addr)
+		}
+	} else if data != nil {
+		h, err := data.Hash()
+		//log.Info("sanity check acc before", "addr", addr.String(), "key", leaf.Sibling.Text(16), "hash", h.Text(16))
+
+		if err != nil {
+			return fmt.Errorf("fail to hash account: %v", err)
+		}
+		if !bytes.Equal(zkt.NewHashFromBigInt(h)[:], leaf.Value) {
+			return fmt.Errorf("unmatch data in leaf for address %s", addr)
+		}
+	}
+	return nil
+}
+
+// for sanity check
+func verifyStorage(key *zkt.Byte32, data *zkt.Byte32, leaf *SMTPathNode) error {
+
+	emptyData := bytes.Equal(data[:], common.Hash{}.Bytes())
+
+	if leaf == nil {
+		if !emptyData {
+			return fmt.Errorf("path has no corresponding leaf for storage")
+		} else {
+			return nil
+		}
+	}
+
+	keyHash, err := key.Hash()
+	if err != nil {
+		return err
+	}
+
+	if !bytes.Equal(zkt.NewHashFromBigInt(keyHash)[:], leaf.Sibling) {
+		if !emptyData {
+			return fmt.Errorf("unmatch leaf node in storage: %x", key[:])
+		}
+	} else {
+		h, err := data.Hash()
+		//log.Info("sanity check acc before", "addr", addr.String(), "key", leaf.Sibling.Text(16), "hash", h.Text(16))
+
+		if err != nil {
+			return fmt.Errorf("fail to hash data: %v", err)
+		}
+		if !bytes.Equal(zkt.NewHashFromBigInt(h)[:], leaf.Value) {
+			return fmt.Errorf("unmatch data in leaf for storage %x", key[:])
+		}
+	}
+	return nil
+}
+
+// update traced account state, and return the corresponding trace object which
+// is still opened for more infos
+// the updated accData state is obtained by a closure which enable it being derived from current status
+func (w *zktrieProofWriter) traceAccountUpdate(addr common.Address, updateAccData func(*types.StateAccount) *types.StateAccount) (*StorageTrace, error) {
+
+	out := new(StorageTrace)
+	//account trie
+	out.Address = addr.Bytes()
+	out.AccountPath = [2]*SMTPath{{}, {}}
+	//fill dummy
+	out.AccountUpdate = [2]*StateAccount{}
+
+	accDataBefore, existed := w.tracingAccounts[addr]
+	if !existed {
+		//sanity check
+		panic(fmt.Errorf("code do not add initialized status for account %s", addr))
+	}
+
+	var proof proofList
+	if err := w.tracingZktrie.Prove(addr.Bytes32(), 0, &proof); err != nil {
+		return nil, fmt.Errorf("prove BEFORE state fail: %s", err)
+	}
+
+	decodeProofForMPTPath(proof, out.AccountPath[0])
+	if err := verifyAccount(addr, accDataBefore, out.AccountPath[0].Leaf); err != nil {
+		panic(fmt.Errorf("code fail to trace account status correctly: %s", err))
+	}
+	if accDataBefore != nil {
+		// we have ensured the nBefore has a key corresponding to the query one
+		out.AccountKey = out.AccountPath[0].Leaf.Sibling
+		out.AccountUpdate[0] = &StateAccount{
+			Nonce:    int(accDataBefore.Nonce),
+			Balance:  (*hexutil.Big)(big.NewInt(0).Set(accDataBefore.Balance)),
+			CodeHash: accDataBefore.CodeHash,
+		}
+	}
+
+	accData := updateAccData(accDataBefore)
+	if accData != nil {
+		out.AccountUpdate[1] = &StateAccount{
+			Nonce:    int(accData.Nonce),
+			Balance:  (*hexutil.Big)(big.NewInt(0).Set(accData.Balance)),
+			CodeHash: accData.CodeHash,
+		}
+	}
+
+	if accData != nil {
+		if err := w.tracingZktrie.TryUpdateAccount(addr.Bytes32(), accData); err != nil {
+			return nil, fmt.Errorf("update zktrie account state fail: %s", err)
+		}
+		w.tracingAccounts[addr] = accData
+	} else {
+		if err := w.tracingZktrie.TryDelete(addr.Bytes32()); err != nil {
+			return nil, fmt.Errorf("delete zktrie account state fail: %s", err)
+		}
+		delete(w.tracingAccounts, addr)
+	}
+
+	proof = proofList{}
+	if err := w.tracingZktrie.Prove(addr.Bytes32(), 0, &proof); err != nil {
+		return nil, fmt.Errorf("prove AFTER state fail: %s", err)
+	}
+
+	decodeProofForMPTPath(proof, out.AccountPath[1])
+	if err := verifyAccount(addr, accData, out.AccountPath[1].Leaf); err != nil {
+		panic(fmt.Errorf("state AFTER has no valid account: %s", err))
+	}
+	if accData != nil {
+		if out.AccountKey == nil {
+			out.AccountKey = out.AccountPath[1].Leaf.Sibling[:]
+		}
+		//now accountKey must has been filled
+	}
+
+	return out, nil
+}
+
+// update traced storage state, and return the corresponding trace object
+func (w *zktrieProofWriter) traceStorageUpdate(addr common.Address, key, value []byte) (*StorageTrace, error) {
+
+	trie := w.tracingStorageTries[addr]
+	if trie == nil {
+		return nil, fmt.Errorf("no trace storage trie for %s", addr)
+	}
+
+	statePath := [2]*SMTPath{{}, {}}
+	stateUpdate := [2]*StateStorage{}
+
+	storeKey := zkt.NewByte32FromBytesPaddingZero(common.BytesToHash(key).Bytes())
+	storeValueBefore := trie.Get(storeKey[:])
+	storeValue := zkt.NewByte32FromBytes(value)
+
+	if storeValueBefore != nil && !bytes.Equal(storeValueBefore[:], common.Hash{}.Bytes()) {
+		stateUpdate[0] = &StateStorage{
+			Key:   storeKey.Bytes(),
+			Value: storeValueBefore,
+		}
+	}
+
+	var storageBeforeProof, storageAfterProof proofList
+	if err := trie.Prove(storeKey.Bytes(), 0, &storageBeforeProof); err != nil {
+		return nil, fmt.Errorf("prove BEFORE storage state fail: %s", err)
+	}
+
+	decodeProofForMPTPath(storageBeforeProof, statePath[0])
+	if err := verifyStorage(storeKey, zkt.NewByte32FromBytes(storeValueBefore), statePath[0].Leaf); err != nil {
+		panic(fmt.Errorf("storage BEFORE has no valid data: %s (%v)", err, statePath[0]))
+	}
+
+	if !bytes.Equal(storeValue.Bytes(), common.Hash{}.Bytes()) {
+		if err := trie.TryUpdate(storeKey.Bytes(), storeValue.Bytes()); err != nil {
+			return nil, fmt.Errorf("update zktrie storage fail: %s", err)
+		}
+		stateUpdate[1] = &StateStorage{
+			Key:   storeKey.Bytes(),
+			Value: storeValue.Bytes(),
+		}
+	} else {
+		if err := trie.TryDelete(storeKey.Bytes()); err != nil {
+			return nil, fmt.Errorf("delete zktrie storage fail: %s", err)
+		}
+	}
+
+	if err := trie.Prove(storeKey.Bytes(), 0, &storageAfterProof); err != nil {
+		return nil, fmt.Errorf("prove AFTER storage state fail: %s", err)
+	}
+	decodeProofForMPTPath(storageAfterProof, statePath[1])
+	if err := verifyStorage(storeKey, storeValue, statePath[1].Leaf); err != nil {
+		panic(fmt.Errorf("storage AFTER has no valid data: %s (%v)", err, statePath[1]))
+	}
+
+	out, err := w.traceAccountUpdate(addr,
+		func(acc *types.StateAccount) *types.StateAccount {
+			if acc == nil {
+				panic("unexpected")
+			}
+
+			//sanity check
+			if accRootFromState := zkt.ReverseByteOrder(statePath[0].Root); !bytes.Equal(acc.Root[:], accRootFromState) {
+				panic(fmt.Errorf("unexpected storage root before: [%s] vs [%x]", acc.Root, accRootFromState))
+			}
+			return &types.StateAccount{
+				Nonce:    acc.Nonce,
+				Balance:  acc.Balance,
+				CodeHash: acc.CodeHash,
+				Root:     common.BytesToHash(zkt.ReverseByteOrder(statePath[1].Root)),
+			}
+		})
+	if err != nil {
+		return nil, fmt.Errorf("update account %s in SSTORE fail: %s", addr, err)
+	}
+
+	if stateUpdate[1] != nil {
+		out.StateKey = statePath[1].Leaf.Sibling
+	} else if stateUpdate[0] != nil {
+		out.StateKey = statePath[0].Leaf.Sibling
+	} else {
+		// it occurs when we are handling SLOAD with non-exist value
+		// still no pretty idea, had to touch the internal behavior in zktrie ....
+		if h, err := storeKey.Hash(); err != nil {
+			return nil, fmt.Errorf("hash storekey fail: %s", err)
+		} else {
+			out.StateKey = zkt.NewHashFromBigInt(h)[:]
+		}
+	}
+
+	out.StatePath = statePath
+	out.StateUpdate = stateUpdate
+	return out, nil
+}
+
+func (w *zktrieProofWriter) HandleNewState(accountState *types.AccountWrapper) (*StorageTrace, error) {
+
+	if accountState.Storage != nil {
+		storeAddr := hexutil.MustDecode(accountState.Storage.Key)
+		storeValue := hexutil.MustDecode(accountState.Storage.Value)
+		return w.traceStorageUpdate(accountState.Address, storeAddr, storeValue)
+	} else {
+
+		accData := getAccountDataFromLogState(accountState)
+
+		out, err := w.traceAccountUpdate(accountState.Address, func(accBefore *types.StateAccount) *types.StateAccount {
+			if accBefore != nil {
+				accData.Root = accBefore.Root
+			}
+			return accData
+		})
+		if err != nil {
+			return nil, fmt.Errorf("update account state %s fail: %s", accountState.Address, err)
+		}
+		hash, err := zkt.NewHashFromBytes(accData.Root[:])
+		if err != nil {
+			return nil, fmt.Errorf("malform of state root in account %s", accountState.Address)
+		}
+		out.CommonStateRoot = hash[:]
+		return out, nil
+	}
+
+}
+
+func handleLogs(od opOrderer, currentContract common.Address, logs []*types.StructLogRes) {
+	logStack := []int{0}
+	contractStack := map[int]common.Address{}
+	skipDepth := 0
+	callEnterAddress := currentContract
+
+	// now trace every OP which could cause changes on state:
+	for i, sLog := range logs {
+
+		//trace log stack by depth rather than scanning specified op
+		if sl := len(logStack); sl < sLog.Depth {
+			logStack = append(logStack, i)
+			//update currentContract according to previous op
+			contractStack[sl] = currentContract
+			currentContract = callEnterAddress
+		} else if sl > sLog.Depth {
+			logStack = logStack[:sl-1]
+			currentContract = contractStack[sLog.Depth]
+			resumePos := logStack[len(logStack)-1]
+			calledLog := logs[resumePos]
+
+			//no need to handle fail calling
+			if !(calledLog.ExtraData != nil && calledLog.ExtraData.CallFailed) {
+				//reentry the last log which "cause" the calling, some handling may needed
+				switch calledLog.Op {
+				case "CREATE", "CREATE2":
+					//addr, accDataBefore := getAccountDataFromProof(calledLog, posCALLBefore)
+					od.absorb(getAccountState(calledLog, posCREATEAfter))
+				}
+			}
+
+		} else {
+			logStack[sl-1] = i
+		}
+		//sanity check
+		if len(logStack) != sLog.Depth {
+			panic("tracking log stack failure")
+		}
+		callEnterAddress = currentContract
+
+		if skipDepth != 0 {
+			if skipDepth < sLog.Depth {
+				continue
+			} else {
+				skipDepth = 0
+			}
+		}
+
+		if exD := sLog.ExtraData; exD != nil && exD.CallFailed {
+			//mark current op and next ops with more depth skippable
+			skipDepth = sLog.Depth
+			continue
+		}
+
+		switch sLog.Op {
+		case "CREATE", "CREATE2":
+			state := getAccountState(sLog, posCREATE)
+			od.absorb(state)
+			//update contract to CREATE addr
+
+			callEnterAddress = state.Address
+		case "CALL", "CALLCODE":
+			state := getAccountState(sLog, posCALL)
+			od.absorb(state)
+			callEnterAddress = state.Address
+		case "STATICCALL":
+			//static call has no update on target address
+			callEnterAddress = getAccountState(sLog, posSTATICCALL).Address
+		case "SLOAD":
+			accountState := getAccountState(sLog, posSSTOREBefore)
+			od.absorb(accountState)
+		case "SSTORE":
+			log.Debug("build SSTORE", "pc", sLog.Pc, "key", sLog.Stack[len(sLog.Stack)-1])
+			accountState := copyAccountState(getAccountState(sLog, posSSTOREBefore))
+			// notice the log only provide the value BEFORE store and it is not suitable for our protocol,
+			// here we change it into value AFTER update
+			accountState.Storage = &types.StorageWrapper{
+				Key:   sLog.Stack[len(sLog.Stack)-1],
+				Value: sLog.Stack[len(sLog.Stack)-2],
+			}
+			od.absorb(accountState)
+
+		default:
+		}
+	}
+}
+
+func handleTx(od opOrderer, txResult *types.ExecutionResult) {
+
+	// handle failed tx
+	if txResult.Failed {
+		handled := false
+		for _, state := range txResult.AccountsAfter {
+			if state.Address != txResult.From.Address {
+				continue
+			}
+			od.absorb(state)
+			handled = true
+		}
+		if !handled {
+			panic(fmt.Errorf("no caller account in postTx status"))
+		}
+		return
+	}
+
+	var toAddr common.Address
+
+	if state := txResult.AccountCreated; state != nil {
+		od.absorb(state)
+		toAddr = state.Address
+	} else {
+		toAddr = txResult.To.Address
+	}
+
+	handleLogs(od, toAddr, txResult.StructLogs)
+
+	for _, state := range txResult.AccountsAfter {
+		od.absorb(state)
+	}
+
+}
+
+const defaultOrdererScheme = MPTWitnessRWTbl
+
+var usedOrdererScheme = defaultOrdererScheme
+
+// HandleBlockResult only for backward compatibility
+func HandleBlockResult(block *types.BlockResult) ([]*StorageTrace, error) {
+	return HandleBlockResultEx(block, usedOrdererScheme)
+}
+
+func HandleBlockResultEx(block *types.BlockResult, ordererScheme MPTWitnessType) ([]*StorageTrace, error) {
+
+	writer, err := NewZkTrieProofWriter(block.StorageTrace)
+	if err != nil {
+		return nil, err
+	}
+
+	var od opOrderer
+	switch ordererScheme {
+	case MPTWitnessNothing:
+		panic("should not come here when scheme is 0")
+	case MPTWitnessNatural:
+		od = &simpleOrderer{}
+	case MPTWitnessRWTbl:
+		od = newRWTblOrderer(writer.tracingAccounts)
+	default:
+		return nil, fmt.Errorf("unrecognized scheme %d", ordererScheme)
+	}
+
+	for _, tx := range block.ExecutionResults {
+		handleTx(od, tx)
+	}
+
+	// notice some coinbase addr (like all zero) is in fact not exist and should not be update
+	// TODO: not a good solution, just for patch ...
+	if coinbaseData := writer.tracingAccounts[block.BlockTrace.Coinbase.Address]; coinbaseData != nil {
+		od.absorb(block.BlockTrace.Coinbase)
+	}
+
+	opDisp := od.end_absorb()
+	var outTrace []*StorageTrace
+
+	for op := opDisp.next(); op != nil; op = opDisp.next() {
+		trace, err := writer.HandleNewState(op)
+		if err != nil {
+			return nil, err
+		}
+		outTrace = append(outTrace, trace)
+	}
+
+	finalHash := writer.tracingZktrie.Hash()
+	if !bytes.Equal(finalHash.Bytes(), block.StorageTrace.RootAfter.Bytes()) {
+		return outTrace, fmt.Errorf("unmatch hash: [%x] vs [%x]", finalHash.Bytes(), block.StorageTrace.RootAfter.Bytes())
+	}
+
+	return outTrace, nil
+
+}
+
+func FillBlockResultForMPTWitness(order MPTWitnessType, block *types.BlockResult) error {
+
+	if order == MPTWitnessNothing {
+		return nil
+	}
+
+	trace, err := HandleBlockResultEx(block, order)
+	if err != nil {
+		return err
+	}
+
+	msg, err := json.Marshal(trace)
+	if err != nil {
+		return err
+	}
+
+	rawmsg := json.RawMessage(msg)
+
+	block.MPTWitness = &rawmsg
+	return nil
+}

trie/zkproof/writer_test.go

@@ -0,0 +1,182 @@
+package zkproof
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"testing"
+
+	"github.com/scroll-tech/go-ethereum/common"
+	"github.com/scroll-tech/go-ethereum/core/types"
+)
+
+func init() {
+	orderScheme := os.Getenv("OP_ORDER")
+	var orderSchemeI int
+	if orderScheme != "" {
+		if n, err := fmt.Sscanf(orderScheme, "%d", &orderSchemeI); err == nil && n == 1 {
+			usedOrdererScheme = MPTWitnessType(orderSchemeI)
+		}
+	}
+}
+
+func loadStaff(t *testing.T, fname string) *types.BlockResult {
+	f, err := os.Open(fname)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+	bt, err := io.ReadAll(f)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	out := new(types.BlockResult)
+
+	err = json.Unmarshal(bt, out)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return out
+}
+
+func TestWriterCreation(t *testing.T) {
+	trace := loadStaff(t, "deploy_trace.json")
+	writer, err := NewZkTrieProofWriter(trace.StorageTrace)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(writer.tracingAccounts) != 4 {
+		t.Error("unexpected tracing account data", writer.tracingAccounts)
+	}
+
+	if v, existed := writer.tracingAccounts[common.HexToAddress("0x000000000000000000636F6e736F6c652e6c6f67")]; !existed || v != nil {
+		t.Error("wrong tracing status for uninited address", v, existed)
+	}
+
+	if v, existed := writer.tracingAccounts[common.HexToAddress("0xb36feAEaF76c2A33335b73bEF9aEf7a23d9af1e3")]; !existed || v != nil {
+		t.Error("wrong tracing status for uninited address", v, existed)
+	}
+
+	if v, existed := writer.tracingAccounts[common.HexToAddress("0x4cb1aB63aF5D8931Ce09673EbD8ae2ce16fD6571")]; !existed || v == nil {
+		t.Error("wrong tracing status for establied address", v, existed)
+	}
+
+	if len(writer.tracingStorageTries) != 1 {
+		t.Error("unexpected tracing storage data", writer.tracingStorageTries)
+	}
+
+	if v, existed := writer.tracingStorageTries[common.HexToAddress("0xb36feAEaF76c2A33335b73bEF9aEf7a23d9af1e3")]; !existed || v == nil {
+		t.Error("wrong tracing storage statu", existed, v)
+	}
+
+}
+
+func TestGreeterTx(t *testing.T) {
+	trace := loadStaff(t, "greeter_trace.json")
+	writer, err := NewZkTrieProofWriter(trace.StorageTrace)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	od := &simpleOrderer{}
+	theTx := trace.ExecutionResults[0]
+	handleTx(od, theTx)
+
+	t.Log(od)
+
+	for _, op := range od.savedOp {
+		_, err = writer.HandleNewState(op)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	traces, err := HandleBlockResult(trace)
+	t.Log("traces: ", len(traces))
+	outObj, _ := json.Marshal(traces)
+	t.Log(string(outObj))
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestTokenTx(t *testing.T) {
+	trace := loadStaff(t, "token_trace.json")
+	traces, err := HandleBlockResult(trace)
+	outObj, _ := json.Marshal(traces)
+	t.Log(string(outObj))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+}
+
+func TestCallTx(t *testing.T) {
+	trace := loadStaff(t, "call_trace.json")
+	traces, err := HandleBlockResult(trace)
+	outObj, _ := json.Marshal(traces)
+	t.Log(string(outObj))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	trace = loadStaff(t, "staticcall_trace.json")
+	traces, err = HandleBlockResult(trace)
+	outObj, _ = json.Marshal(traces)
+	t.Log(string(outObj))
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestCreateTx(t *testing.T) {
+	trace := loadStaff(t, "create_trace.json")
+	traces, err := HandleBlockResult(trace)
+	outObj, _ := json.Marshal(traces)
+	t.Log(string(outObj))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	trace = loadStaff(t, "deploy_trace.json")
+	traces, err = HandleBlockResult(trace)
+	outObj, _ = json.Marshal(traces)
+	t.Log(string(outObj))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+}
+
+func TestFailedCallTx(t *testing.T) {
+	trace := loadStaff(t, "fail_call_trace.json")
+	traces, err := HandleBlockResult(trace)
+	outObj, _ := json.Marshal(traces)
+	t.Log(string(outObj))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	trace = loadStaff(t, "fail_create_trace.json")
+	traces, err = HandleBlockResult(trace)
+	outObj, _ = json.Marshal(traces)
+	t.Log(string(outObj))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+}
+
+func TestMutipleTx(t *testing.T) {
+	trace := loadStaff(t, "multi_txs.json")
+	traces, err := HandleBlockResult(trace)
+	outObj, _ := json.Marshal(traces)
+	t.Log(string(outObj))
+	if err != nil {
+		t.Fatal(err)
+	}
+}