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

cmd/geth/main.go

@@ -201,6 +201,11 @@ var (
 		utils.MetricsInfluxDBBucketFlag,
 		utils.MetricsInfluxDBOrganizationFlag,
 	}
+
+	traceFlags = []cli.Flag{
+		utils.TraceCacheLimit,
+		utils.MPTWitness,
+	}
 )
 
 func init() {
@@ -247,6 +252,7 @@ func init() {
 	app.Flags = append(app.Flags, consoleFlags...)
 	app.Flags = append(app.Flags, debug.Flags...)
 	app.Flags = append(app.Flags, metricsFlags...)
+	app.Flags = append(app.Flags, traceFlags...)
 
 	app.Before = func(ctx *cli.Context) error {
 		return debug.Setup(ctx)

cmd/geth/usage.go

@@ -61,6 +61,7 @@ var AppHelpFlagGroups = []flags.FlagGroup{
 		Name: "NOTRACE CLIENT",
 		Flags: []cli.Flag{
 			utils.TraceCacheLimit,
+			utils.MPTWitness,
 		},
 	},
 	{

cmd/utils/flags.go

@@ -255,6 +255,12 @@ var (
 		Usage: "Handle the latest several blockResults",
 		Value: ethconfig.Defaults.TraceCacheLimit,
 	}
+	// mpt witness settings
+	MPTWitness = cli.IntFlag{
+		Name:  "trace.mptwitness",
+		Usage: "Output witness for mpt circuit with Specified order (default = no output, 1 = by executing order",
+		Value: ethconfig.Defaults.MPTWitness,
+	}
 	// Light server and client settings
 	LightServeFlag = cli.IntFlag{
 		Name:  "light.serve",
@@ -1105,6 +1111,8 @@ func setTrace(ctx *cli.Context, cfg *ethconfig.Config) {
 	if ctx.GlobalIsSet(TraceCacheLimit.Name) {
 		cfg.TraceCacheLimit = ctx.GlobalInt(TraceCacheLimit.Name)
 	}
+
+	cfg.MPTWitness = ctx.GlobalInt(MPTWitness.Name)
 }
 
 // setEtherbase retrieves the etherbase either from the directly specified

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:      0,
 }
 
 // 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(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

@@ -93,6 +93,7 @@ var Defaults = Config{
 	GPO:             FullNodeGPO,
 	RPCTxFeeCap:     1, // 1 ether
 	TraceCacheLimit: 32,
+	MPTWitness:      0,
 }
 
 func init() {
@@ -208,6 +209,7 @@ type Config struct {
 
 	// Trace option
 	TraceCacheLimit int
+	MPTWitness      int
 }
 
 // CreateConsensusEngine creates a consensus engine for the given chain configuration.

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()