Apply zero traces to the SMT

smt/pkg/smt/entity_storage.go

@@ -252,7 +252,6 @@ func (s *SMT) SetStorage(ctx context.Context, logPrefix string, accChanges map[l
 			if err = s.DeleteKeySource(&keyBalance); err != nil {
 				return nil, nil, err
 			}
-
 		}
 
 		keysBatchStorage = append(keysBatchStorage, &keyNonce)

smt/pkg/smt/smt.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/ledgerwatch/erigon/smt/pkg/db"
 	"github.com/ledgerwatch/erigon/smt/pkg/utils"
+	"github.com/ledgerwatch/erigon/turbo/trie"
 	"github.com/ledgerwatch/log/v3"
 )
 
@@ -328,7 +329,11 @@ func (s *SMT) insert(k utils.NodeKey, v utils.NodeValue8, newValH [4]uint64, old
 
 				oldKey := utils.RemoveKeyBits(*foundKey, level2+1)
 				oldLeafHash, err := s.hashcalcAndSave(utils.ConcatArrays4(oldKey, foundOldValHash), utils.LeafCapacity)
-				s.Db.InsertHashKey(oldLeafHash, *foundKey)
+				if err != nil {
+					return nil, err
+				}
+
+				err = s.Db.InsertHashKey(oldLeafHash, *foundKey)
 				if err != nil {
 					return nil, err
 				}
@@ -350,7 +355,10 @@ func (s *SMT) insert(k utils.NodeKey, v utils.NodeValue8, newValH [4]uint64, old
 					return nil, err
 				}
 
-				s.Db.InsertHashKey(newLeafHash, k)
+				err = s.Db.InsertHashKey(newLeafHash, k)
+				if err != nil {
+					return nil, err
+				}
 
 				var node [8]uint64
 				for i := 0; i < 8; i++ {
@@ -416,7 +424,10 @@ func (s *SMT) insert(k utils.NodeKey, v utils.NodeValue8, newValH [4]uint64, old
 				return nil, err
 			}
 
-			s.Db.InsertHashKey(newLeafHash, k)
+			err = s.Db.InsertHashKey(newLeafHash, k)
+			if err != nil {
+				return nil, err
+			}
 
 			proofHashCounter += 2
 
@@ -471,10 +482,15 @@ func (s *SMT) insert(k utils.NodeKey, v utils.NodeValue8, newValH [4]uint64, old
 
 					oldKey := utils.RemoveKeyBits(*insKey, level+1)
 					oldLeafHash, err := s.hashcalcAndSave(utils.ConcatArrays4(oldKey, *valH), utils.LeafCapacity)
-					s.Db.InsertHashKey(oldLeafHash, *insKey)
 					if err != nil {
 						return nil, err
 					}
+
+					err = s.Db.InsertHashKey(oldLeafHash, *insKey)
+					if err != nil {
+						return nil, err
+					}
+
 					proofHashCounter += 1
 
 					if level >= 0 {
@@ -785,36 +801,39 @@ func (s *SMT) insertHashNode(path []int, hash [4]uint64, root utils.NodeKey) (ut
 	}
 
 	childIndex := path[0]
-
 	childOldRoot := rootVal[childIndex*4 : childIndex*4+4]
 
 	childNewRoot, err := s.insertHashNode(path[1:], hash, utils.NodeKeyFromBigIntArray(childOldRoot))
-
 	if err != nil {
 		return utils.NodeKey{}, err
 	}
 
 	var newIn [8]uint64
 
 	emptyRootVal := utils.NodeValue12{}
+	sibling := utils.BranchCapacity
 
-	if childIndex == 0 {
-		var sibling [4]uint64
-		if rootVal == emptyRootVal {
-			sibling = [4]uint64{0, 0, 0, 0}
-		} else {
+	if childIndex == int(utils.Left) {
+		if rootVal != emptyRootVal {
 			sibling = *rootVal.Get4to8()
 		}
 		newIn = utils.ConcatArrays4(childNewRoot, sibling)
 	} else {
-		var sibling [4]uint64
-		if rootVal == emptyRootVal {
-			sibling = [4]uint64{0, 0, 0, 0}
-		} else {
+		if rootVal != emptyRootVal {
 			sibling = *rootVal.Get0to4()
 		}
 		newIn = utils.ConcatArrays4(sibling, childNewRoot)
 	}
 
 	return s.hashcalcAndSave(newIn, utils.BranchCapacity)
 }
+
+// Copy copies the SMT, by converting it into the witness first and then back to the SMT from witness.
+func (s *SMT) Copy(ctx context.Context, rd trie.RetainDecider) (*SMT, error) {
+	witness, err := BuildWitness(s, rd, ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return BuildSMTFromWitness(witness)
+}

smt/pkg/smt/smt_state_reader.go

@@ -43,9 +43,9 @@ func (s *SMT) ReadAccountData(address libcommon.Address) (*accounts.Account, err
 	return account, nil
 }
 
-// ReadAccountStorage reads account storage from the SMT (not implemented for SMT)
+// ReadAccountStorage reads account storage from the SMT
 func (s *SMT) ReadAccountStorage(address libcommon.Address, incarnation uint64, key *libcommon.Hash) ([]byte, error) {
-	value, err := s.getValue(0, address, key)
+	value, err := s.getValue(utils.SC_STORAGE, address, key)
 	if err != nil {
 		return []byte{}, err
 	}
@@ -70,7 +70,7 @@ func (s *SMT) ReadAccountCodeSize(address libcommon.Address, _ uint64, _ libcomm
 		return 0, err
 	}
 
-	sizeBig := big.NewInt(0).SetBytes(valueInBytes)
+	sizeBig := new(big.Int).SetBytes(valueInBytes)
 
 	if !sizeBig.IsInt64() {
 		err = errors.New("code size value is too large to fit into an int")

smt/pkg/smt/smt_state_writer.go

@@ -0,0 +1,122 @@
+package smt
+
+import (
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"math/big"
+
+	"github.com/holiman/uint256"
+	"github.com/ledgerwatch/erigon-lib/common"
+	libcommon "github.com/ledgerwatch/erigon-lib/common"
+	"github.com/ledgerwatch/erigon/core/state"
+	"github.com/ledgerwatch/erigon/core/types"
+	"github.com/ledgerwatch/erigon/core/types/accounts"
+	"github.com/ledgerwatch/erigon/smt/pkg/utils"
+	"github.com/ledgerwatch/erigon/turbo/trie"
+)
+
+var _ state.StateWriter = (*SMT)(nil)
+
+func (s *SMT) UpdateAccountData(address common.Address, _ *accounts.Account, account *accounts.Account) error {
+	return s.SetAccountStorage(address, account)
+}
+
+func (s *SMT) UpdateAccountCode(address common.Address, _ uint64, _ common.Hash, code []byte) error {
+	if len(code) == 0 {
+		return nil
+	}
+
+	return s.SetContractBytecode(address.Hex(), fmt.Sprintf("0x%s", hex.EncodeToString(code)))
+}
+
+func (s *SMT) DeleteAccount(address common.Address, original *accounts.Account) error {
+	return errors.New("DeleteAccount is not implemented for the SMT")
+}
+
+func (s *SMT) WriteAccountStorage(address common.Address, _ uint64, key *common.Hash, _ *uint256.Int, value *uint256.Int) error {
+	if key == nil || value == nil {
+		return nil
+	}
+
+	storageKeyString := fmt.Sprintf("0x%s", hex.EncodeToString(key.Bytes()))
+	storageMap := map[string]string{storageKeyString: value.Hex()}
+	_, err := s.SetContractStorage(address.Hex(), storageMap, nil)
+
+	return err
+}
+
+func (s *SMT) CreateContract(address common.Address) error {
+	return s.SetAccountStorage(address, nil)
+}
+
+// ApplyTraces applies a map of traces on the given SMT and returns new instance of SMT, without altering the original one.
+func (s *SMT) ApplyTraces(traces map[libcommon.Address]*types.TxnTrace, rd trie.RetainDecider) (*SMT, error) {
+	// result, err := s.Copy(context.Background(), rd)
+	// if err != nil {
+	// 	return nil, err
+	// }
+
+	result := s
+
+	for addr, trace := range traces {
+		if trace.SelfDestructed != nil && *trace.SelfDestructed {
+			nodeVal, err := s.getValue(utils.KEY_NONCE, addr, nil)
+			if err != nil {
+				return nil, err
+			}
+
+			if len(nodeVal) > 0 {
+				return nil, fmt.Errorf("account %s is annotated to be self-destructed, but it already exists in the SMT", addr)
+			}
+
+			continue
+		}
+
+		addrString := addr.Hex()
+
+		// Set account balance and nonce
+		balanceBig := big.NewInt(0)
+		if trace.Balance != nil {
+			balanceBig = trace.Balance.ToBig()
+		}
+
+		nonceBig := big.NewInt(0)
+		if trace.Nonce != nil {
+			nonceBig = trace.Nonce.ToBig()
+		}
+
+		if _, err := result.SetAccountState(addrString, balanceBig, nonceBig); err != nil {
+			return nil, err
+		}
+
+		// Set account nonce
+		if trace.Nonce != nil {
+			if _, err := result.SetAccountNonce(addrString, trace.Nonce.ToBig()); err != nil {
+				return nil, err
+			}
+		}
+
+		// Set account storage map
+		storageMap := make(map[string]string)
+		for hash, storageSlot := range trace.StorageWritten {
+			storageKey := fmt.Sprintf("0x%s", hex.EncodeToString(hash[:]))
+			storageMap[storageKey] = storageSlot.Hex()
+		}
+
+		if len(storageMap) > 0 {
+			if _, err := result.SetContractStorage(addrString, storageMap, nil); err != nil {
+				return nil, err
+			}
+		}
+
+		// Set account bytecode
+		if trace.CodeUsage != nil {
+			if err := result.SetContractBytecode(addrString, hex.EncodeToString(trace.CodeUsage.Write)); err != nil {
+				return nil, err
+			}
+		}
+	}
+
+	return result, nil
+}

smt/pkg/smt/smt_state_writer_test.go

@@ -0,0 +1,85 @@
+package smt
+
+import (
+	"bytes"
+	"encoding/hex"
+	"encoding/json"
+	"os"
+	"strings"
+	"testing"
+
+	libcommon "github.com/ledgerwatch/erigon-lib/common"
+	"github.com/ledgerwatch/erigon-lib/kv/memdb"
+	"github.com/ledgerwatch/erigon/core/state"
+	"github.com/ledgerwatch/erigon/core/types"
+	"github.com/ledgerwatch/erigon/turbo/trie"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSMTApplyTraces(t *testing.T) {
+	type testData struct {
+		IsFullWitness bool                                  `json:"isFullWitness"`
+		Witness       string                                `json:"witness"`
+		StateRoot     libcommon.Hash                        `json:"stateRoot"`
+		Traces        map[libcommon.Address]*types.TxnTrace `json:"traces,omitempty"`
+	}
+
+	cases := []struct {
+		name string
+		file string
+	}{
+		{
+			name: "Single EOA tx full witness",
+			file: "./testdata/zerotraces/full-witness-single-eoa-tx.json",
+		},
+		{
+			name: "Empty block full witness",
+			file: "./testdata/zerotraces/full-witness-empty-block.json",
+		},
+		// {
+		// 	name: "SC deployment full witness",
+		// 	file: "./testdata/zerotraces/full-witness-contract-deployment.json",
+		// },
+		// {
+		// 	name: "SC interaction full witness",
+		// 	file: "./testdata/zerotraces/full-witness-contract-interaction.json",
+		// },
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			rawTestData, err := os.ReadFile(c.file)
+			require.NoError(t, err)
+
+			var td testData
+			err = json.Unmarshal(rawTestData, &td)
+			require.NoError(t, err)
+
+			decodedWitnessRaw, err := hex.DecodeString(strings.TrimPrefix(td.Witness, "0x"))
+			require.NoError(t, err)
+
+			witness, err := trie.NewWitnessFromReader(bytes.NewReader(decodedWitnessRaw), false)
+			require.NoError(t, err)
+
+			smt, err := BuildSMTFromWitness(witness)
+			require.NoError(t, err)
+
+			var rd trie.RetainDecider
+			if td.IsFullWitness {
+				rd = &trie.AlwaysTrueRetainDecider{}
+			} else {
+				_, tx := memdb.NewTestTx(t)
+				tds := state.NewTrieDbState(td.StateRoot, tx, 0, state.NewPlainStateReader(tx))
+				tds.StartNewBuffer()
+
+				rd, err = tds.ResolveSMTRetainList(nil)
+				require.NoError(t, err)
+			}
+
+			newSMT, err := smt.ApplyTraces(td.Traces, rd)
+			require.NoError(t, err)
+
+			require.Equal(t, td.StateRoot, libcommon.BigToHash(newSMT.LastRoot()))
+		})
+	}
+}