overlay benchmark: enable proof generation

core/state/access_witness.go

@@ -17,9 +17,15 @@
 package state
 
 import (
+	"fmt"
+	"time"
+
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/params"
+	"github.com/ethereum/go-ethereum/trie"
 	"github.com/ethereum/go-ethereum/trie/utils"
+	"github.com/gballet/go-verkle"
 	"github.com/holiman/uint256"
 )
 
@@ -43,6 +49,25 @@ type AccessWitness struct {
 	chunks   map[chunkAccessKey]mode
 
 	pointCache *utils.PointCache
+
+	// Fields below are only used when the witness tree loader
+	// is enabled by using the NewAccessWitnessWithTreeLoader constructor.
+
+	// treeLoaderChunks is a channel used to schedule chunks to be loaded.
+	treeLoaderChunks chan chunkAccessKey
+	// tree is the witness tree used to load chunks.
+	tree *trie.VerkleTrie
+	// treeLoaderQuit is used to signal the background tree loader to quit.
+	treeLoaderQuit chan struct{}
+	// treeLoaderClosed is used to signal that the background tree loader has quit.
+	treeLoaderClosed chan struct{}
+	// treeLoaderErr is used to signal that the background tree loader has
+	// encountered an error.
+	treeLoaderErr error
+	// witnessKeys is the list of keys that have been accessed.
+	witnessKeys [][]byte
+	// witnessKeyValues is the list of values that have been accessed.
+	witnessKeyValues map[string][]byte
 }
 
 func NewAccessWitness(pointCache *utils.PointCache) *AccessWitness {
@@ -53,6 +78,19 @@ func NewAccessWitness(pointCache *utils.PointCache) *AccessWitness {
 	}
 }
 
+// NewAccessWitnessWithTreeLoader creates a new AccessWitness that will
+// construct a witness tree in the background when merging child witnesses.
+// The caller is expected to call ProveAndSerialize() to get the proof for
+// the witness tree, which will also clean up background tasks.
+func NewAccessWitnessWithTreeLoader(pointCache *utils.PointCache, trie *trie.VerkleTrie) *AccessWitness {
+	aw := NewAccessWitness(pointCache)
+	aw.treeLoaderChunks = make(chan chunkAccessKey, 64)
+	aw.tree = trie
+	go aw.backgroundTreeLoader()
+
+	return aw
+}
+
 // Merge is used to merge the witness that got generated during the execution
 // of a tx, with the accumulation of witnesses that were generated during the
 // execution of all the txs preceding this one in a given block.
@@ -61,6 +99,13 @@ func (aw *AccessWitness) Merge(other *AccessWitness) {
 		aw.branches[k] |= other.branches[k]
 	}
 	for k, chunk := range other.chunks {
+		// If the to-be-merged chunk isn't already present,
+		// schedule to be warmed up in the loding witness tree.
+		if aw.treeLoaderChunks != nil {
+			if _, ok := aw.chunks[k]; !ok {
+				aw.treeLoaderChunks <- k
+			}
+		}
 		aw.chunks[k] |= chunk
 	}
 }
@@ -229,6 +274,50 @@ func (aw *AccessWitness) touchAddress(addr []byte, treeIndex uint256.Int, subInd
 	return branchRead, chunkRead, branchWrite, chunkWrite, chunkFill
 }
 
+func (aw *AccessWitness) GenerateProofAndSerialize() (*verkle.VerkleProof, verkle.StateDiff, error) {
+	start := time.Now()
+	// Signal that we are done sending chunks to load the tree.
+	aw.treeLoaderQuit <- struct{}{}
+	// Wait for the background loader to finish loading pending keys.
+	<-aw.treeLoaderClosed
+	log.Debug("Waiting for the tree loader to finish took %s", time.Since(start))
+
+	// If the background loader encountered an error, return it.
+	if aw.treeLoaderErr != nil {
+		return nil, nil, fmt.Errorf("loading the witness tree failed: %s", aw.treeLoaderErr)
+	}
+
+	// The tree is fully loaded and ready to support proof generation.
+	proof, stateDiff, err := aw.tree.ProveAndSerialize(aw.witnessKeys, aw.witnessKeyValues)
+	if err != nil {
+		return nil, nil, fmt.Errorf("generating the witness proof failed: %s", err)
+	}
+
+	return proof, stateDiff, nil
+}
+
+func (aw *AccessWitness) backgroundTreeLoader() {
+	defer close(aw.treeLoaderClosed)
+	for {
+		select {
+		// If we receive a new chunk, load it into the tree.
+		case chunk := <-aw.treeLoaderChunks:
+			basePoint := aw.pointCache.GetTreeKeyHeader(chunk.addr[:])
+			key := utils.GetTreeKeyWithEvaluatedAddess(basePoint, &chunk.treeIndex, chunk.leafKey)
+			value, err := aw.tree.GetAndLoadForProof(key)
+			if err != nil {
+				aw.treeLoaderErr = err
+				return
+			}
+			aw.witnessKeys = append(aw.witnessKeys, key)
+			aw.witnessKeyValues[string(key)] = value
+		// If we receive a quit signal, stop the loader.
+		case <-aw.treeLoaderQuit:
+			return
+		}
+	}
+}
+
 type branchAccessKey struct {
 	addr      common.Address
 	treeIndex uint256.Int

core/state/statedb.go

@@ -212,7 +212,8 @@ func (s *StateDB) NewAccessWitness() *AccessWitness {
 
 func (s *StateDB) Witness() *AccessWitness {
 	if s.witness == nil {
-		s.witness = s.NewAccessWitness()
+		verkleTree := s.GetTrie().(*trie.TransitionTrie).Overlay()
+		s.witness = NewAccessWitnessWithTreeLoader(s.db.(*cachingDB).addrToPoint, verkleTree)
 	}
 	return s.witness
 }

core/state_processor.go

@@ -115,6 +115,12 @@ func (p *StateProcessor) Process(block *types.Block, statedb *state.StateDB, cfg
 	// verkle transition: if the conversion process is in progress, move
 	// N values from the MPT into the verkle tree.
 	if migrdb.InTransition() {
+		proof, stateDiff, err := statedb.Witness().GenerateProofAndSerialize()
+		if err != nil {
+			return nil, nil, 0, fmt.Errorf("could not generate proof: %w", err)
+		}
+		_, _ = proof, stateDiff
+
 		var (
 			now             = time.Now()
 			tt              = statedb.GetTrie().(*trie.TransitionTrie)

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/cloudflare/cloudflare-go v0.14.0
 	github.com/cockroachdb/pebble v0.0.0-20230209160836-829675f94811
 	github.com/consensys/gnark-crypto v0.10.0
+	github.com/crate-crypto/go-ipa v0.0.0-20230710183535-d5eb1c4661bd
 	github.com/crate-crypto/go-kzg-4844 v0.3.0
 	github.com/davecgh/go-spew v1.1.1
 	github.com/deckarep/golang-set/v2 v2.1.0
@@ -25,7 +26,7 @@ require (
 	github.com/fjl/memsize v0.0.0-20190710130421-bcb5799ab5e5
 	github.com/fsnotify/fsnotify v1.6.0
 	github.com/gballet/go-libpcsclite v0.0.0-20190607065134-2772fd86a8ff
-	github.com/gballet/go-verkle v0.0.0-20230725193842-b2d852dc666b
+	github.com/gballet/go-verkle v0.0.0-20230809181720-8d8980709295
 	github.com/go-stack/stack v1.8.1
 	github.com/gofrs/flock v0.8.1
 	github.com/golang-jwt/jwt/v4 v4.3.0
@@ -92,7 +93,6 @@ require (
 	github.com/cockroachdb/redact v1.1.3 // indirect
 	github.com/consensys/bavard v0.1.13 // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect
-	github.com/crate-crypto/go-ipa v0.0.0-20230710183535-d5eb1c4661bd // indirect
 	github.com/decred/dcrd/dcrec/secp256k1/v4 v4.0.1 // indirect
 	github.com/deepmap/oapi-codegen v1.8.2 // indirect
 	github.com/dlclark/regexp2 v1.7.0 // indirect

go.sum

@@ -84,8 +84,6 @@ github.com/cpuguy83/go-md2man v1.0.10/go.mod h1:SmD6nW6nTyfqj6ABTjUi3V3JVMnlJmwc
 github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU=
 github.com/cpuguy83/go-md2man/v2 v2.0.2 h1:p1EgwI/C7NhT0JmVkwCD2ZBK8j4aeHQX2pMHHBfMQ6w=
 github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
-github.com/crate-crypto/go-ipa v0.0.0-20230601170251-1830d0757c80 h1:DuBDHVjgGMPki7bAyh91+3cF1Vh34sAEdH8JQgbc2R0=
-github.com/crate-crypto/go-ipa v0.0.0-20230601170251-1830d0757c80/go.mod h1:gzbVz57IDJgQ9rLQwfSk696JGWof8ftznEL9GoAv3NI=
 github.com/crate-crypto/go-ipa v0.0.0-20230710183535-d5eb1c4661bd h1:jgf65Q4+jHFuLlhVApaVfTUwcU7dAdXK+GESow2UlaI=
 github.com/crate-crypto/go-ipa v0.0.0-20230710183535-d5eb1c4661bd/go.mod h1:gzbVz57IDJgQ9rLQwfSk696JGWof8ftznEL9GoAv3NI=
 github.com/crate-crypto/go-kzg-4844 v0.3.0 h1:UBlWE0CgyFqqzTI+IFyCzA7A3Zw4iip6uzRv5NIXG0A=
@@ -146,10 +144,8 @@ github.com/garslo/gogen v0.0.0-20170306192744-1d203ffc1f61/go.mod h1:Q0X6pkwTILD
 github.com/gavv/httpexpect v2.0.0+incompatible/go.mod h1:x+9tiU1YnrOvnB725RkpoLv1M62hOWzwo5OXotisrKc=
 github.com/gballet/go-libpcsclite v0.0.0-20190607065134-2772fd86a8ff h1:tY80oXqGNY4FhTFhk+o9oFHGINQ/+vhlm8HFzi6znCI=
 github.com/gballet/go-libpcsclite v0.0.0-20190607065134-2772fd86a8ff/go.mod h1:x7DCsMOv1taUwEWCzT4cmDeAkigA5/QCwUodaVOe8Ww=
-github.com/gballet/go-verkle v0.0.0-20230607174250-df487255f46b h1:vMT47RYsrftsHSTQhqXwC3BYflo38OLC3Y4LtXtLyU0=
-github.com/gballet/go-verkle v0.0.0-20230607174250-df487255f46b/go.mod h1:CDncRYVRSDqwakm282WEkjfaAj1hxU/v5RXxk5nXOiI=
-github.com/gballet/go-verkle v0.0.0-20230725193842-b2d852dc666b h1:2lDzSxjCii8FxrbuxtlFtFiw6c4nTPl9mhaZ6lgpwws=
-github.com/gballet/go-verkle v0.0.0-20230725193842-b2d852dc666b/go.mod h1:+k9fzNguudDonU5q4/TUaTdmiHw3h3oGOIVmqyhaA3E=
+github.com/gballet/go-verkle v0.0.0-20230809181720-8d8980709295 h1:4vXMHXZuVdrpVnUXYC224zIC3XHilsZz3H8UvwETpoc=
+github.com/gballet/go-verkle v0.0.0-20230809181720-8d8980709295/go.mod h1:+k9fzNguudDonU5q4/TUaTdmiHw3h3oGOIVmqyhaA3E=
 github.com/getkin/kin-openapi v0.53.0/go.mod h1:7Yn5whZr5kJi6t+kShccXS8ae1APpYTW6yheSwk8Yi4=
 github.com/getkin/kin-openapi v0.61.0/go.mod h1:7Yn5whZr5kJi6t+kShccXS8ae1APpYTW6yheSwk8Yi4=
 github.com/getsentry/sentry-go v0.12.0/go.mod h1:NSap0JBYWzHND8oMbyi0+XZhUalc1TBdRL1M71JZW2c=
@@ -576,8 +572,6 @@ golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.9.0 h1:KS/R3tvhPqvJvwcKfnBHJwwthS11LRhmM5D59eEXa0s=
-golang.org/x/sys v0.9.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.10.0 h1:SqMFp9UcQJZa+pmYuAKjd9xq1f0j5rLcDIk0mj4qAsA=
 golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=

trie/verkle.go

@@ -19,13 +19,16 @@ package trie
 import (
 	"bytes"
 	"encoding/binary"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"math/big"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/ethdb"
+	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/trie/trienode"
 	"github.com/ethereum/go-ethereum/trie/utils"
 	"github.com/gballet/go-verkle"
@@ -316,21 +319,32 @@ func (trie *VerkleTrie) Copy() *VerkleTrie {
 	}
 }
 
+func (trie *VerkleTrie) GetAndLoadForProof(hashedKey []byte) ([]byte, error) {
+	return trie.root.(*verkle.InternalNode).GetAndLoadForProof(hashedKey, trie.flatdbNodeResolver)
+}
+
 func (trie *VerkleTrie) IsVerkle() bool {
 	return true
 }
 
 func (trie *VerkleTrie) ProveAndSerialize(keys [][]byte, kv map[string][]byte) (*verkle.VerkleProof, verkle.StateDiff, error) {
+	// TODO: remove verbose logging.
+	startProofGen := time.Now()
 	proof, _, _, _, err := verkle.MakeVerkleMultiProof(trie.root, keys)
 	if err != nil {
 		return nil, nil, err
 	}
 
+	startProofSerialization := time.Now()
 	p, kvps, err := verkle.SerializeProof(proof)
 	if err != nil {
 		return nil, nil, err
 	}
-
+	proofJson, err := json.Marshal(proof)
+	if err != nil {
+		return nil, nil, err
+	}
+	log.Debug("Generating the proof", "number of keys", len(keys), "building proof", startProofSerialization.Sub(startProofGen), "serializing proof", time.Since(startProofSerialization), "json encoded size", common.StorageSize(len(proofJson)))
 	return p, kvps, nil
 }