trie: make value node resolvable via a callback

Author: jwasinger

cmd/devp2p/README.md

@@ -121,7 +121,7 @@ with our test chain. The chain files are located in `./cmd/devp2p/internal/ethte
         --nat=none                     \
         --networkid 3503995874084926   \
         --verbosity 5                  \
-        --authrpc.jwtsecret 0x7365637265747365637265747365637265747365637265747365637265747365
+        --authrpc.jwtsecret jwt.secret
 
 Note that the tests also require access to the engine API.
 The test suite can now be executed using the devp2p tool.
@@ -130,7 +130,7 @@ The test suite can now be executed using the devp2p tool.
         --chain internal/ethtest/testdata   \
         --node enode://....                 \
         --engineapi http://127.0.0.1:8551   \
-        --jwtsecret 0x7365637265747365637265747365637265747365637265747365637265747365
+        --jwtsecret $(cat jwt.secret)
 
 Repeat the above process (re-initialising the node) in order to run the Eth Protocol test suite again.
 

cmd/devp2p/rlpxcmd.go

@@ -143,9 +143,6 @@ type testParams struct {
 
 func cliTestParams(ctx *cli.Context) *testParams {
 	nodeStr := ctx.String(testNodeFlag.Name)
-	if nodeStr == "" {
-		exit(fmt.Errorf("missing -%s", testNodeFlag.Name))
-	}
 	node, err := parseNode(nodeStr)
 	if err != nil {
 		exit(err)
@@ -156,14 +153,5 @@ func cliTestParams(ctx *cli.Context) *testParams {
 		jwt:       ctx.String(testNodeJWTFlag.Name),
 		chainDir:  ctx.String(testChainDirFlag.Name),
 	}
-	if p.engineAPI == "" {
-		exit(fmt.Errorf("missing -%s", testNodeEngineFlag.Name))
-	}
-	if p.jwt == "" {
-		exit(fmt.Errorf("missing -%s", testNodeJWTFlag.Name))
-	}
-	if p.chainDir == "" {
-		exit(fmt.Errorf("missing -%s", testChainDirFlag.Name))
-	}
 	return &p
 }

cmd/devp2p/runtest.go

@@ -39,26 +39,29 @@ var (
 	}
 
 	// for eth/snap tests
-	testChainDirFlag = &cli.StringFlag{
+	testChainDirFlag = &cli.PathFlag{
 		Name:     "chain",
 		Usage:    "Test chain directory (required)",
 		Category: flags.TestingCategory,
+		Required: true,
 	}
 	testNodeFlag = &cli.StringFlag{
 		Name:     "node",
 		Usage:    "Peer-to-Peer endpoint (ENR) of the test node (required)",
 		Category: flags.TestingCategory,
+		Required: true,
 	}
 	testNodeJWTFlag = &cli.StringFlag{
 		Name:     "jwtsecret",
 		Usage:    "JWT secret for the engine API of the test node (required)",
 		Category: flags.TestingCategory,
-		Value:    "0x7365637265747365637265747365637265747365637265747365637265747365",
+		Required: true,
 	}
 	testNodeEngineFlag = &cli.StringFlag{
 		Name:     "engineapi",
 		Usage:    "Engine API endpoint of the test node (required)",
 		Category: flags.TestingCategory,
+		Required: true,
 	}
 
 	// These two are specific to the discovery tests.

cmd/utils/flags.go

@@ -20,7 +20,6 @@ package utils
 import (
 	"context"
 	"crypto/ecdsa"
-	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -1341,15 +1340,10 @@ func setEtherbase(ctx *cli.Context, cfg *ethconfig.Config) {
 		return
 	}
 	addr := ctx.String(MinerPendingFeeRecipientFlag.Name)
-	if strings.HasPrefix(addr, "0x") || strings.HasPrefix(addr, "0X") {
-		addr = addr[2:]
-	}
-	b, err := hex.DecodeString(addr)
-	if err != nil || len(b) != common.AddressLength {
+	if !common.IsHexAddress(addr) {
 		Fatalf("-%s: invalid pending block producer address %q", MinerPendingFeeRecipientFlag.Name, addr)
-		return
 	}
-	cfg.Miner.PendingFeeRecipient = common.BytesToAddress(b)
+	cfg.Miner.PendingFeeRecipient = common.HexToAddress(addr)
 }
 
 func SetP2PConfig(ctx *cli.Context, cfg *p2p.Config) {

core/state/database.go

@@ -99,6 +99,8 @@ type Trie interface {
 	// in the trie with provided address.
 	UpdateAccount(address common.Address, account *types.StateAccount, codeLen int) error
 
+	UpdateAccountAsync(address common.Address, accountResolver func() *types.StateAccount) error
+
 	// UpdateStorage associates key with value in the trie. If value has length zero,
 	// any existing value is deleted from the trie. The value bytes must not be modified
 	// by the caller while they are stored in the trie. If a node was not found in the

core/state/state_sizer_test.go

@@ -94,6 +94,14 @@ func TestSizeTracker(t *testing.T) {
 	}
 	baselineRoot := currentRoot
 
+	// Close and reopen the trie database so all async flushes triggered by the
+	// baseline commits are written before we measure the baseline snapshot.
+	if err := tdb.Close(); err != nil {
+		t.Fatalf("Failed to close triedb before baseline measurement: %v", err)
+	}
+	tdb = triedb.NewDatabase(db, &triedb.Config{PathDB: pathdb.Defaults})
+	sdb = NewDatabase(tdb, nil)
+
 	// Wait for snapshot completion
 	for !tdb.SnapshotCompleted() {
 		time.Sleep(100 * time.Millisecond)
@@ -215,13 +223,12 @@ func TestSizeTracker(t *testing.T) {
 	if actualStats.ContractCodeBytes != expectedStats.ContractCodeBytes {
 		t.Errorf("Contract code bytes mismatch: expected %d, got %d", expectedStats.ContractCodeBytes, actualStats.ContractCodeBytes)
 	}
-	// TODO: failed on github actions, need to investigate
-	// if actualStats.AccountTrienodes != expectedStats.AccountTrienodes {
-	// 	t.Errorf("Account trie nodes mismatch: expected %d, got %d", expectedStats.AccountTrienodes, actualStats.AccountTrienodes)
-	// }
-	// if actualStats.AccountTrienodeBytes != expectedStats.AccountTrienodeBytes {
-	// 	t.Errorf("Account trie node bytes mismatch: expected %d, got %d", expectedStats.AccountTrienodeBytes, actualStats.AccountTrienodeBytes)
-	// }
+	if actualStats.AccountTrienodes != expectedStats.AccountTrienodes {
+		t.Errorf("Account trie nodes mismatch: expected %d, got %d", expectedStats.AccountTrienodes, actualStats.AccountTrienodes)
+	}
+	if actualStats.AccountTrienodeBytes != expectedStats.AccountTrienodeBytes {
+		t.Errorf("Account trie node bytes mismatch: expected %d, got %d", expectedStats.AccountTrienodeBytes, actualStats.AccountTrienodeBytes)
+	}
 	if actualStats.StorageTrienodes != expectedStats.StorageTrienodes {
 		t.Errorf("Storage trie nodes mismatch: expected %d, got %d", expectedStats.StorageTrienodes, actualStats.StorageTrienodes)
 	}

core/state/statedb.go

@@ -577,6 +577,13 @@ func (s *StateDB) updateStateObject(obj *stateObject) {
 	}
 }
 
+// updateStateObject writes the given object to the trie.
+func (s *StateDB) updateStateObjectAsync(addr common.Address, resolver func() *types.StateAccount) {
+	if err := s.trie.UpdateAccountAsync(addr, resolver); err != nil {
+		s.setError(fmt.Errorf("updateStateObject (%x) error: %v", addr, err))
+	}
+}
+
 // deleteStateObject removes the given object from the state trie.
 func (s *StateDB) deleteStateObject(addr common.Address) {
 	if err := s.trie.DeleteAccount(addr); err != nil {
@@ -829,11 +836,14 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 		// later time.
 		workers.SetLimit(1)
 	}
+
+	stateObjectsResolve := make(map[common.Address]func() *types.StateAccount)
 	for addr, op := range s.mutations {
 		if op.applied || op.isDelete() {
 			continue
 		}
 		obj := s.stateObjects[addr] // closure for the task runner below
+		complete := make(chan *types.StateAccount)
 		workers.Go(func() error {
 			if s.db.TrieDB().IsVerkle() {
 				obj.updateTrie()
@@ -846,8 +856,13 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 					s.witness.AddState(obj.trie.Witness())
 				}
 			}
+			complete <- &obj.data
 			return nil
 		})
+
+		stateObjectsResolve[addr] = func() *types.StateAccount {
+			return <-complete
+		}
 	}
 	// If witness building is enabled, gather all the read-only accesses.
 	// Skip witness collection in Verkle mode, they will be gathered
@@ -898,7 +913,6 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 			}
 		}
 	}
-	workers.Wait()
 	s.StorageUpdates += time.Since(start)
 
 	// Now we're about to start to write changes to the trie. The trie is so far
@@ -939,7 +953,11 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 		if op.isDelete() {
 			deletedAddrs = append(deletedAddrs, addr)
 		} else {
-			s.updateStateObject(s.stateObjects[addr])
+			if s.db.TrieDB().IsVerkle() {
+				s.updateStateObject(s.stateObjects[addr])
+			} else {
+				s.updateStateObjectAsync(addr, stateObjectsResolve[addr])
+			}
 			s.AccountUpdated += 1
 		}
 		usedAddrs = append(usedAddrs, addr) // Copy needed for closure
@@ -966,6 +984,7 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 			s.witnessStats.Add(witness, common.Hash{})
 		}
 	}
+
 	return hash
 }
 

crypto/keccak_ziren.go

@@ -21,14 +21,72 @@ package crypto
 import (
 	"github.com/ProjectZKM/Ziren/crates/go-runtime/zkvm_runtime"
 	"github.com/ethereum/go-ethereum/common"
-	"golang.org/x/crypto/sha3"
 )
 
+// zirenKeccakState implements the KeccakState interface using the Ziren zkvm_runtime.
+// It accumulates data written to it and uses the zkvm's Keccak256 system call for hashing.
+type zirenKeccakState struct {
+	buf    []byte // accumulated data
+	result []byte // cached result
+	dirty  bool   // whether new data has been written since last hash
+}
+
+func newZirenKeccakState() KeccakState {
+	return &zirenKeccakState{
+		buf: make([]byte, 0, 512), // pre-allocate reasonable capacity
+	}
+}
+
+func (s *zirenKeccakState) Write(p []byte) (n int, err error) {
+	s.buf = append(s.buf, p...)
+	s.dirty = true
+	return len(p), nil
+}
+
+func (s *zirenKeccakState) Sum(b []byte) []byte {
+	s.computeHashIfNeeded()
+	return append(b, s.result...)
+}
+
+func (s *zirenKeccakState) Reset() {
+	s.buf = s.buf[:0]
+	s.result = nil
+	s.dirty = false
+}
+
+func (s *zirenKeccakState) Size() int {
+	return 32
+}
+
+func (s *zirenKeccakState) BlockSize() int {
+	return 136 // Keccak256 rate
+}
+
+func (s *zirenKeccakState) Read(p []byte) (n int, err error) {
+	s.computeHashIfNeeded()
+
+	if len(p) == 0 {
+		return 0, nil
+	}
+
+	// After computeHashIfNeeded(), s.result is always a 32-byte slice
+	n = copy(p, s.result)
+	return n, nil
+}
+
+func (s *zirenKeccakState) computeHashIfNeeded() {
+	if s.dirty || s.result == nil {
+		// Use the zkvm_runtime Keccak256 which uses SyscallKeccakSponge
+		hashArray := zkvm_runtime.Keccak256(s.buf)
+		s.result = hashArray[:]
+		s.dirty = false
+	}
+}
+
 // NewKeccakState creates a new KeccakState
-// For now, we fallback to the original implementation for the stateful interface.
-// TODO: Implement a stateful wrapper around zkvm_runtime.Keccak256 if needed.
+// This uses a Ziren-optimized implementation that leverages the zkvm_runtime.Keccak256 system call.
 func NewKeccakState() KeccakState {
-	return sha3.NewLegacyKeccak256().(KeccakState)
+	return newZirenKeccakState()
 }
 
 // Keccak256 calculates and returns the Keccak256 hash using the Ziren zkvm_runtime implementation.

p2p/server.go

@@ -813,7 +813,9 @@ func (srv *Server) listenLoop() {
 				time.Sleep(time.Millisecond * 200)
 				continue
 			} else if err != nil {
-				srv.log.Debug("Read error", "err", err)
+				if !errors.Is(err, net.ErrClosed) {
+					srv.log.Debug("Read error", "err", err)
+				}
 				slots <- struct{}{}
 				return
 			}

rpc/client_test.go

@@ -973,7 +973,7 @@ func (l *flakeyListener) Accept() (net.Conn, error) {
 
 	c, err := l.Listener.Accept()
 	if err == nil {
-		timeout := time.Duration(rand.Int63n(int64(l.maxKillTimeout)))
+		timeout := max(time.Millisecond*10, time.Duration(rand.Int63n(int64(l.maxKillTimeout))))
 		time.AfterFunc(timeout, func() {
 			log.Debug(fmt.Sprintf("killing conn %v after %v", c.LocalAddr(), timeout))
 			c.Close()