feat: Initialize the new app with a deep SMT instead of populating the empty SMT

baseapp/baseapp.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	"github.com/gogo/protobuf/proto"
+	smtlib "github.com/lazyledger/smt"
 	abci "github.com/tendermint/tendermint/abci/types"
 	"github.com/tendermint/tendermint/crypto/tmhash"
 	"github.com/tendermint/tendermint/libs/log"
@@ -774,8 +775,7 @@ func makeABCIData(msgResponses []*codectypes.Any) ([]byte, error) {
 // It returns the tracing-enabled app along with the trace buffers used
 func (app *BaseApp) enableFraudProofGenerationMode(storeKeys []types.StoreKey, routerOpts map[string]AppOptionFunc) (*BaseApp, map[string]*bytes.Buffer, error) {
 	cms := app.cms.(*multi.Store)
-	lastVersion := cms.LastCommitID().Version
-	previousCMS, err := cms.GetVersion(lastVersion)
+	previousCMS, err := cms.GetLastStore()
 	if err != nil {
 		return nil, nil, err
 	}
@@ -787,11 +787,13 @@ func (app *BaseApp) enableFraudProofGenerationMode(storeKeys []types.StoreKey, r
 
 	// Initialize params from previousCMS
 	storeToLoadFrom := make(map[string]types.KVStore)
+	storeKeyToSMT := make(map[string]*smtlib.SparseMerkleTree)
 	storeKeyNames := make([]string, 0, len(storeKeys))
 	for _, storeKey := range storeKeys {
 		storeKeyName := storeKey.Name()
 		storeKeyNames = append(storeKeyNames, storeKeyName)
 		storeToLoadFrom[storeKeyName] = previousCMS.GetKVStore(storeKey)
+		storeKeyToSMT[storeKeyName] = previousCMS.GetSubstoreSMT(storeKeyName).GetTree()
 		storeKeyToSubstoreTraceBuf[storeKeyName] = &bytes.Buffer{}
 	}
 
@@ -808,7 +810,7 @@ func (app *BaseApp) enableFraudProofGenerationMode(storeKeys []types.StoreKey, r
 			options = append(options, AppOptionFunc(routerOpt))
 		}
 	}
-	newApp, err := SetupBaseAppFromParams(app.name+"WithTracing", app.logger, dbm.NewMemDB(), app.txDecoder, storeKeyNames, app.LastBlockHeight(), storeToLoadFrom, options...)
+	newApp, err := SetupBaseAppFromParams(app.name+"WithTracing", app.logger, dbm.NewMemDB(), app.txDecoder, storeKeyNames, storeKeyToSMT, app.LastBlockHeight(), storeToLoadFrom, options...)
 
 	// Need to reset all the buffers to remove anything logged while setting up baseapp
 	storeTraceBuf.Reset()
@@ -870,20 +872,27 @@ func (app *BaseApp) generateFraudProof(storeKeyToSubstoreTraceBuf map[string]*by
 }
 
 // set up a new baseapp from given params
-func SetupBaseAppFromParams(appName string, logger log.Logger, db dbm.Connection, txDecoder sdk.TxDecoder, storeKeyNames []string, blockHeight int64, storeToLoadFrom map[string]types.KVStore, options ...AppOption) (*BaseApp, error) {
+func SetupBaseAppFromParams(appName string, logger log.Logger, db dbm.Connection, txDecoder sdk.TxDecoder, storeKeyNames []string, storeKeyToSMT map[string]*smtlib.SparseMerkleTree, blockHeight int64, storeToLoadFrom map[string]types.KVStore, options ...AppOption) (*BaseApp, error) {
 	storeKeys := make([]types.StoreKey, 0, len(storeKeyNames))
-
+	storeKeyToSubstoreHash := make(map[string][]byte)
 	for _, storeKeyName := range storeKeyNames {
 		storeKey := sdk.NewKVStoreKey(storeKeyName)
 		storeKeys = append(storeKeys, storeKey)
 		subStore := storeToLoadFrom[storeKeyName]
 		it := subStore.Iterator(nil, nil)
+		substoreSMT := storeKeyToSMT[storeKeyName]
 		for ; it.Valid(); it.Next() {
 			key, val := it.Key(), it.Value()
-			options = append(options, SetSubstoreKVPair(storeKey, key, val))
+			proof, err := substoreSMT.Prove(key)
+			if err != nil {
+				return nil, err
+			}
+			options = append(options, SetDeepSMTBranchKVPair(storeKey, substoreSMT.Root(), proof, key, val))
 		}
+		storeKeyToSubstoreHash[storeKeyName] = substoreSMT.Root()
 	}
-	options = append(options, SetSubstores(storeKeys...))
+
+	options = append(options, SetSubstoresWithRoots(storeKeyToSubstoreHash, storeKeys...))
 
 	// This initial height is used in `BeginBlock` in `validateHeight`
 	options = append(options, SetInitialHeight(blockHeight))
@@ -898,5 +907,9 @@ func SetupBaseAppFromParams(appName string, logger log.Logger, db dbm.Connection
 
 // set up a new baseapp from a fraudproof
 func SetupBaseAppFromFraudProof(appName string, logger log.Logger, db dbm.Connection, txDecoder sdk.TxDecoder, fraudProof FraudProof, options ...AppOption) (*BaseApp, error) {
-	return SetupBaseAppFromParams(appName, logger, db, txDecoder, fraudProof.getModules(), fraudProof.blockHeight, fraudProof.extractStore(), options...)
+	storeKeyToSMT, err := fraudProof.getSubstoreSMTs()
+	if err != nil {
+		return nil, err
+	}
+	return SetupBaseAppFromParams(appName, logger, db, txDecoder, fraudProof.getModules(), storeKeyToSMT, fraudProof.blockHeight, fraudProof.extractStore(), options...)
 }

baseapp/baseapp_test.go

@@ -2328,6 +2328,8 @@ func TestEndToEndFraudProof(t *testing.T) {
 	appHashB2, err := appB2.cms.(*multi.Store).GetAppHash()
 	require.Nil(t, err)
 
+	storeHashB2 := cmsB2.GetSubstoreSMT(capKey2.Name()).Root()
+
 	fraudProof, err := appB2.generateFraudProof(storeKeyToSubstoreTraceBuf, appB2.LastBlockHeight())
 	require.Nil(t, err)
 
@@ -2337,6 +2339,16 @@ func TestEndToEndFraudProof(t *testing.T) {
 	require.Nil(t, err)
 	require.True(t, fraudProofVerified)
 
+	// Now we take contents of the fraud proof and create a new baseapp with its contents
+	codec := codec.NewLegacyAmino()
+	registerTestCodec(codec)
+	appB3, err := SetupBaseAppFromFraudProof(t.Name(), defaultLogger(), dbm.NewMemDB(), testTxDecoder(codec), fraudProof, AppOptionFunc(routerOpt))
+	require.Nil(t, err)
+
+	// Check if the the resulting app's substore is the same as the one fraudproof initialized it with
+	storeHashB3 := appB3.cms.(*multi.Store).GetSubstoreSMT(capKey2.Name()).Root()
+	require.Equal(t, storeHashB2, storeHashB3)
+
 	// Check if fraudproof only contains keys from that fraudulent block or not, if yes, pass
 	for _, moduleName := range fraudProof.getModules() {
 		witnessKeys, witnessVals := make([]string, 0), make([]string, 0)
@@ -2426,7 +2438,7 @@ func TestFraudProofGenerationMode(t *testing.T) {
 	}
 
 	txs1 := executeBlockWithArbitraryTxs(t, appB2, numTransactions, 1)
-  
+
 	for _, storeKey := range cmsB2.GetStoreKeys() {
 		subStoreBuf := storeKeyToSubstoreTraceBuf[storeKey.Name()]
 		tracedKeys := appB2.cms.GetKVStore(storeKey).(*tracekv.Store).GetAllKeysUsedInTrace(*subStoreBuf).Values()
@@ -2516,8 +2528,7 @@ func TestGenerateAndLoadFraudProof(t *testing.T) {
 	executeBlockWithRequests(t, appB1, nil, []*abci.RequestDeliverTx{fraudDeliverRequest}, nil, 0)
 	// Now, subStoreBuf knows about all the keys accessed in (S2 -> S3)
 
-	//make new appFraudGen app which creates app with previous state
-
+	// Make new appFraudGen app which creates app with previous state
 	storeKeys := []types.StoreKey{capKey1, capKey2}
 	routerOpts := make(map[string]AppOptionFunc)
 	newRouterOpt := func(bapp *BaseApp) {
@@ -2552,14 +2563,12 @@ func TestGenerateAndLoadFraudProof(t *testing.T) {
 	require.Nil(t, err)
 	require.True(t, fraudProofVerified)
 
-	// Now we take contents of the fraud proof which was recorded with S1 and try to populate a fresh baseapp B2 with it
-	// B2 <- S1
+	// Now we take contents of the fraud proof which was recorded with S2 and try to populate a fresh baseapp B2 with it
+	// B2 <- S2
 	codec := codec.NewLegacyAmino()
 	registerTestCodec(codec)
 	appB2, err := SetupBaseAppFromFraudProof(t.Name(), defaultLogger(), dbm.NewMemDB(), testTxDecoder(codec), fraudProof, AppOptionFunc(routerOpt))
 	require.Nil(t, err)
 	storeHashB2 := appB2.cms.(*multi.Store).GetSubstoreSMT(capKey2.Name()).Root()
-	_, _ = storeHashB1, storeHashB2
-	// TODO: need to initialize the new app with deep subtree instead of populating the empty SMT on startup
-	// require.Equal(t, storeHashB1, storeHashB2)
+	require.Equal(t, storeHashB1, storeHashB2)
 }

baseapp/fraudproof.go

@@ -2,12 +2,13 @@ package baseapp
 
 import (
 	"bytes"
+	"crypto/sha256"
 	"fmt"
 
 	"github.com/cosmos/cosmos-sdk/store/mem"
 	"github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/cosmos/cosmos-sdk/store/v2alpha1/smt"
-
+	smtlib "github.com/lazyledger/smt"
 	tmcrypto "github.com/tendermint/tendermint/proto/tendermint/crypto"
 )
 
@@ -45,6 +46,25 @@ func (fraudProof *FraudProof) getModules() []string {
 	return keys
 }
 
+func (fraudProof *FraudProof) getSubstoreSMTs() (map[string]*smtlib.SparseMerkleTree, error) {
+	storeKeyToSMT := make(map[string]*smtlib.SparseMerkleTree)
+	for storeKey, stateWitness := range fraudProof.stateWitness {
+		rootHash := stateWitness.rootHash
+		substoreDeepSMT := smtlib.NewDeepSparseMerkleSubTree(smtlib.NewSimpleMap(), smtlib.NewSimpleMap(), sha256.New(), rootHash)
+		for _, witnessData := range stateWitness.WitnessData {
+			proofOp, key, val := witnessData.proof, witnessData.Key, witnessData.Value
+			proof, err := smt.ProofDecoder(proofOp)
+			if err != nil {
+				return nil, err
+			}
+			smtProof := proof.(*smt.ProofOp).GetProof()
+			substoreDeepSMT.AddBranch(smtProof, key, val)
+		}
+		storeKeyToSMT[storeKey] = substoreDeepSMT.SparseMerkleTree
+	}
+	return storeKeyToSMT, nil
+}
+
 func (fraudProof *FraudProof) extractStore() map[string]types.KVStore {
 	store := make(map[string]types.KVStore)
 	for storeKey, stateWitness := range fraudProof.stateWitness {

baseapp/options.go

@@ -11,6 +11,7 @@ import (
 	storetypes "github.com/cosmos/cosmos-sdk/store/v2alpha1"
 	"github.com/cosmos/cosmos-sdk/store/v2alpha1/multi"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/lazyledger/smt"
 )
 
 // File for storing in-package BaseApp optional functions,
@@ -43,7 +44,20 @@ func SetTracerFor(skey string, w io.Writer) StoreOption {
 // Only works for v2alpha1/multi
 func SetSubstoreKVPair(skey storetypes.StoreKey, key, val []byte) AppOptionOrdered {
 	return AppOptionOrdered{
-		func(bapp *BaseApp) { bapp.cms.(*multi.Store).SetSubstoreKVPair(skey, key, val) },
+		func(bapp *BaseApp) {
+			bapp.cms.(*multi.Store).SetSubstoreKVPair(skey, key, val)
+		},
+		OptionOrderAfterStore,
+	}
+}
+
+// SetDeepSMTBranchKVPair adds a branch with a key, value pair for the given substore's deep SMT inside a multistore
+// Only works for v2alpha1/multi
+func SetDeepSMTBranchKVPair(skey storetypes.StoreKey, root []byte, proof smt.SparseMerkleProof, key, val []byte) AppOptionOrdered {
+	return AppOptionOrdered{
+		func(bapp *BaseApp) {
+			bapp.cms.(*multi.Store).SetDeepSMTBranchKVPair(skey.Name(), root, proof, key, val)
+		},
 		OptionOrderAfterStore,
 	}
 }
@@ -116,6 +130,24 @@ func SetSubstores(keys ...storetypes.StoreKey) StoreOption {
 	}
 }
 
+// SetSubstoresWithRoots registers substores according to app configuration
+// It also assigns the given substoreHashes for each key using deep SMTs
+func SetSubstoresWithRoots(storeKeyToSubstoreHash map[string][]byte, keys ...storetypes.StoreKey) StoreOption {
+	return func(config *multi.StoreParams, _ uint64) error {
+		for _, key := range keys {
+			typ, err := storetypes.StoreKeyToType(key)
+			if err != nil {
+				return err
+			}
+			if err = config.RegisterSubstore(key, typ); err != nil {
+				return err
+			}
+			config.SetSubStoreHashMap(storeKeyToSubstoreHash)
+		}
+		return nil
+	}
+}
+
 func SetSubstoresFromMaps(
 	keys map[string]*storetypes.KVStoreKey,
 	tkeys map[string]*storetypes.TransientStoreKey,

store/v2alpha1/multi/params.go

@@ -17,6 +17,7 @@ func DefaultStoreParams() StoreParams {
 		storeKeys:                 storeKeys{},
 		traceListenMixin:          newTraceListenMixin(),
 		substoreTraceListenMixins: make(map[string]*traceListenMixin),
+		storeKeyToSubstoreHash:    make(map[string][]byte),
 	}
 }
 
@@ -50,6 +51,10 @@ func (par *StoreParams) SetTracerFor(skey string, w io.Writer) {
 	par.substoreTraceListenMixins[skey] = tlm
 }
 
+func (par *StoreParams) SetSubStoreHashMap(storeKeyToSubstoreHash map[string][]byte) {
+	par.storeKeyToSubstoreHash = storeKeyToSubstoreHash
+}
+
 func (par *StoreParams) storeKey(key string) (types.StoreKey, error) {
 	skey, ok := par.storeKeys[key]
 	if !ok {

store/v2alpha1/multi/store.go

@@ -27,6 +27,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/store/v2alpha1/transient"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/cosmos/cosmos-sdk/types/kv"
+	smtlib "github.com/lazyledger/smt"
 	tmcrypto "github.com/tendermint/tendermint/proto/tendermint/crypto"
 )
 
@@ -76,6 +77,7 @@ type StoreParams struct {
 	// Contains The trace context and listeners that can also be set from store methods.
 	*traceListenMixin
 	substoreTraceListenMixins map[string]*traceListenMixin
+	storeKeyToSubstoreHash    map[string][]byte
 }
 
 // StoreSchema defineds a mapping of substore keys to store types
@@ -353,8 +355,8 @@ func NewStore(db dbm.Connection, opts StoreParams) (ret *Store, err error) {
 			return
 		}
 		ret.schema[skey] = typ
+		ret.createSubstoreAsDeepSMT(skey.Name(), opts.storeKeyToSubstoreHash[skey.Name()])
 	}
-
 	return
 }
 
@@ -494,6 +496,36 @@ func (rs *Store) HasKVStore(skey types.StoreKey) bool {
 	return has
 }
 
+func (s *Store) SetDeepSMTBranchKVPair(skey string, root []byte, proof smtlib.SparseMerkleProof, key []byte, value []byte) {
+	pfx := prefixSubstore(skey)
+	stateCommitmentRW := prefixdb.NewReadWriter(s.stateCommitmentTxn, pfx)
+	smtdb := prefixdb.NewReadWriter(stateCommitmentRW, smtPrefix)
+	smt.AddBranchToStoreWithDeepSMT(smtdb, root, proof, key, value)
+}
+
+// Creates a persistent substore. This reads, but does not update the substore cache.
+// Uses the given subStoreHash to initialize the substore with an underlying deep SMT
+func (s *Store) createSubstoreAsDeepSMT(skey string, substoreHash []byte) (*substore, error) {
+	pfx := prefixSubstore(skey)
+	stateRW := prefixdb.NewReadWriter(s.stateTxn, pfx)
+	stateCommitmentRW := prefixdb.NewReadWriter(s.stateCommitmentTxn, pfx)
+	var stateCommitmentStore *smt.Store
+
+	err := stateRW.Set(substoreMerkleRootKey, substoreHash)
+	if err != nil {
+		return nil, err
+	}
+	smtdb := prefixdb.NewReadWriter(stateCommitmentRW, smtPrefix)
+	stateCommitmentStore = smt.NewStoreWithDeepSMT(smtdb, substoreHash)
+
+	return &substore{
+		root:                 s,
+		name:                 skey,
+		dataBucket:           prefixdb.NewReadWriter(stateRW, dataPrefix),
+		stateCommitmentStore: stateCommitmentStore,
+	}, nil
+}
+
 // Gets a persistent substore. This reads, but does not update the substore cache.
 // Use it in cases where we need to access a store internally (e.g. read/write Merkle keys, queries)
 func (s *Store) getSubstore(key string) (*substore, error) {
@@ -1072,6 +1104,19 @@ func (s *Store) SetPruning(po pruningtypes.PruningOptions) {
 	s.pruningManager.SetOptions(po)
 }
 
+func (s *Store) GetLastStore() (*viewStore, error) {
+	versions, err := s.stateDB.Versions()
+	if err != nil {
+		return nil, err
+	}
+	lastVersion := int64(versions.Last())
+	lastStore, err := s.GetVersion(lastVersion)
+	if err != nil {
+		return nil, err
+	}
+	return lastStore.(*viewStore), nil
+}
+
 func (s *Store) GetSubstoreSMT(key string) *smt.Store {
 	sub, err := s.getSubstore(key)
 	if err != nil {

store/v2alpha1/multi/view_store.go

@@ -79,7 +79,7 @@ func (s *viewStore) GetSubstoreSMT(key string) *smt.Store {
 	return sub.stateCommitmentStore
 }
 
-func (s *viewStore) GetSubStoreProof(key string) (*tmcrypto.ProofOp, []byte, error) {
+func (s *viewStore) GetSubstoreProof(key string) (*tmcrypto.ProofOp, []byte, error) {
 	storeHashes, err := s.getMerkleRoots()
 	if err != nil {
 		return nil, nil, err

store/v2alpha1/smt/proof.go

@@ -62,6 +62,10 @@ func (p *ProofOp) GetKey() []byte {
 	return p.Key
 }
 
+func (p *ProofOp) GetProof() smt.SparseMerkleProof {
+	return p.Proof
+}
+
 func (p *ProofOp) ProofOp() tmmerkle.ProofOp {
 	var data bytes.Buffer
 	enc := gob.NewEncoder(&data)