Fix witness generator for failed call/create op and tx (#146)

* update test examples

* update tests

* solidify account deleting

* update zkwriter for reverting

* solid handling for call, and update test examples

* fix status capture for CREATE issue, update test examples

* add fail case for deploy

* update test example (call non-exist address)

* move test and stuffs to internal path, make required exports

* fix issue in selfdestruct

* update destruct test

* lint

* rename traces

core/state/statedb.go

@@ -599,10 +599,10 @@ func (s *StateDB) getDeletedStateObject(addr common.Address) *stateObject {
 		if len(enc) == 0 {
 			return nil
 		}
-		data = new(types.StateAccount)
 		if s.IsZktrie() {
 			data, err = types.UnmarshalStateAccount(enc)
 		} else {
+			data = new(types.StateAccount)
 			err = rlp.DecodeBytes(enc, data)
 		}
 		if err != nil {

core/types/l2trace.go

@@ -124,6 +124,9 @@ type ExtraData struct {
 	// STATICCALL: [stack.nth_last(1) (i.e. callee) address’s account,
 	//					  callee contract address's account (before called)]
 	StateList []*AccountWrapper `json:"proofList,omitempty"`
+	// The status of caller, it would be captured twice:
+	// 1. before execution and 2. updated in CaptureEnter (for CALL/CALLCODE it duplicated with StateList[0])
+	Caller []*AccountWrapper `json:"caller,omitempty"`
 }
 
 type AccountWrapper struct {

core/vm/logger.go

@@ -280,6 +280,12 @@ func (l *StructLogger) CaptureState(pc uint64, op OpCode, gas, cost uint64, scop
 			}
 		}
 	}
+	// for each "calling" op, pick the caller's state
+	switch op {
+	case CALL, CALLCODE, STATICCALL, DELEGATECALL, CREATE, CREATE2:
+		extraData := structlog.getOrInitExtraData()
+		extraData.Caller = append(extraData.Caller, getWrappedAccountForAddr(l, scope.Contract.Address()))
+	}
 
 	structlog.RefundCounter = l.env.StateDB.GetRefund()
 	l.logs = append(l.logs, structlog)
@@ -323,6 +329,10 @@ func (l *StructLogger) CaptureEnter(typ OpCode, from common.Address, to common.A
 	// append extraData part for the log, capture the account status (the nonce / balance has been updated in capture enter)
 	wrappedStatus := getWrappedAccountForAddr(l, to)
 	theLog.ExtraData.StateList = append(theLog.ExtraData.StateList, wrappedStatus)
+	// finally we update the caller's status (it is possible that nonce and balance being updated)
+	if len(theLog.ExtraData.Caller) == 1 {
+		theLog.ExtraData.Caller = append(theLog.ExtraData.Caller, getWrappedAccountForAddr(l, from))
+	}
 }
 
 // in CaptureExit phase, a CREATE has its target address's code being set and queryable

core/vm/logger_trace.go

@@ -15,8 +15,8 @@ var (
 		CALLCODE:     {traceToAddressCode, traceLastNAddressCode(1), traceContractAccount, traceLastNAddressAccount(1)}, // contract account is the caller, stack.nth_last(1) is the callee's address
 		DELEGATECALL: {traceToAddressCode, traceLastNAddressCode(1)},
 		STATICCALL:   {traceToAddressCode, traceLastNAddressCode(1), traceLastNAddressAccount(1)},
-		CREATE:       {}, // sender is already recorded in ExecutionResult, callee is recorded in CaptureEnter&CaptureExit
+		CREATE:       {}, // caller is already recorded in ExtraData.Caller, callee is recorded in CaptureEnter&CaptureExit
-		CREATE2:      {}, // sender is already recorded in ExecutionResult, callee is recorded in CaptureEnter&CaptureExit
+		CREATE2:      {}, // caller is already recorded in ExtraData.Caller, callee is recorded in CaptureEnter&CaptureExit
 		SLOAD:        {}, // trace storage in `captureState` instead of here, to handle `l.cfg.DisableStorage` flag
 		SSTORE:       {}, // trace storage in `captureState` instead of here, to handle `l.cfg.DisableStorage` flag
 		SELFDESTRUCT: {traceContractAccount, traceLastNAddressAccount(0)},

internal/utesting/blocktraces/zktrie_gen_witness_test.go

@@ -1,4 +1,4 @@
-package zkproof
+package blocktraces
 
 import (
 	"encoding/json"
@@ -9,6 +9,7 @@ import (
 
 	"github.com/scroll-tech/go-ethereum/common"
 	"github.com/scroll-tech/go-ethereum/core/types"
+	"github.com/scroll-tech/go-ethereum/trie/zkproof"
 )
 
 func init() {
@@ -16,7 +17,7 @@ func init() {
 	var orderSchemeI int
 	if orderScheme != "" {
 		if n, err := fmt.Sscanf(orderScheme, "%d", &orderSchemeI); err == nil && n == 1 {
-			usedOrdererScheme = MPTWitnessType(orderSchemeI)
+			zkproof.SetOrderScheme(zkproof.MPTWitnessType(orderSchemeI))
 		}
 	}
 }
@@ -43,59 +44,47 @@ func loadStaff(t *testing.T, fname string) *types.BlockResult {
 }
 
 func TestWriterCreation(t *testing.T) {
-	trace := loadStaff(t, "deploy_trace.json")
+	trace := loadStaff(t, "deploy.json")
-	writer, err := NewZkTrieProofWriter(trace.StorageTrace)
+	writer, err := zkproof.NewZkTrieProofWriter(trace.StorageTrace)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	if len(writer.tracingAccounts) != 4 {
+	if len(writer.TracingAccounts()) != 3 {
-		t.Error("unexpected tracing account data", writer.tracingAccounts)
+		t.Error("unexpected tracing account data", writer.TracingAccounts())
 	}
 
-	if v, existed := writer.tracingAccounts[common.HexToAddress("0x000000000000000000636F6e736F6c652e6c6f67")]; !existed || v != nil {
+	if v, existed := writer.TracingAccounts()[common.HexToAddress("0x08c683b684d1e24cab8ce6de5c8c628d993ac140")]; !existed || v != nil {
 		t.Error("wrong tracing status for uninited address", v, existed)
 	}
 
-	if v, existed := writer.tracingAccounts[common.HexToAddress("0xb36feAEaF76c2A33335b73bEF9aEf7a23d9af1e3")]; !existed || v != nil {
+	if v, existed := writer.TracingAccounts()[common.HexToAddress("0x4cb1aB63aF5D8931Ce09673EbD8ae2ce16fD6571")]; !existed || v == nil {
-		t.Error("wrong tracing status for uninited address", v, existed)
-	}
-
-	if v, existed := writer.tracingAccounts[common.HexToAddress("0x4cb1aB63aF5D8931Ce09673EbD8ae2ce16fD6571")]; !existed || v == nil {
 		t.Error("wrong tracing status for establied address", v, existed)
 	}
 
-	if len(writer.tracingStorageTries) != 1 {
-		t.Error("unexpected tracing storage data", writer.tracingStorageTries)
-	}
-
-	if v, existed := writer.tracingStorageTries[common.HexToAddress("0xb36feAEaF76c2A33335b73bEF9aEf7a23d9af1e3")]; !existed || v == nil {
-		t.Error("wrong tracing storage statu", existed, v)
-	}
-
 }
 
 func TestGreeterTx(t *testing.T) {
-	trace := loadStaff(t, "greeter_trace.json")
+	trace := loadStaff(t, "greeter.json")
-	writer, err := NewZkTrieProofWriter(trace.StorageTrace)
+	writer, err := zkproof.NewZkTrieProofWriter(trace.StorageTrace)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	od := &simpleOrderer{}
+	od := zkproof.NewSimpleOrderer()
 	theTx := trace.ExecutionResults[0]
-	handleTx(od, theTx)
+	zkproof.HandleTx(od, theTx)
 
 	t.Log(od)
 
-	for _, op := range od.savedOp {
+	for _, op := range od.SavedOp() {
 		_, err = writer.HandleNewState(op)
 		if err != nil {
 			t.Fatal(err)
 		}
 	}
 
-	traces, err := HandleBlockResult(trace)
+	traces, err := zkproof.HandleBlockResult(trace)
 	t.Log("traces: ", len(traces))
 	outObj, _ := json.Marshal(traces)
 	t.Log(string(outObj))
@@ -105,8 +94,8 @@ func TestGreeterTx(t *testing.T) {
 }
 
 func TestTokenTx(t *testing.T) {
-	trace := loadStaff(t, "token_trace.json")
+	trace := loadStaff(t, "token.json")
-	traces, err := HandleBlockResult(trace)
+	traces, err := zkproof.HandleBlockResult(trace)
 	outObj, _ := json.Marshal(traces)
 	t.Log(string(outObj))
 	if err != nil {
@@ -116,16 +105,16 @@ func TestTokenTx(t *testing.T) {
 }
 
 func TestCallTx(t *testing.T) {
-	trace := loadStaff(t, "call_trace.json")
+	trace := loadStaff(t, "call.json")
-	traces, err := HandleBlockResult(trace)
+	traces, err := zkproof.HandleBlockResult(trace)
 	outObj, _ := json.Marshal(traces)
 	t.Log(string(outObj))
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	trace = loadStaff(t, "staticcall_trace.json")
+	trace = loadStaff(t, "call_edge.json")
-	traces, err = HandleBlockResult(trace)
+	traces, err = zkproof.HandleBlockResult(trace)
 	outObj, _ = json.Marshal(traces)
 	t.Log(string(outObj))
 	if err != nil {
@@ -134,16 +123,16 @@ func TestCallTx(t *testing.T) {
 }
 
 func TestCreateTx(t *testing.T) {
-	trace := loadStaff(t, "create_trace.json")
+	trace := loadStaff(t, "create.json")
-	traces, err := HandleBlockResult(trace)
+	traces, err := zkproof.HandleBlockResult(trace)
 	outObj, _ := json.Marshal(traces)
 	t.Log(string(outObj))
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	trace = loadStaff(t, "deploy_trace.json")
+	trace = loadStaff(t, "deploy.json")
-	traces, err = HandleBlockResult(trace)
+	traces, err = zkproof.HandleBlockResult(trace)
 	outObj, _ = json.Marshal(traces)
 	t.Log(string(outObj))
 	if err != nil {
@@ -153,16 +142,16 @@ func TestCreateTx(t *testing.T) {
 }
 
 func TestFailedCallTx(t *testing.T) {
-	trace := loadStaff(t, "fail_call_trace.json")
+	trace := loadStaff(t, "fail_call.json")
-	traces, err := HandleBlockResult(trace)
+	traces, err := zkproof.HandleBlockResult(trace)
 	outObj, _ := json.Marshal(traces)
 	t.Log(string(outObj))
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	trace = loadStaff(t, "fail_create_trace.json")
+	trace = loadStaff(t, "fail_create.json")
-	traces, err = HandleBlockResult(trace)
+	traces, err = zkproof.HandleBlockResult(trace)
 	outObj, _ = json.Marshal(traces)
 	t.Log(string(outObj))
 	if err != nil {
@@ -173,8 +162,19 @@ func TestFailedCallTx(t *testing.T) {
 
 //notice: now only work with OP_ORDER=2
 func TestDeleteTx(t *testing.T) {
-	trace := loadStaff(t, "delete_trace.json")
+	trace := loadStaff(t, "delete.json")
-	traces, err := HandleBlockResult(trace)
+	traces, err := zkproof.HandleBlockResult(trace)
+	outObj, _ := json.Marshal(traces)
+	t.Log(string(outObj))
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+//notice: now only work with OP_ORDER=2
+func TestDestructTx(t *testing.T) {
+	trace := loadStaff(t, "destruct.json")
+	traces, err := zkproof.HandleBlockResult(trace)
 	outObj, _ := json.Marshal(traces)
 	t.Log(string(outObj))
 	if err != nil {
@@ -184,7 +184,7 @@ func TestDeleteTx(t *testing.T) {
 
 func TestMutipleTx(t *testing.T) {
 	trace := loadStaff(t, "multi_txs.json")
-	traces, err := HandleBlockResult(trace)
+	traces, err := zkproof.HandleBlockResult(trace)
 	outObj, _ := json.Marshal(traces)
 	t.Log(string(outObj))
 	if err != nil {

trie/zk_trie.go

@@ -162,11 +162,8 @@ func (t *ZkTrie) TryDelete(key []byte) error {
 		return nil
 	}
 
-	zeroBt := common.Hash{}
+	// FIXME: when tryDelete get more solid test, use it instead
-	// FIXME: delete should not be implemented as Update(0)
+	return t.tree.tryDeleteLite(zkt.NewHashFromBigInt(k))
-	return t.tree.TryUpdate(k.Bytes(), zeroBt[:])
-	//kPreimage := smt.NewByte32FromBytesPadding(key)
-	//return t.tree.DeleteWord(kPreimage)
 }
 
 // GetKey returns the preimage of a hashed key that was

trie/zk_trie_impl.go

@@ -444,76 +444,6 @@ func (mt *ZkTrieImpl) GetLeafNodeByWord(kPreimage *zkt.Byte32) (*Node, error) {
 	return n, err
 }
 
-/*
-// Update function updates the value of a specified key in the ZkTrieImpl, and updates
-// the path from the leaf to the Root with the new values,and returns the
-// CircomProcessorProof.
-func (mt *ZkTrieImpl) Update(k, v *big.Int, kPreimage *zkt.Byte32, vPreimage []byte) error {
-	// verify that the ZkTrieImpl is writable
-	if !mt.writable {
-		return ErrNotWritable
-	}
-
-	// verify that k & are valid and fit inside the Finite Field.
-	if !cryptoUtils.CheckBigIntInField(k) {
-		return errors.New("Key not inside the Finite Field")
-	}
-	if !cryptoUtils.CheckBigIntInField(v) {
-		return errors.New("Key not inside the Finite Field")
-	}
-
-	kHash := zkt.NewHashFromBigInt(k)
-	vHash := zkt.NewHashFromBigInt(v)
-	path := getPath(mt.maxLevels, kHash[:])
-
-	nextKey := mt.rootKey
-	siblings := []*zkt.Hash{}
-	for i := 0; i < mt.maxLevels; i++ {
-		n, err := mt.GetNode(nextKey)
-		if err != nil {
-			return err
-		}
-		switch n.Type {
-		case NodeTypeEmpty:
-			return ErrKeyNotFound
-		case NodeTypeLeaf:
-			if bytes.Equal(kHash[:], n.NodeKey[:]) {
-				// update leaf and upload to the root
-				newNodeLeaf := NewNodeLeaf(kHash, vHash, kPreimage, vPreimage)
-				_, err := mt.updateNode(newNodeLeaf)
-				if err != nil {
-					return err
-				}
-				newRootKey, err :=
-					mt.recalculatePathUntilRoot(path, newNodeLeaf, siblings)
-				if err != nil {
-					return err
-				}
-				mt.rootKey = newRootKey
-				err = mt.dbInsert(dbKeyRootNode, DBEntryTypeRoot, mt.rootKey[:])
-				if err != nil {
-					return err
-				}
-				return nil
-			}
-			return ErrKeyNotFound
-		case NodeTypeMiddle:
-			if path[i] {
-				nextKey = n.ChildR
-				siblings = append(siblings, n.ChildL)
-			} else {
-				nextKey = n.ChildL
-				siblings = append(siblings, n.ChildR)
-			}
-		default:
-			return ErrInvalidNodeFound
-		}
-	}
-
-	return ErrKeyNotFound
-}
-*/
-
 // Deprecated: only for testing
 func (mt *ZkTrieImpl) UpdateWord(kPreimage, vPreimage *zkt.Byte32) error {
 	k, err := kPreimage.Hash()
@@ -589,6 +519,57 @@ func (mt *ZkTrieImpl) UpdateVarWord(kPreimage *zkt.Byte32, vHash *big.Int, vPrei
 }
 */
 
+// only update the corresponding leaf for
+func (mt *ZkTrieImpl) tryDeleteLite(kHash *zkt.Hash) error {
+	// verify that the ZkTrieImpl is writable
+	if !mt.writable {
+		return ErrNotWritable
+	}
+
+	// verify that k is valid and fit inside the Finite Field.
+	if !cryptoUtils.CheckBigIntInField(kHash.BigInt()) {
+		return errors.New("Key not inside the Finite Field")
+	}
+
+	path := getPath(mt.maxLevels, kHash[:])
+
+	nextKey := mt.rootKey
+	siblings := []*zkt.Hash{}
+	for i := 0; i < mt.maxLevels; i++ {
+		n, err := mt.GetNode(nextKey)
+		if err != nil {
+			return err
+		}
+		switch n.Type {
+		case NodeTypeEmpty:
+			return ErrKeyNotFound
+		case NodeTypeLeaf:
+			if bytes.Equal(kHash[:], n.NodeKey[:]) {
+				// remove and go up with the sibling
+				newRootKey, err := mt.recalculatePathUntilRoot(path, NewNodeEmpty(), siblings)
+				if err != nil {
+					return err
+				}
+				mt.rootKey = newRootKey
+				return mt.dbInsert(dbKeyRootNode, DBEntryTypeRoot, mt.rootKey[:])
+			}
+			return ErrKeyNotFound
+		case NodeTypeMiddle:
+			if path[i] {
+				nextKey = n.ChildR
+				siblings = append(siblings, n.ChildL)
+			} else {
+				nextKey = n.ChildL
+				siblings = append(siblings, n.ChildR)
+			}
+		default:
+			return ErrInvalidNodeFound
+		}
+	}
+
+	return ErrKeyNotFound
+}
+
 // Delete removes the specified Key from the ZkTrieImpl and updates the path
 // from the deleted key to the Root with the new values.  This method removes
 // the key from the ZkTrieImpl, but does not remove the old nodes from the
@@ -645,6 +626,23 @@ func (mt *ZkTrieImpl) tryDelete(kHash *zkt.Hash) error {
 	return ErrKeyNotFound
 }
 
+// Delete removes the specified Key from the ZkTrieImpl and updates the path
+// from the deleted key to the Root with the new values.  This method removes
+// the key from the ZkTrieImpl, but does not remove the old nodes from the
+// key-value database; this means that if the tree is accessed by an old Root
+// where the key was not deleted yet, the key will still exist. If is desired
+// to remove the key-values from the database that are not under the current
+// Root, an option could be to dump all the leafs (using mt.DumpLeafs) and
+// import them in a new ZkTrieImpl in a new database (using
+// mt.ImportDumpedLeafs), but this will lose all the Root history of the
+// ZkTrieImpl
+// Delete removes any existing value for key from the trie.
+func (t *ZkTrieImpl) Delete(key []byte) {
+	if err := t.TryDelete(key); err != nil {
+		log.Error(fmt.Sprintf("Unhandled trie error: %v", err))
+	}
+}
+
 // TryDelete removes any existing value for key from the trie.
 // If a node was not found in the database, a MissingNodeError is returned.
 func (mt *ZkTrieImpl) TryDelete(key []byte) error {

trie/zkproof/orderer.go

@@ -14,7 +14,9 @@ type opIterator interface {
 }
 
 type opOrderer interface {
+	readonly(bool)
 	absorb(*types.AccountWrapper)
+	absorbStorage(*types.AccountWrapper, *types.StorageWrapper)
 	end_absorb() opIterator
 }
 
@@ -33,13 +35,33 @@ func (ops *iterateOp) next() *types.AccountWrapper {
 }
 
 type simpleOrderer struct {
+	readOnly int
 	savedOp  []*types.AccountWrapper
 }
 
+func (od *simpleOrderer) SavedOp() []*types.AccountWrapper { return od.savedOp }
+
+func (od *simpleOrderer) readonly(mode bool) {
+	if mode {
+		od.readOnly += 1
+	} else if od.readOnly == 0 {
+		panic("unexpected readonly mode stack pop")
+	} else {
+		od.readOnly -= 1
+	}
+}
+
 func (od *simpleOrderer) absorb(st *types.AccountWrapper) {
+	if od.readOnly > 0 {
+		return
+	}
 	od.savedOp = append(od.savedOp, st)
 }
 
+func (od *simpleOrderer) absorbStorage(st *types.AccountWrapper, _ *types.StorageWrapper) {
+	od.absorb(st)
+}
+
 func (od *simpleOrderer) end_absorb() opIterator {
 	ret := iterateOp(od.savedOp)
 	return &ret
@@ -69,6 +91,15 @@ func (opss *multiOpIterator) next() *types.AccountWrapper {
 }
 
 type rwTblOrderer struct {
+	readOnly         int
+	readOnlySnapshot struct {
+		accounts map[string]*types.AccountWrapper
+		storages map[string]map[string]*types.StorageWrapper
+	}
+	initedData map[common.Address]*types.AccountWrapper
+
+	// help to track all accounts being touched, and provide the
+	// completed account status for storage updating
 	traced map[string]*types.AccountWrapper
 
 	opAccNonce    map[string]*types.AccountWrapper
@@ -77,21 +108,24 @@ type rwTblOrderer struct {
 	opStorage     map[string]map[string]*types.StorageWrapper
 }
 
-func newRWTblOrderer(inited map[common.Address]*types.StateAccount) *rwTblOrderer {
+func NewSimpleOrderer() *simpleOrderer { return &simpleOrderer{} }
 
-	traced := make(map[string]*types.AccountWrapper)
+func NewRWTblOrderer(inited map[common.Address]*types.StateAccount) *rwTblOrderer {
 
+	initedAcc := make(map[common.Address]*types.AccountWrapper)
 	for addr, data := range inited {
 		if data == nil {
-			continue
+			initedAcc[addr] = &types.AccountWrapper{
+				Address: addr,
+				Balance: (*hexutil.Big)(big.NewInt(0)),
 			}
-
+		} else {
 			bl := data.Balance
 			if bl == nil {
 				bl = big.NewInt(0)
 			}
 
-		traced[addr.String()] = &types.AccountWrapper{
+			initedAcc[addr] = &types.AccountWrapper{
 				Address:  addr,
 				Nonce:    data.Nonce,
 				Balance:  (*hexutil.Big)(bl),
@@ -99,8 +133,11 @@ func newRWTblOrderer(inited map[common.Address]*types.StateAccount) *rwTblOrdere
 			}
 		}
 
+	}
+
 	return &rwTblOrderer{
-		traced:        traced,
+		initedData:    initedAcc,
+		traced:        make(map[string]*types.AccountWrapper),
 		opAccNonce:    make(map[string]*types.AccountWrapper),
 		opAccBalance:  make(map[string]*types.AccountWrapper),
 		opAccCodeHash: make(map[string]*types.AccountWrapper),
@@ -108,16 +145,101 @@ func newRWTblOrderer(inited map[common.Address]*types.StateAccount) *rwTblOrdere
 	}
 }
 
+func (od *rwTblOrderer) readonly(mode bool) {
+	if mode {
+		if od.readOnly == 0 {
+			od.readOnlySnapshot.accounts = make(map[string]*types.AccountWrapper)
+			od.readOnlySnapshot.storages = make(map[string]map[string]*types.StorageWrapper)
+		}
+		od.readOnly += 1
+	} else if od.readOnly == 0 {
+		panic("unexpected readonly mode stack pop")
+	} else {
+		od.readOnly -= 1
+		if od.readOnly == 0 {
+			for addrS, st := range od.readOnlySnapshot.accounts {
+				od.absorb(st)
+				if m, existed := od.readOnlySnapshot.storages[addrS]; existed {
+					for _, stg := range m {
+						st.Storage = stg
+						od.absorbStorage(st, nil)
+					}
+				}
+			}
+		}
+	}
+}
+
+func (od *rwTblOrderer) absorbStorage(st *types.AccountWrapper, before *types.StorageWrapper) {
+	if st.Storage == nil {
+		panic("do not call absorbStorage ")
+	}
+
+	od.absorb(st)
+	addrStr := st.Address.String()
+
+	if stg := st.Storage; stg != nil {
+		m, existed := od.opStorage[addrStr]
+		if !existed {
+			m = make(map[string]*types.StorageWrapper)
+			od.opStorage[addrStr] = m
+		}
+
+		// key must be unified into 32 bytes
+		keyBytes := hexutil.MustDecode(stg.Key)
+		keyStr := common.BytesToHash(keyBytes).String()
+
+		// trace every "touched" status for readOnly
+		if od.readOnly > 0 {
+			m, existed := od.readOnlySnapshot.storages[addrStr]
+			if !existed {
+				m = make(map[string]*types.StorageWrapper)
+				od.readOnlySnapshot.storages[addrStr] = m
+			}
+			if _, hashTraced := m[keyStr]; !hashTraced {
+				if before != nil {
+					m[keyStr] = before
+				} else {
+					m[keyStr] = stg
+				}
+
+			}
+		}
+
+		m[keyStr] = stg
+	}
+
+}
+
 func (od *rwTblOrderer) absorb(st *types.AccountWrapper) {
 
+	initedRef, existed := od.initedData[st.Address]
+	if !existed {
+		panic("encounter unprepared status")
+	}
+
 	addrStr := st.Address.String()
 
-	start, existed := od.traced[addrStr]
+	// trace every "touched" status for readOnly
+	if od.readOnly > 0 {
+		snapShot, existed := od.traced[addrStr]
 		if !existed {
-		start = &types.AccountWrapper{
+			snapShot = initedRef
+		}
+
+		if _, hasTraced := od.readOnlySnapshot.accounts[addrStr]; !hasTraced {
+			od.readOnlySnapshot.accounts[addrStr] = copyAccountState(snapShot)
+		}
+	}
+
+	if isDeletedAccount(st) {
+		// for account delete, made a safer data for status
+		st = &types.AccountWrapper{
+			Address: st.Address,
 			Balance: (*hexutil.Big)(big.NewInt(0)),
 		}
 	}
+
 	od.traced[addrStr] = st
 
 	// notice there would be at least one entry for all 3 fields when accessing an address
@@ -128,8 +250,8 @@ func (od *rwTblOrderer) absorb(st *types.AccountWrapper) {
 
 	if traced, existed := od.opAccNonce[addrStr]; !existed {
 		traced = copyAccountState(st)
-		traced.Balance = start.Balance
+		traced.Balance = initedRef.Balance
-		traced.CodeHash = start.CodeHash
+		traced.CodeHash = initedRef.CodeHash
 		traced.Storage = nil
 		od.opAccNonce[addrStr] = traced
 	} else {
@@ -138,7 +260,7 @@ func (od *rwTblOrderer) absorb(st *types.AccountWrapper) {
 
 	if traced, existed := od.opAccBalance[addrStr]; !existed {
 		traced = copyAccountState(st)
-		traced.CodeHash = start.CodeHash
+		traced.CodeHash = initedRef.CodeHash
 		traced.Storage = nil
 		od.opAccBalance[addrStr] = traced
 	} else {
@@ -156,17 +278,6 @@ func (od *rwTblOrderer) absorb(st *types.AccountWrapper) {
 		traced.CodeHash = st.CodeHash
 	}
 
-	if stg := st.Storage; stg != nil {
-		m, existed := od.opStorage[addrStr]
-		if !existed {
-			m = make(map[string]*types.StorageWrapper)
-			od.opStorage[addrStr] = m
-		}
-
-		// key must be unified into 32 bytes
-		keyBytes := hexutil.MustDecode(stg.Key)
-		m[common.BytesToHash(keyBytes).String()] = stg
-	}
 }
 
 func (od *rwTblOrderer) end_absorb() opIterator {
@@ -184,6 +295,7 @@ func (od *rwTblOrderer) end_absorb() opIterator {
 	var iterStorage []*types.AccountWrapper
 
 	for _, addrStr := range sortedAddrs {
+
 		if v, existed := od.opAccNonce[addrStr]; existed {
 			iterNonce = append(iterNonce, v)
 		}

trie/zkproof/writer.go

@@ -143,6 +143,10 @@ type zktrieProofWriter struct {
 	tracingAccounts     map[common.Address]*types.StateAccount
 }
 
+func (wr *zktrieProofWriter) TracingAccounts() map[common.Address]*types.StateAccount {
+	return wr.tracingAccounts
+}
+
 func NewZkTrieProofWriter(storage *types.StorageTrace) (*zktrieProofWriter, error) {
 
 	underlayerDb := memorydb.New()
@@ -232,6 +236,8 @@ const (
 	posCREATEAfter  = 1
 	posCALL         = 2
 	posSTATICCALL   = 0
+
+//	posSELFDESTRUCT = 2
 )
 
 func getAccountState(l *types.StructLogRes, pos int) *types.AccountWrapper {
@@ -263,7 +269,16 @@ func copyAccountState(st *types.AccountWrapper) *types.AccountWrapper {
 	}
 }
 
+func isDeletedAccount(state *types.AccountWrapper) bool {
+	return state.Nonce == 0 && bytes.Equal(state.CodeHash.Bytes(), common.Hash{}.Bytes())
+}
+
 func getAccountDataFromLogState(state *types.AccountWrapper) *types.StateAccount {
+
+	if isDeletedAccount(state) {
+		return nil
+	}
+
 	return &types.StateAccount{
 		Nonce:    state.Nonce,
 		Balance:  (*big.Int)(state.Balance),
@@ -392,7 +407,8 @@ func (w *zktrieProofWriter) traceAccountUpdate(addr common.Address, updateAccDat
 		if err := w.tracingZktrie.TryDelete(addr.Bytes32()); err != nil {
 			return nil, fmt.Errorf("delete zktrie account state fail: %s", err)
 		}
-		delete(w.tracingAccounts, addr)
+		w.tracingAccounts[addr] = nil
+
 	}
 
 	proof = proofList{}
@@ -411,6 +427,18 @@ func (w *zktrieProofWriter) traceAccountUpdate(addr common.Address, updateAccDat
 		//now accountKey must has been filled
 	}
 
+	// notice we have change that no leaf (account data) exist in either before or after,
+	// for that case we had to calculate the nodeKey here
+	if out.AccountKey == nil {
+		word := zkt.NewByte32FromBytesPaddingZero(addr.Bytes())
+		k, err := word.Hash()
+		if err != nil {
+			panic(fmt.Errorf("unexpected hash error for address: %s", err))
+		}
+		kHash := zkt.NewHashFromBigInt(k)
+		out.AccountKey = hexutil.Bytes(kHash[:])
+	}
+
 	return out, nil
 }
 
@@ -516,18 +544,24 @@ func (w *zktrieProofWriter) HandleNewState(accountState *types.AccountWrapper) (
 		return w.traceStorageUpdate(accountState.Address, storeAddr, storeValue)
 	} else {
 
+		var stateRoot common.Hash
 		accData := getAccountDataFromLogState(accountState)
 
 		out, err := w.traceAccountUpdate(accountState.Address, func(accBefore *types.StateAccount) *types.StateAccount {
 			if accBefore != nil {
-				accData.Root = accBefore.Root
+				stateRoot = accBefore.Root
+			}
+			// we need to restore stateRoot from before
+			if accData != nil {
+				accData.Root = stateRoot
 			}
 			return accData
 		})
 		if err != nil {
 			return nil, fmt.Errorf("update account state %s fail: %s", accountState.Address, err)
 		}
-		hash, err := zkt.NewHashFromBytes(accData.Root[:])
+
+		hash, err := zkt.NewHashFromBytes(stateRoot[:])
 		if err != nil {
 			return nil, fmt.Errorf("malform of state root in account %s", accountState.Address)
 		}
@@ -540,7 +574,6 @@ func (w *zktrieProofWriter) HandleNewState(accountState *types.AccountWrapper) (
 func handleLogs(od opOrderer, currentContract common.Address, logs []*types.StructLogRes) {
 	logStack := []int{0}
 	contractStack := map[int]common.Address{}
-	skipDepth := 0
 	callEnterAddress := currentContract
 
 	// now trace every OP which could cause changes on state:
@@ -552,6 +585,7 @@ func handleLogs(od opOrderer, currentContract common.Address, logs []*types.Stru
 			//update currentContract according to previous op
 			contractStack[sl] = currentContract
 			currentContract = callEnterAddress
+
 		} else if sl > sLog.Depth {
 			logStack = logStack[:sl-1]
 			currentContract = contractStack[sLog.Depth]
@@ -559,13 +593,17 @@ func handleLogs(od opOrderer, currentContract common.Address, logs []*types.Stru
 			calledLog := logs[resumePos]
 
 			//no need to handle fail calling
-			if !(calledLog.ExtraData != nil && calledLog.ExtraData.CallFailed) {
+			if calledLog.ExtraData != nil {
+				if !calledLog.ExtraData.CallFailed {
 					//reentry the last log which "cause" the calling, some handling may needed
 					switch calledLog.Op {
 					case "CREATE", "CREATE2":
 						//addr, accDataBefore := getAccountDataFromProof(calledLog, posCALLBefore)
 						od.absorb(getAccountState(calledLog, posCREATEAfter))
 					}
+				} else {
+					od.readonly(false)
+				}
 			}
 
 		} else {
@@ -577,73 +615,96 @@ func handleLogs(od opOrderer, currentContract common.Address, logs []*types.Stru
 		}
 		callEnterAddress = currentContract
 
-		if skipDepth != 0 {
+		//check extra status for current op if it is a call
-			if skipDepth < sLog.Depth {
+		if extraData := sLog.ExtraData; extraData != nil {
-				continue
+			if extraData.CallFailed || len(sLog.ExtraData.Caller) < 2 {
-			} else {
+				// no enough caller data (2) is being capture indicate we are in an immediate failure
-				skipDepth = 0
+				// i.e. it fail before stack entry (like no enough balance for a "call with value"),
+				// or we just not handle this calling op correctly yet
+
+				// for a failed option, now we just purpose nothing happens (FIXME: it is inconsentent with mpt_table)
+				// except for CREATE, for which the callee's nonce would be increased
+				switch sLog.Op {
+				case "CREATE", "CREATE2":
+					st := copyAccountState(extraData.Caller[0])
+					st.Nonce += 1
+					od.absorb(st)
 				}
 			}
 
-		if exD := sLog.ExtraData; exD != nil && exD.CallFailed {
+			if extraData.CallFailed {
-			//mark current op and next ops with more depth skippable
+				od.readonly(true)
-			skipDepth = sLog.Depth
+			}
-			continue
+			// now trace caller's status first
+			if caller := extraData.Caller; len(caller) >= 2 {
+				od.absorb(caller[1])
+			}
 		}
 
 		switch sLog.Op {
+		case "SELFDESTRUCT":
+			// NOTE: this op code has been disabled so we treat it as nothing now
+
+			//in SELFDESTRUCT, a call on target address is made so the balance would be updated
+			//in the last item
+			//stateTarget := getAccountState(sLog, posSELFDESTRUCT)
+			//od.absorb(stateTarget)
+			//then build an "deleted state", only address and other are default
+			//od.absorb(&types.AccountWrapper{Address: currentContract})
+
 		case "CREATE", "CREATE2":
+			// notice in immediate failure we have no enough tracing in extraData
+			if len(sLog.ExtraData.StateList) >= 2 {
 				state := getAccountState(sLog, posCREATE)
 				od.absorb(state)
 				//update contract to CREATE addr
-
 				callEnterAddress = state.Address
+			}
+
 		case "CALL", "CALLCODE":
+			// notice in immediate failure we have no enough tracing in extraData
+			if len(sLog.ExtraData.StateList) >= 3 {
 				state := getAccountState(sLog, posCALL)
 				od.absorb(state)
 				callEnterAddress = state.Address
+			}
 		case "STATICCALL":
-			//static call has no update on target address
+			//static call has no update on target address (and no immediate failure?)
 			callEnterAddress = getAccountState(sLog, posSTATICCALL).Address
+		case "DELEGATECALL":
+
 		case "SLOAD":
 			accountState := getAccountState(sLog, posSSTOREBefore)
-			od.absorb(accountState)
+			od.absorbStorage(accountState, nil)
 		case "SSTORE":
 			log.Debug("build SSTORE", "pc", sLog.Pc, "key", sLog.Stack[len(sLog.Stack)-1])
 			accountState := copyAccountState(getAccountState(sLog, posSSTOREBefore))
 			// notice the log only provide the value BEFORE store and it is not suitable for our protocol,
 			// here we change it into value AFTER update
+			before := accountState.Storage
 			accountState.Storage = &types.StorageWrapper{
 				Key:   sLog.Stack[len(sLog.Stack)-1],
 				Value: sLog.Stack[len(sLog.Stack)-2],
 			}
-			od.absorb(accountState)
+			od.absorbStorage(accountState, before)
 
 		default:
 		}
 	}
 }
 
-func handleTx(od opOrderer, txResult *types.ExecutionResult) {
+func HandleTx(od opOrderer, txResult *types.ExecutionResult) {
+
+	// the from state is read before tx is handled and nonce is added, we combine both
+	preTxSt := copyAccountState(txResult.From)
+	preTxSt.Nonce += 1
+	od.absorb(preTxSt)
 
-	// handle failed tx
 	if txResult.Failed {
-		handled := false
+		od.readonly(true)
-		for _, state := range txResult.AccountsAfter {
-			if state.Address != txResult.From.Address {
-				continue
-			}
-			od.absorb(state)
-			handled = true
-		}
-		if !handled {
-			panic(fmt.Errorf("no caller account in postTx status"))
-		}
-		return
 	}
 
 	var toAddr common.Address
-
 	if state := txResult.AccountCreated; state != nil {
 		od.absorb(state)
 		toAddr = state.Address
@@ -652,8 +713,18 @@ func handleTx(od opOrderer, txResult *types.ExecutionResult) {
 	}
 
 	handleLogs(od, toAddr, txResult.StructLogs)
+	if txResult.Failed {
+		od.readonly(false)
+	}
 
 	for _, state := range txResult.AccountsAfter {
+		// special case: for suicide, the state has been captured in SELFDESTRUCT
+		// and we skip it here
+		if isDeletedAccount(state) {
+			log.Debug("skip suicide address", "address", state.Address)
+			continue
+		}
+
 		od.absorb(state)
 	}
 
@@ -663,6 +734,8 @@ const defaultOrdererScheme = MPTWitnessRWTbl
 
 var usedOrdererScheme = defaultOrdererScheme
 
+func SetOrderScheme(t MPTWitnessType) { usedOrdererScheme = t }
+
 // HandleBlockResult only for backward compatibility
 func HandleBlockResult(block *types.BlockResult) ([]*StorageTrace, error) {
 	return HandleBlockResultEx(block, usedOrdererScheme)
@@ -682,13 +755,13 @@ func HandleBlockResultEx(block *types.BlockResult, ordererScheme MPTWitnessType)
 	case MPTWitnessNatural:
 		od = &simpleOrderer{}
 	case MPTWitnessRWTbl:
-		od = newRWTblOrderer(writer.tracingAccounts)
+		od = NewRWTblOrderer(writer.tracingAccounts)
 	default:
 		return nil, fmt.Errorf("unrecognized scheme %d", ordererScheme)
 	}
 
 	for _, tx := range block.ExecutionResults {
-		handleTx(od, tx)
+		HandleTx(od, tx)
 	}
 
 	// notice some coinbase addr (like all zero) is in fact not exist and should not be update