fix(zktrie): fix deletion proofs and collect them in commiting phase (#263)

* new deletion proof

* complete tracer and test

* extend proof for parallel tracing

* integrating into blocktrace

* lint

* deduplication of deletion proofs

* fix an issue on marking deletion

* fixs since last review

* Update version.go

---------

Co-authored-by: HAOYUatHZ <37070449+HAOYUatHZ@users.noreply.github.com>
Co-authored-by: Ubuntu <ubuntu@ip-10-2-9-63.us-west-2.compute.internal>

core/state/state_prove.go

@@ -0,0 +1,88 @@
+package state
+
+import (
+	"fmt"
+
+	zkt "github.com/scroll-tech/zktrie/types"
+
+	zktrie "github.com/scroll-tech/go-ethereum/trie"
+
+	"github.com/scroll-tech/go-ethereum/common"
+	"github.com/scroll-tech/go-ethereum/crypto"
+	"github.com/scroll-tech/go-ethereum/ethdb"
+)
+
+type TrieProve interface {
+	Prove(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) error
+}
+
+type ZktrieProofTracer struct {
+	*zktrie.ProofTracer
+}
+
+// MarkDeletion overwrite the underlayer method with secure key
+func (t ZktrieProofTracer) MarkDeletion(key common.Hash) {
+	key_s, _ := zkt.ToSecureKeyBytes(key.Bytes())
+	t.ProofTracer.MarkDeletion(key_s.Bytes())
+}
+
+// Merge overwrite underlayer method with proper argument
+func (t ZktrieProofTracer) Merge(another ZktrieProofTracer) {
+	t.ProofTracer.Merge(another.ProofTracer)
+}
+
+func (t ZktrieProofTracer) Available() bool {
+	return t.ProofTracer != nil
+}
+
+// NewProofTracer is not in Db interface and used explictily for reading proof in storage trie (not updated by the dirty value)
+func (s *StateDB) NewProofTracer(trieS Trie) ZktrieProofTracer {
+	if s.IsZktrie() {
+		zkTrie := trieS.(*zktrie.ZkTrie)
+		if zkTrie == nil {
+			panic("unexpected trie type for zktrie")
+		}
+		return ZktrieProofTracer{zkTrie.NewProofTracer()}
+	}
+	return ZktrieProofTracer{}
+}
+
+// GetStorageTrieForProof is not in Db interface and used explictily for reading proof in storage trie (not updated by the dirty value)
+func (s *StateDB) GetStorageTrieForProof(addr common.Address) (Trie, error) {
+
+	// try the trie in stateObject first, else we would create one
+	stateObject := s.getStateObject(addr)
+	if stateObject == nil {
+		// still return a empty trie
+		addrHash := crypto.Keccak256Hash(addr[:])
+		dummy_trie, _ := s.db.OpenStorageTrie(addrHash, common.Hash{})
+		return dummy_trie, nil
+	}
+
+	trie := stateObject.trie
+	var err error
+	if trie == nil {
+		// use a new, temporary trie
+		trie, err = s.db.OpenStorageTrie(stateObject.addrHash, stateObject.data.Root)
+		if err != nil {
+			return nil, fmt.Errorf("can't create storage trie on root %s: %v ", stateObject.data.Root, err)
+		}
+	}
+
+	return trie, nil
+}
+
+// GetSecureTrieProof handle any interface with Prove (should be a Trie in most case) and
+// deliver the proof in bytes
+func (s *StateDB) GetSecureTrieProof(trieProve TrieProve, key common.Hash) ([][]byte, error) {
+
+	var proof proofList
+	var err error
+	if s.IsZktrie() {
+		key_s, _ := zkt.ToSecureKeyBytes(key.Bytes())
+		err = trieProve.Prove(key_s.Bytes(), 0, &proof)
+	} else {
+		err = trieProve.Prove(crypto.Keccak256(key.Bytes()), 0, &proof)
+	}
+	return proof, err
+}

core/state/statedb.go

@@ -350,56 +350,13 @@ func (s *StateDB) GetRootHash() common.Hash {
 	return s.trie.Hash()
 }
 
-// StorageTrieProof is not in Db interface and used explictily for reading proof in storage trie (not the dirty value)
-// For zktrie it also provide required data for predict the deletion, else it just fallback to GetStorageProof
-func (s *StateDB) GetStorageTrieProof(a common.Address, key common.Hash) ([][]byte, []byte, error) {
-
-	// try the trie in stateObject first, else we would create one
-	stateObject := s.getStateObject(a)
-	if stateObject == nil {
-		return nil, nil, errors.New("storage trie for requested address does not exist")
-	}
-
-	trieS := stateObject.trie
-	var err error
-	if trieS == nil {
-		// use a new, temporary trie
-		trieS, err = s.db.OpenStorageTrie(stateObject.addrHash, stateObject.data.Root)
-		if err != nil {
-			return nil, nil, fmt.Errorf("can't create storage trie on root %s: %v ", stateObject.data.Root, err)
-		}
-	}
-
-	var proof proofList
-	var sibling []byte
-	if s.IsZktrie() {
-		zkTrie := trieS.(*trie.ZkTrie)
-		if zkTrie == nil {
-			panic("unexpected trie type for zktrie")
-		}
-		key_s, _ := zkt.ToSecureKeyBytes(key.Bytes())
-		sibling, err = zkTrie.ProveWithDeletion(key_s.Bytes(), 0, &proof)
-	} else {
-		err = trieS.Prove(crypto.Keccak256(key.Bytes()), 0, &proof)
-	}
-	return proof, sibling, err
-}
-
 // GetStorageProof returns the Merkle proof for given storage slot.
 func (s *StateDB) GetStorageProof(a common.Address, key common.Hash) ([][]byte, error) {
-	var proof proofList
 	trie := s.StorageTrie(a)
 	if trie == nil {
-		return proof, errors.New("storage trie for requested address does not exist")
+		return nil, errors.New("storage trie for requested address does not exist")
 	}
-	var err error
-	if s.IsZktrie() {
-		key_s, _ := zkt.ToSecureKeyBytes(key.Bytes())
-		err = trie.Prove(key_s.Bytes(), 0, &proof)
-	} else {
-		err = trie.Prove(crypto.Keccak256(key.Bytes()), 0, &proof)
-	}
-	return proof, err
+	return s.GetSecureTrieProof(trie, key)
 }
 
 // GetCommittedState retrieves a value from the given account's committed storage trie.

eth/tracers/api_blocktrace.go

@@ -1,6 +1,7 @@
 package tracers
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -42,6 +43,8 @@ type traceEnv struct {
 	// this lock is used to protect StorageTrace's read and write mutual exclusion.
 	sMu sync.Mutex
 	*types.StorageTrace
+	// zktrie tracer is used for zktrie storage to build additional deletion proof
+	zkTrieTracer     map[string]state.ZktrieProofTracer
 	executionResults []*types.ExecutionResult
 }
 
@@ -119,6 +122,7 @@ func (api *API) createTraceEnv(ctx context.Context, config *TraceConfig, block *
 			Proofs:        make(map[string][]hexutil.Bytes),
 			StorageProofs: make(map[string]map[string][]hexutil.Bytes),
 		},
+		zkTrieTracer:     make(map[string]state.ZktrieProofTracer),
 		executionResults: make([]*types.ExecutionResult, block.Transactions().Len()),
 	}
 
@@ -189,6 +193,18 @@ func (api *API) getBlockTrace(block *types.Block, env *traceEnv) (*types.BlockTr
 	close(jobs)
 	pend.Wait()
 
+	// after all tx has been traced, collect "deletion proof" for zktrie
+	for _, tracer := range env.zkTrieTracer {
+		delProofs, err := tracer.GetDeletionProofs()
+		if err != nil {
+			log.Error("deletion proof failure", "error", err)
+		} else {
+			for _, proof := range delProofs {
+				env.DeletionProofs = append(env.DeletionProofs, proof)
+			}
+		}
+	}
+
 	// If execution failed in between, abort
 	select {
 	case err := <-errCh:
@@ -299,22 +315,47 @@ func (api *API) getTxResult(env *traceEnv, state *state.StateDB, index int, bloc
 
 	proofStorages := tracer.UpdatedStorages()
 	for addr, keys := range proofStorages {
-		for key := range keys {
+		env.sMu.Lock()
+		trie, err := state.GetStorageTrieForProof(addr)
+		if err != nil {
+			// but we still continue to next address
+			log.Error("Storage trie not available", "error", err, "address", addr)
+			env.sMu.Unlock()
+			continue
+		}
+		zktrieTracer := state.NewProofTracer(trie)
+		env.sMu.Unlock()
+
+		for key, values := range keys {
 			addrStr := addr.String()
 			keyStr := key.String()
+			isDelete := bytes.Equal(values.Bytes(), common.Hash{}.Bytes())
 
 			env.sMu.Lock()
 			m, existed := env.StorageProofs[addrStr]
 			if !existed {
 				m = make(map[string][]hexutil.Bytes)
 				env.StorageProofs[addrStr] = m
+				if zktrieTracer.Available() {
+					env.zkTrieTracer[addrStr] = zktrieTracer
+				}
 			} else if _, existed := m[keyStr]; existed {
+				// still need to touch tracer for deletion
+				if isDelete && zktrieTracer.Available() {
+					env.zkTrieTracer[addrStr].MarkDeletion(key)
+				}
 				env.sMu.Unlock()
 				continue
 			}
 			env.sMu.Unlock()
 
-			proof, sibling, err := state.GetStorageTrieProof(addr, key)
+			var proof [][]byte
+			var err error
+			if zktrieTracer.Available() {
+				proof, err = state.GetSecureTrieProof(zktrieTracer, key)
+			} else {
+				proof, err = state.GetSecureTrieProof(trie, key)
+			}
 			if err != nil {
 				log.Error("Storage proof not available", "error", err, "address", addrStr, "key", keyStr)
 				// but we still mark the proofs map with nil array
@@ -325,8 +366,11 @@ func (api *API) getTxResult(env *traceEnv, state *state.StateDB, index int, bloc
 			}
 			env.sMu.Lock()
 			m[keyStr] = wrappedProof
-			if sibling != nil {
-				env.DeletionProofs = append(env.DeletionProofs, sibling)
+			if zktrieTracer.Available() {
+				if isDelete {
+					zktrieTracer.MarkDeletion(key)
+				}
+				env.zkTrieTracer[addrStr].Merge(zktrieTracer)
 			}
 			env.sMu.Unlock()
 		}

params/version.go

@@ -24,7 +24,7 @@ import (
 const (
 	VersionMajor = 3       // Major version component of the current release
 	VersionMinor = 1       // Minor version component of the current release
-	VersionPatch = 10      // Patch version component of the current release
+	VersionPatch = 11      // Patch version component of the current release
 	VersionMeta  = "alpha" // Version metadata to append to the version string
 )
 

trie/zk_trie.go

@@ -174,49 +174,29 @@ func (t *ZkTrie) NodeIterator(start []byte) NodeIterator {
 // nodes of the longest existing prefix of the key (at least the root node), ending
 // with the node that proves the absence of the key.
 func (t *ZkTrie) Prove(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) error {
-	// omit sibling, which is not required for proving only
-	_, err := t.ProveWithDeletion(key, fromLevel, proofDb)
-	return err
-}
+	err := t.ZkTrie.Prove(key, fromLevel, func(n *zktrie.Node) error {
+		nodeHash, err := n.NodeHash()
+		if err != nil {
+			return err
+		}
 
-// ProveWithDeletion is the implement of Prove, it also return possible sibling node
-// (if there is, i.e. the node of key exist and is not the only node in trie)
-// so witness generator can predict the final state root after deletion of this key
-// the returned sibling node has no key along with it for witness generator must decode
-// the node for its purpose
-func (t *ZkTrie) ProveWithDeletion(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) (sibling []byte, err error) {
-	err = t.ZkTrie.ProveWithDeletion(key, fromLevel,
-		func(n *zktrie.Node) error {
-			nodeHash, err := n.NodeHash()
-			if err != nil {
-				return err
+		if n.Type == zktrie.NodeTypeLeaf {
+			preImage := t.GetKey(n.NodeKey.Bytes())
+			if len(preImage) > 0 {
+				n.KeyPreimage = &zkt.Byte32{}
+				copy(n.KeyPreimage[:], preImage)
+				//return fmt.Errorf("key preimage not found for [%x] ref %x", n.NodeKey.Bytes(), k.Bytes())
 			}
-
-			if n.Type == zktrie.NodeTypeLeaf {
-				preImage := t.GetKey(n.NodeKey.Bytes())
-				if len(preImage) > 0 {
-					n.KeyPreimage = &zkt.Byte32{}
-					copy(n.KeyPreimage[:], preImage)
-					//return fmt.Errorf("key preimage not found for [%x] ref %x", n.NodeKey.Bytes(), k.Bytes())
-				}
-			}
-			return proofDb.Put(nodeHash[:], n.Value())
-		},
-		func(_ *zktrie.Node, n *zktrie.Node) {
-			// the sibling for each leaf should be unique except for EmptyNode
-			if n != nil && n.Type != zktrie.NodeTypeEmpty {
-				sibling = n.Value()
-			}
-		},
-	)
+		}
+		return proofDb.Put(nodeHash[:], n.Value())
+	})
 	if err != nil {
-		return
+		return err
 	}
 
 	// we put this special kv pair in db so we can distinguish the type and
 	// make suitable Proof
-	err = proofDb.Put(magicHash, zktrie.ProofMagicBytes())
-	return
+	return proofDb.Put(magicHash, zktrie.ProofMagicBytes())
 }
 
 // VerifyProof checks merkle proofs. The given proof must contain the value for

trie/zk_trie_proof_test.go

@@ -196,7 +196,7 @@ func randomZktrie(t *testing.T, n int) (*ZkTrie, map[string]*kv) {
 	return tr, vals
 }
 
-// Tests that new "proof with deletion" feature
+// Tests that new "proof trace" feature
 func TestProofWithDeletion(t *testing.T) {
 	tr, _ := NewZkTrie(common.Hash{}, NewZktrieDatabase((memorydb.New())))
 	mt := &zkTrieImplTestWrapper{tr.Tree()}
@@ -217,20 +217,66 @@ func TestProofWithDeletion(t *testing.T) {
 	s_key1, err := zkt.ToSecureKeyBytes(key1)
 	assert.NoError(t, err)
 
-	sibling1, err := tr.ProveWithDeletion(s_key1.Bytes(), 0, proof)
+	proofTracer := tr.NewProofTracer()
+
+	err = proofTracer.Prove(s_key1.Bytes(), 0, proof)
 	assert.NoError(t, err)
 	nd, err := tr.TryGet(key2)
 	assert.NoError(t, err)
-	l := len(sibling1)
-	// a hacking to grep the value part directly from the encoded leaf node,
-	// notice the sibling of key `k*32`` is just the leaf of key `m*32`
-	assert.Equal(t, sibling1[l-33:l-1], nd)
 
 	s_key2, err := zkt.ToSecureKeyBytes(bytes.Repeat([]byte("x"), 32))
 	assert.NoError(t, err)
 
-	sibling2, err := tr.ProveWithDeletion(s_key2.Bytes(), 0, proof)
+	err = proofTracer.Prove(s_key2.Bytes(), 0, proof)
 	assert.NoError(t, err)
-	assert.Nil(t, sibling2)
+	// assert.Equal(t, len(sibling1), len(delTracer.GetProofs()))
 
+	siblings, err := proofTracer.GetDeletionProofs()
+	assert.NoError(t, err)
+	assert.Equal(t, 0, len(siblings))
+
+	proofTracer.MarkDeletion(s_key1.Bytes())
+	siblings, err = proofTracer.GetDeletionProofs()
+	assert.NoError(t, err)
+	assert.Equal(t, 1, len(siblings))
+	l := len(siblings[0])
+	// a hacking to grep the value part directly from the encoded leaf node,
+	// notice the sibling of key `k*32`` is just the leaf of key `m*32`
+	assert.Equal(t, siblings[0][l-33:l-1], nd)
+
+	// no effect
+	proofTracer.MarkDeletion(s_key2.Bytes())
+	siblings, err = proofTracer.GetDeletionProofs()
+	assert.NoError(t, err)
+	assert.Equal(t, 1, len(siblings))
+
+	key3 := bytes.Repeat([]byte("x"), 32)
+	err = mt.UpdateWord(
+		zkt.NewByte32FromBytesPaddingZero(key3),
+		zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("z"), 32)),
+	)
+	assert.NoError(t, err)
+
+	proofTracer = tr.NewProofTracer()
+	err = proofTracer.Prove(s_key1.Bytes(), 0, proof)
+	assert.NoError(t, err)
+	err = proofTracer.Prove(s_key2.Bytes(), 0, proof)
+	assert.NoError(t, err)
+
+	proofTracer.MarkDeletion(s_key1.Bytes())
+	siblings, err = proofTracer.GetDeletionProofs()
+	assert.NoError(t, err)
+	assert.Equal(t, 1, len(siblings))
+
+	proofTracer.MarkDeletion(s_key2.Bytes())
+	siblings, err = proofTracer.GetDeletionProofs()
+	assert.NoError(t, err)
+	assert.Equal(t, 2, len(siblings))
+
+	// one of the siblings is just leaf for key2, while
+	// another one must be a middle node
+	match1 := bytes.Equal(siblings[0][l-33:l-1], nd)
+	match2 := bytes.Equal(siblings[1][l-33:l-1], nd)
+	assert.True(t, match1 || match2)
+	assert.False(t, match1 && match2)
 }

trie/zktrie_deletionproof.go

@@ -0,0 +1,163 @@
+package trie
+
+import (
+	"bytes"
+
+	zktrie "github.com/scroll-tech/zktrie/trie"
+	zkt "github.com/scroll-tech/zktrie/types"
+
+	"github.com/scroll-tech/go-ethereum/ethdb"
+)
+
+// Pick Node from its hash directly from database, notice it has different
+// interface with the function of same name in `trie`
+func (t *ZkTrie) TryGetNode(nodeHash *zkt.Hash) (*zktrie.Node, error) {
+	if bytes.Equal(nodeHash[:], zkt.HashZero[:]) {
+		return zktrie.NewEmptyNode(), nil
+	}
+	nBytes, err := t.db.Get(nodeHash[:])
+	if err == zktrie.ErrKeyNotFound {
+		return nil, zktrie.ErrKeyNotFound
+	} else if err != nil {
+		return nil, err
+	}
+	return zktrie.NewNodeFromBytes(nBytes)
+}
+
+type ProofTracer struct {
+	*ZkTrie
+	deletionTracer map[zkt.Hash]struct{}
+	rawPaths       map[string][]*zktrie.Node
+}
+
+// NewProofTracer create a proof tracer object
+func (t *ZkTrie) NewProofTracer() *ProofTracer {
+	return &ProofTracer{
+		ZkTrie: t,
+		// always consider 0 is "deleted"
+		deletionTracer: map[zkt.Hash]struct{}{zkt.HashZero: {}},
+		rawPaths:       make(map[string][]*zktrie.Node),
+	}
+}
+
+// Merge merge the input tracer into current and return current tracer
+func (t *ProofTracer) Merge(another *ProofTracer) *ProofTracer {
+
+	// sanity checking
+	if !bytes.Equal(t.Hash().Bytes(), another.Hash().Bytes()) {
+		panic("can not merge two proof tracer base on different trie")
+	}
+
+	for k := range another.deletionTracer {
+		t.deletionTracer[k] = struct{}{}
+	}
+
+	for k, v := range another.rawPaths {
+		t.rawPaths[k] = v
+	}
+
+	return t
+}
+
+// GetDeletionProofs generate current deletionTracer and collect deletion proofs
+// which is possible to be used from all rawPaths, which enabling witness generator
+// to predict the final state root after executing any deletion
+// along any of the rawpath, no matter of the deletion occurs in any position of the mpt ops
+// Note the collected sibling node has no key along with it since witness generator would
+// always decode the node for its purpose
+func (t *ProofTracer) GetDeletionProofs() ([][]byte, error) {
+
+	retMap := map[zkt.Hash][]byte{}
+
+	// check each path: reversively, skip the final leaf node
+	for _, path := range t.rawPaths {
+
+		checkPath := path[:len(path)-1]
+		for i := len(checkPath); i > 0; i-- {
+			n := checkPath[i-1]
+			_, deletedL := t.deletionTracer[*n.ChildL]
+			_, deletedR := t.deletionTracer[*n.ChildR]
+			if deletedL && deletedR {
+				nodeHash, _ := n.NodeHash()
+				t.deletionTracer[*nodeHash] = struct{}{}
+			} else {
+				var siblingHash *zkt.Hash
+				if deletedL {
+					siblingHash = n.ChildR
+				} else if deletedR {
+					siblingHash = n.ChildL
+				}
+				if siblingHash != nil {
+					sibling, err := t.TryGetNode(siblingHash)
+					if err != nil {
+						return nil, err
+					}
+					if sibling.Type != zktrie.NodeTypeEmpty {
+						retMap[*siblingHash] = sibling.Value()
+					}
+				}
+				break
+			}
+		}
+	}
+
+	var ret [][]byte
+	for _, bt := range retMap {
+		ret = append(ret, bt)
+	}
+
+	return ret, nil
+
+}
+
+// MarkDeletion mark a key has been involved into deletion
+func (t *ProofTracer) MarkDeletion(key []byte) {
+	if path, existed := t.rawPaths[string(key)]; existed {
+		// sanity check
+		leafNode := path[len(path)-1]
+		if leafNode.Type != zktrie.NodeTypeLeaf {
+			panic("all path recorded in proofTrace should be ended with leafNode")
+		}
+
+		nodeHash, _ := leafNode.NodeHash()
+		t.deletionTracer[*nodeHash] = struct{}{}
+	}
+}
+
+// Prove act the same as zktrie.Prove, while also collect the raw path
+// for collecting deletion proofs in a post-work
+func (t *ProofTracer) Prove(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) error {
+	var mptPath []*zktrie.Node
+	err := t.ZkTrie.ProveWithDeletion(key, fromLevel,
+		func(n *zktrie.Node) error {
+			nodeHash, err := n.NodeHash()
+			if err != nil {
+				return err
+			}
+
+			if n.Type == zktrie.NodeTypeLeaf {
+				preImage := t.GetKey(n.NodeKey.Bytes())
+				if len(preImage) > 0 {
+					n.KeyPreimage = &zkt.Byte32{}
+					copy(n.KeyPreimage[:], preImage)
+				}
+			} else if n.Type == zktrie.NodeTypeParent {
+				mptPath = append(mptPath, n)
+			}
+
+			return proofDb.Put(nodeHash[:], n.Value())
+		},
+		func(n *zktrie.Node, _ *zktrie.Node) {
+			// only "hit" path (i.e. the leaf node corresponding the input key can be found)
+			// would be add into tracer
+			mptPath = append(mptPath, n)
+			t.rawPaths[string(key)] = mptPath
+		},
+	)
+	if err != nil {
+		return err
+	}
+	// we put this special kv pair in db so we can distinguish the type and
+	// make suitable Proof
+	return proofDb.Put(magicHash, zktrie.ProofMagicBytes())
+}