zktrie part3: include zktrie witness in block trace; add demo for gen…

…erating 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)
-	if err == nil {
-		extraData.StateList = append(extraData.StateList, storage)
-	}
-	return err
+	storage := getWrappedAccountForStorage(l, scope.Contract.Address(), key)
+	extraData.StateList = append(extraData.StateList, storage)
+
+	return nil
 }
 
 // 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())
-	if err == nil {
-		extraData.StateList = append(extraData.StateList, state)
-		l.statesAffected[scope.Contract.Address()] = struct{}{}
-	}
-	return err
+	state := getWrappedAccountForAddr(l, scope.Contract.Address())
+	extraData.StateList = append(extraData.StateList, state)
+	l.statesAffected[scope.Contract.Address()] = struct{}{}
+
+	return nil
 }
 
 // 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)
-		if err == nil {
-			extraData.StateList = append(extraData.StateList, state)
-			l.statesAffected[scope.Contract.Address()] = struct{}{}
-		}
-		return err
+		state := getWrappedAccountForAddr(l, address)
+		extraData.StateList = append(extraData.StateList, state)
+		l.statesAffected[scope.Contract.Address()] = struct{}{}
+
+		return nil
 	}
 }
 
 // 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)
-	if err == nil {
-		extraData.StateList = append(extraData.StateList, state)
-		l.statesAffected[scope.Contract.Address()] = struct{}{}
-	}
-	return err
+	state := getWrappedAccountForAddr(l, address)
+
+	extraData.StateList = append(extraData.StateList, state)
+	l.statesAffected[scope.Contract.Address()] = struct{}{}
+
+	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,
-	RPCGasCap:       50000000,
-	RPCEVMTimeout:   5 * time.Second,
-	GPO:             FullNodeGPO,
-	RPCTxFeeCap:     1, // 1 ether
-	TraceCacheLimit: 32,
+	TxPool:        core.DefaultTxPoolConfig,
+	RPCGasCap:     50000000,
+	RPCEVMTimeout: 5 * time.Second,
+	GPO:           FullNodeGPO,
+	RPCTxFeeCap:   1, // 1 ether
 }
 
 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()