feat: state migration from IAVL to SMT (ADR-040)

store/rootmulti/snapshot_test.go

@@ -210,7 +210,8 @@ func TestMultistoreSnapshotRestore(t *testing.T) {
 	require.Equal(t, *dummyExtensionItem.GetExtension(), *nextItem.GetExtension())
 
 	assert.Equal(t, source.LastCommitID(), target.LastCommitID())
-	for key, sourceStore := range source.GetStores() {
+	for _, key := range source.GetStoreNames() {
+		sourceStore := source.GetStoreByName(key.Name()).(types.CommitKVStore)
 		targetStore := target.GetStoreByName(key.Name()).(types.CommitKVStore)
 		switch sourceStore.GetStoreType() {
 		case types.StoreTypeTransient:
@@ -234,7 +235,7 @@ func benchmarkMultistoreSnapshot(b *testing.B, stores uint8, storeKeys uint64) {
 
 	for i := 0; i < b.N; i++ {
 		target := rootmulti.NewStore(dbm.NewMemDB())
-		for key := range source.GetStores() {
+		for _, key := range source.GetStoreNames() {
 			target.MountStoreWithDB(key, types.StoreTypeIAVL, nil)
 		}
 		err := target.LoadLatestVersion()
@@ -269,7 +270,7 @@ func benchmarkMultistoreSnapshotRestore(b *testing.B, stores uint8, storeKeys ui
 
 	for i := 0; i < b.N; i++ {
 		target := rootmulti.NewStore(dbm.NewMemDB())
-		for key := range source.GetStores() {
+		for _, key := range source.GetStoreNames() {
 			target.MountStoreWithDB(key, types.StoreTypeIAVL, nil)
 		}
 		err := target.LoadLatestVersion()

store/rootmulti/store.go

@@ -148,9 +148,9 @@ func (rs *Store) GetCommitKVStore(key types.StoreKey) types.CommitKVStore {
 	return rs.stores[key]
 }
 
-// GetStores returns mounted stores
-func (rs *Store) GetStores() map[types.StoreKey]types.CommitKVStore {
-	return rs.stores
+// GetStoreNames returns the storeNames and StoreKeys
+func (rs *Store) GetStoreNames() map[string]types.StoreKey {
+	return rs.keysByName
 }
 
 // LoadLatestVersionAndUpgrade implements CommitMultiStore

store/v2/multi/migration.go

@@ -0,0 +1,68 @@
+package multi
+
+import (
+	dbm "github.com/cosmos/cosmos-sdk/db"
+	"github.com/cosmos/cosmos-sdk/store/iavl"
+	"github.com/cosmos/cosmos-sdk/store/mem"
+	v1Store "github.com/cosmos/cosmos-sdk/store/rootmulti"
+	"github.com/cosmos/cosmos-sdk/store/transient"
+	"github.com/cosmos/cosmos-sdk/store/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+)
+
+// MigrateV2 will migrate the state from iavl to smt
+func MigrateV2(rootMultiStore *v1Store.Store, store2db dbm.DBConnection, storeConfig StoreConfig) (*Store, error) {
+	type namedStore struct {
+		*iavl.Store
+		name string
+	}
+	var stores []namedStore
+	for _, storeKey := range rootMultiStore.GetStoreNames() {
+		keyName := storeKey.Name()
+		switch store := rootMultiStore.GetStoreByName(keyName).(type) {
+		case *iavl.Store:
+			err := storeConfig.RegisterSubstore(keyName, types.StoreTypePersistent)
+			if err != nil {
+				return nil, err
+			}
+			stores = append(stores, namedStore{name: keyName, Store: store})
+		case *transient.Store, *mem.Store:
+			continue
+		default:
+			return nil, sdkerrors.Wrapf(sdkerrors.ErrLogic, "don't know how to migrate store %q of type %T", keyName, store)
+		}
+	}
+
+	// creating the new store of smt tree
+	rootStore, err := NewStore(store2db, storeConfig)
+	if err != nil {
+		return nil, err
+	}
+
+	// if version is 0 there is no state data to commit
+	if rootMultiStore.LastCommitID().Version == 0 {
+		return rootStore, nil
+	}
+
+	// iterate through the rootmulti stores and save the key/values into smt tree
+	for _, store := range stores {
+		subStore, err := rootStore.getSubstore(store.name)
+		if err != nil {
+			return nil, err
+		}
+		// iterate all iavl tree node key/values
+		iterator := store.Iterator(nil, nil)
+		for ; iterator.Valid(); iterator.Next() {
+			// set the iavl key,values into smt node
+			subStore.Set(iterator.Key(), iterator.Value())
+		}
+	}
+
+	// commit the all key/values from iavl to smt tree (SMT Store)
+	_, err = rootStore.commit(uint64(rootMultiStore.LastCommitID().Version))
+	if err != nil {
+		return nil, err
+	}
+
+	return rootStore, nil
+}

store/v2/multi/migration_test.go

@@ -0,0 +1,104 @@
+package multi
+
+import (
+	"encoding/binary"
+	"fmt"
+	"github.com/cosmos/cosmos-sdk/db/memdb"
+	"github.com/cosmos/cosmos-sdk/store/iavl"
+	"github.com/cosmos/cosmos-sdk/store/rootmulti"
+	"github.com/cosmos/cosmos-sdk/store/types"
+	"github.com/stretchr/testify/require"
+	dbm "github.com/tendermint/tm-db"
+	"math/rand"
+	"testing"
+)
+
+func TestMigrationV2(t *testing.T) {
+	r := rand.New(rand.NewSource(49872768940))
+
+	// setup a rootmulti store
+	db := dbm.NewMemDB()
+	v1Store := rootmulti.NewStore(db)
+
+	// mount the kvStores
+	var keys []*types.KVStoreKey
+	for i := uint8(0); i < 10; i++ {
+		key := types.NewKVStoreKey(fmt.Sprintf("store%v", i))
+		v1Store.MountStoreWithDB(key, types.StoreTypeIAVL, nil)
+		keys = append(keys, key)
+	}
+
+	err := v1Store.LoadLatestVersion()
+	require.Nil(t, err)
+
+	// setup a random test data
+	for _, key := range keys {
+		store := v1Store.GetStore(key).(*iavl.Store)
+		store.Set([]byte("temp_data"), []byte("one"))
+
+		for i := 0; i < len(keys); i++ {
+			k := make([]byte, 8)
+			v := make([]byte, 1024)
+			binary.BigEndian.PutUint64(k, uint64(i))
+			_, err := r.Read(v)
+			if err != nil {
+				panic(err)
+			}
+			store.Set(k, v)
+		}
+	}
+
+	testCases := []struct {
+		testName   string
+		emptyStore bool
+	}{
+		{
+			"Migration With Empty Store",
+			true,
+		},
+		{
+			"Migration From Root Multi Store (IAVL) to SMT ",
+			false,
+		},
+	}
+
+	for _, testCase := range testCases {
+		if !testCase.emptyStore {
+			v1Store.Commit()
+		}
+
+		// setup a new root store of smt
+		db2 := memdb.NewDB()
+		storeConfig := DefaultStoreConfig()
+		// migrating the iavl store (v1) to smt store (v2)
+		v2Store, err := MigrateV2(v1Store, db2, storeConfig)
+		require.NoError(t, err)
+
+		for _, key := range keys {
+			v2StoreKVStore := v2Store.GetKVStore(key)
+			if testCase.emptyStore {
+				// check the empty store
+				require.Nil(t, v2StoreKVStore.Get([]byte("temp_data")))
+			} else {
+				require.Equal(t, v2StoreKVStore.Get([]byte("temp_data")), []byte("one"))
+			}
+			require.Equal(t, v2Store.LastCommitID().Version, v1Store.LastCommitID().Version)
+		}
+		err = v2Store.Close()
+		require.NoError(t, err)
+	}
+}
+
+func TestMigrateV2ForEmptyStore(t *testing.T) {
+	// setup a rootmulti store
+	db := dbm.NewMemDB()
+	v1Store := rootmulti.NewStore(db)
+	err := v1Store.LoadLatestVersion()
+	require.Nil(t, err)
+	db2 := memdb.NewDB()
+	storeConfig := DefaultStoreConfig()
+	// migrating the iavl store (v1) to smt store (v2)
+	v2Store, err := MigrateV2(v1Store, db2, storeConfig)
+	require.NoError(t, err)
+	require.Equal(t, v2Store.LastCommitID(), v1Store.LastCommitID())
+}