From e900fb1d7370fb7532b5b182845216f6d6036b55 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Mon, 1 Mar 2021 19:34:07 -0800
Subject: [PATCH 01/26] -added consensus version tracking to x/upgrade

---
 simapp/app.go                       |  6 +++-
 types/module/module.go              | 12 ++++++++
 x/upgrade/abci.go                   |  4 +++
 x/upgrade/abci_test.go              |  6 ++--
 x/upgrade/keeper/grpc_query_test.go |  3 +-
 x/upgrade/keeper/keeper.go          | 48 ++++++++++++++++++++++++++++-
 x/upgrade/keeper/keeper_test.go     |  3 +-
 x/upgrade/module.go                 |  3 +-
 x/upgrade/module_test.go            | 26 ++++++++++++++++
 x/upgrade/types/handler.go          |  3 +-
 x/upgrade/types/keys.go             | 13 ++++++++
 11 files changed, 118 insertions(+), 9 deletions(-)
 create mode 100644 x/upgrade/module_test.go

diff --git a/simapp/app.go b/simapp/app.go
index 4086585e684c..e43018e09807 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -374,6 +374,10 @@ func NewSimApp(
 		authz.NewAppModule(appCodec, app.AuthzKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
 		transferModule,
 	)
+	// give upgrade keeper the module manager
+	app.UpgradeKeeper.SetModuleManager(app.mm)
+	// pass the updated keeper to the module manager
+	app.mm.Modules[upgradetypes.ModuleName] = upgrade.NewAppModule(app.UpgradeKeeper)
 
 	// During begin block slashing happens after distr.BeginBlocker so that
 	// there is nothing left over in the validator fee pool, so as to keep the
@@ -391,7 +395,7 @@ func NewSimApp(
 	// so that other modules that want to create or claim capabilities afterwards in InitChain
 	// can do so safely.
 	app.mm.SetOrderInitGenesis(
-		capabilitytypes.ModuleName, authtypes.ModuleName, banktypes.ModuleName, distrtypes.ModuleName, stakingtypes.ModuleName,
+		upgradetypes.ModuleName, capabilitytypes.ModuleName, authtypes.ModuleName, banktypes.ModuleName, distrtypes.ModuleName, stakingtypes.ModuleName,
 		slashingtypes.ModuleName, govtypes.ModuleName, minttypes.ModuleName, crisistypes.ModuleName,
 		ibchost.ModuleName, genutiltypes.ModuleName, evidencetypes.ModuleName, authztypes.ModuleName, ibctransfertypes.ModuleName,
 		feegranttypes.ModuleName,
diff --git a/types/module/module.go b/types/module/module.go
index 9123c9af7ec9..f6f597013b46 100644
--- a/types/module/module.go
+++ b/types/module/module.go
@@ -395,3 +395,15 @@ func (m *Manager) EndBlock(ctx sdk.Context, req abci.RequestEndBlock) abci.Respo
 		Events:           ctx.EventManager().ABCIEvents(),
 	}
 }
+
+// GetConsensusVersions gets consensus version from all modules
+func (m *Manager) GetConsensusVersions() MigrationMap {
+	migmap := make(map[string]uint64)
+	for _, v := range m.Modules {
+		version := v.ConsensusVersion()
+		name := v.Name()
+		migmap[name] = version
+	}
+
+	return migmap
+}
diff --git a/x/upgrade/abci.go b/x/upgrade/abci.go
index d346decb023c..e5a6ab2eb28b 100644
--- a/x/upgrade/abci.go
+++ b/x/upgrade/abci.go
@@ -73,3 +73,7 @@ func BeginBlocker(k keeper.Keeper, ctx sdk.Context, _ abci.RequestBeginBlock) {
 func BuildUpgradeNeededMsg(plan types.Plan) string {
 	return fmt.Sprintf("UPGRADE \"%s\" NEEDED at %s: %s", plan.Name, plan.DueAt(), plan.Info)
 }
+
+func InitChainer(k keeper.Keeper, ctx sdk.Context) {
+	k.SetConsensusVersions(ctx)
+}
diff --git a/x/upgrade/abci_test.go b/x/upgrade/abci_test.go
index 9eb0fc1a7ab9..e7ba99c0e409 100644
--- a/x/upgrade/abci_test.go
+++ b/x/upgrade/abci_test.go
@@ -128,7 +128,7 @@ func VerifyDoUpgrade(t *testing.T) {
 	})
 
 	t.Log("Verify that the upgrade can be successfully applied with a handler")
-	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan) {})
+	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan, migrationMap module.MigrationMap) error { return nil })
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -144,7 +144,7 @@ func VerifyDoUpgradeWithCtx(t *testing.T, newCtx sdk.Context, proposalName strin
 	})
 
 	t.Log("Verify that the upgrade can be successfully applied with a handler")
-	s.keeper.SetUpgradeHandler(proposalName, func(ctx sdk.Context, plan types.Plan) {})
+	s.keeper.SetUpgradeHandler(proposalName, func(ctx sdk.Context, plan types.Plan, migrationMap module.MigrationMap) error { return nil })
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -156,7 +156,7 @@ func TestHaltIfTooNew(t *testing.T) {
 	s := setupTest(10, map[int64]bool{})
 	t.Log("Verify that we don't panic with registered plan not in database at all")
 	var called int
-	s.keeper.SetUpgradeHandler("future", func(ctx sdk.Context, plan types.Plan) { called++ })
+	s.keeper.SetUpgradeHandler("future", func(ctx sdk.Context, plan types.Plan, migrationMap module.MigrationMap) error { called++; return nil })
 
 	newCtx := s.ctx.WithBlockHeight(s.ctx.BlockHeight() + 1).WithBlockTime(time.Now())
 	req := abci.RequestBeginBlock{Header: newCtx.BlockHeader()}
diff --git a/x/upgrade/keeper/grpc_query_test.go b/x/upgrade/keeper/grpc_query_test.go
index d307157402de..a71f015d8721 100644
--- a/x/upgrade/keeper/grpc_query_test.go
+++ b/x/upgrade/keeper/grpc_query_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/simapp"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
 
@@ -110,7 +111,7 @@ func (suite *UpgradeTestSuite) TestAppliedCurrentPlan() {
 				suite.app.UpgradeKeeper.ScheduleUpgrade(suite.ctx, plan)
 
 				suite.ctx = suite.ctx.WithBlockHeight(expHeight)
-				suite.app.UpgradeKeeper.SetUpgradeHandler(planName, func(ctx sdk.Context, plan types.Plan) {})
+				suite.app.UpgradeKeeper.SetUpgradeHandler(planName, func(_ sdk.Context, _ types.Plan, _ module.MigrationMap) error { return nil })
 				suite.app.UpgradeKeeper.ApplyUpgrade(suite.ctx, plan)
 
 				req = &types.QueryAppliedPlanRequest{Name: planName}
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index e366ae09198e..ff910521f558 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -16,6 +16,7 @@ import (
 	store "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
 
@@ -28,6 +29,7 @@ type Keeper struct {
 	storeKey           sdk.StoreKey
 	cdc                codec.BinaryMarshaler
 	upgradeHandlers    map[string]types.UpgradeHandler
+	ModuleManager      *module.Manager
 }
 
 // NewKeeper constructs an upgrade Keeper
@@ -48,6 +50,50 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
+func (k *Keeper) SetModuleManager(mm *module.Manager) *Keeper {
+	k.ModuleManager = mm
+	return k
+}
+
+// SetConsensusVersions saves the consensus versions retrieved from module.Manager
+func (k Keeper) SetConsensusVersions(ctx sdk.Context) {
+	modules := k.ModuleManager.GetConsensusVersions()
+	for modName, ver := range modules {
+		k.setConsensusVersion(ctx, ver, modName)
+	}
+}
+
+// SetConsensusVersion sets the given module's consensus version to the given version
+func (k Keeper) setConsensusVersion(ctx sdk.Context, version uint64, moduleName string) {
+	store := k.getConsensusVersionPrefixStore(ctx, moduleName)
+	versionBytes := make([]byte, 8)
+	binary.LittleEndian.PutUint64(versionBytes, version)
+	store.Set(types.MigrationMapKeyModule(moduleName), versionBytes)
+}
+
+// getConsensusVersionStore gets the store of the given module name
+func (k Keeper) getConsensusVersionPrefixStore(ctx sdk.Context, moduleName string) prefix.Store {
+	store := ctx.KVStore(k.storeKey)
+	return prefix.NewStore(store, types.MigrationMapKeyModule(moduleName))
+}
+
+// GetConsensusVersions gets a MigrationMap from state
+func (k Keeper) GetConsensusVersions(ctx sdk.Context) module.MigrationMap {
+	store := ctx.KVStore(k.storeKey)
+	it := sdk.KVStorePrefixIterator(store, []byte{types.MigrationMapByte})
+
+	migmap := make(map[string]uint64)
+
+	defer it.Close()
+	for ; it.Valid(); it.Next() {
+		moduleName := string(it.Key())
+		moduleVersion := binary.LittleEndian.Uint64(it.Value())
+		migmap[moduleName] = moduleVersion
+	}
+
+	return migmap
+}
+
 // ScheduleUpgrade schedules an upgrade based on the specified plan.
 // If there is another Plan already scheduled, it will overwrite it
 // (implicitly cancelling the current plan)
@@ -191,7 +237,7 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic("ApplyUpgrade should never be called without first checking HasHandler")
 	}
 
-	handler(ctx, plan)
+	handler(ctx, plan, k.GetConsensusVersions(ctx))
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
 	// This will prevent resubmission of upgrade msg after upgrade is already completed.
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index 48f0301d3ff1..3809eb35229d 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/simapp"
 	store "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/keeper"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
@@ -135,7 +136,7 @@ func (s *KeeperTestSuite) TestScheduleUpgrade() {
 				Height: 123450000,
 			},
 			setup: func() {
-				s.app.UpgradeKeeper.SetUpgradeHandler("all-good", func(_ sdk.Context, _ types.Plan) {})
+				s.app.UpgradeKeeper.SetUpgradeHandler("all-good", func(_ sdk.Context, _ types.Plan, _ module.MigrationMap) error { return nil })
 				s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, types.Plan{
 					Name:   "all-good",
 					Info:   "some text here",
diff --git a/x/upgrade/module.go b/x/upgrade/module.go
index 23311c22d5f3..be5baec3d496 100644
--- a/x/upgrade/module.go
+++ b/x/upgrade/module.go
@@ -101,7 +101,8 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 }
 
 // InitGenesis is ignored, no sense in serializing future upgrades
-func (am AppModule) InitGenesis(_ sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
+func (am AppModule) InitGenesis(ctx sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
+	InitChainer(am.keeper, ctx)
 	return []abci.ValidatorUpdate{}
 }
 
diff --git a/x/upgrade/module_test.go b/x/upgrade/module_test.go
new file mode 100644
index 000000000000..292fb0688b52
--- /dev/null
+++ b/x/upgrade/module_test.go
@@ -0,0 +1,26 @@
+package upgrade_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	abcitypes "github.com/tendermint/tendermint/abci/types"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
+
+	"github.com/cosmos/cosmos-sdk/simapp"
+)
+
+func TestItSavesConsensusVersionToState(t *testing.T) {
+	app := simapp.Setup(false)
+	ctx := app.BaseApp.NewContext(false, tmproto.Header{})
+
+	app.InitChain(
+		abcitypes.RequestInitChain{
+			AppStateBytes: []byte("{}"),
+			ChainId:       "test-chain-id",
+		},
+	)
+
+	migmap := app.UpgradeKeeper.GetConsensusVersions(ctx)
+	require.NotEmpty(t, migmap)
+}
diff --git a/x/upgrade/types/handler.go b/x/upgrade/types/handler.go
index 44e50cff112f..cc1cd78a495e 100644
--- a/x/upgrade/types/handler.go
+++ b/x/upgrade/types/handler.go
@@ -2,7 +2,8 @@ package types
 
 import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
 )
 
 // UpgradeHandler specifies the type of function that is called when an upgrade is applied
-type UpgradeHandler func(ctx sdk.Context, plan Plan)
+type UpgradeHandler func(ctx sdk.Context, plan Plan, migrationMap module.MigrationMap) error
diff --git a/x/upgrade/types/keys.go b/x/upgrade/types/keys.go
index 410f63597c6e..fdffe1ef891a 100644
--- a/x/upgrade/types/keys.go
+++ b/x/upgrade/types/keys.go
@@ -22,6 +22,9 @@ const (
 	// DoneByte is a prefix for to look up completed upgrade plan by name
 	DoneByte = 0x1
 
+	// MigrationMapByte is a prefix to look up module names (key) and versions (value)
+	MigrationMapByte = 0x2
+
 	// KeyUpgradedIBCState is the key under which upgraded ibc state is stored in the upgrade store
 	KeyUpgradedIBCState = "upgradedIBCState"
 
@@ -51,3 +54,13 @@ func UpgradedClientKey(height int64) []byte {
 func UpgradedConsStateKey(height int64) []byte {
 	return []byte(fmt.Sprintf("%s/%d/%s", KeyUpgradedIBCState, height, KeyUpgradedConsState))
 }
+
+func MigrationMapKey() []byte {
+	return []byte{MigrationMapByte}
+}
+
+func MigrationMapKeyModule(moduleName string) []byte {
+	nameBytes := []byte(moduleName)
+	key := append(MigrationMapKey(), nameBytes...)
+	return key
+}

From 036db4f5c1ae7c9709ff2373fa57cb1460d9f370 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Wed, 3 Mar 2021 10:31:45 -0800
Subject: [PATCH 02/26] -added interface to module manager -added e2e test for
 migrations using consensus version store in x/upgrade -cleaned up x/upgrade
 Keeper -handler in apply upgrade now handles errors and setting consensus
 versions -cleaned up migration map keys -removed init chainer method -simapp
 now implements GetConsensusVersions to assist with testing

---
 simapp/app.go                   |  7 ++++-
 types/module/module.go          |  6 ++++-
 x/upgrade/abci.go               |  4 ---
 x/upgrade/keeper/keeper.go      | 45 +++++++++++++++------------------
 x/upgrade/keeper/keeper_test.go | 31 ++++++++++++++++++++++-
 x/upgrade/module.go             |  1 -
 x/upgrade/module_test.go        |  6 ++---
 x/upgrade/types/keys.go         | 10 --------
 8 files changed, 63 insertions(+), 47 deletions(-)

diff --git a/simapp/app.go b/simapp/app.go
index e43018e09807..58b800fee261 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -375,7 +375,7 @@ func NewSimApp(
 		transferModule,
 	)
 	// give upgrade keeper the module manager
-	app.UpgradeKeeper.SetModuleManager(app.mm)
+	app.UpgradeKeeper.SetVersionManager(app.mm)
 	// pass the updated keeper to the module manager
 	app.mm.Modules[upgradetypes.ModuleName] = upgrade.NewAppModule(app.UpgradeKeeper)
 
@@ -631,6 +631,11 @@ func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.Mig
 	return app.mm.RunMigrations(ctx, app.configurator, migrateFromVersions)
 }
 
+// Returns a map (MigrationMap) of key module name and value module consensus version
+func (app *SimApp) GetConsensusVersions() module.MigrationMap {
+	return app.mm.GetConsensusVersions()
+}
+
 // RegisterSwaggerAPI registers swagger route with API Server
 func RegisterSwaggerAPI(ctx client.Context, rtr *mux.Router) {
 	statikFS, err := fs.New()
diff --git a/types/module/module.go b/types/module/module.go
index f6f597013b46..e821723b2677 100644
--- a/types/module/module.go
+++ b/types/module/module.go
@@ -337,6 +337,10 @@ type MigrationHandler func(sdk.Context) error
 // version from which we should perform the migration for each module.
 type MigrationMap map[string]uint64
 
+type VersionManager interface {
+	GetConsensusVersions() MigrationMap
+}
+
 // RunMigrations performs in-place store migrations for all modules.
 func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, migrateFromVersions MigrationMap) error {
 	c, ok := cfg.(configurator)
@@ -398,7 +402,7 @@ func (m *Manager) EndBlock(ctx sdk.Context, req abci.RequestEndBlock) abci.Respo
 
 // GetConsensusVersions gets consensus version from all modules
 func (m *Manager) GetConsensusVersions() MigrationMap {
-	migmap := make(map[string]uint64)
+	migmap := make(MigrationMap)
 	for _, v := range m.Modules {
 		version := v.ConsensusVersion()
 		name := v.Name()
diff --git a/x/upgrade/abci.go b/x/upgrade/abci.go
index e5a6ab2eb28b..d346decb023c 100644
--- a/x/upgrade/abci.go
+++ b/x/upgrade/abci.go
@@ -73,7 +73,3 @@ func BeginBlocker(k keeper.Keeper, ctx sdk.Context, _ abci.RequestBeginBlock) {
 func BuildUpgradeNeededMsg(plan types.Plan) string {
 	return fmt.Sprintf("UPGRADE \"%s\" NEEDED at %s: %s", plan.Name, plan.DueAt(), plan.Info)
 }
-
-func InitChainer(k keeper.Keeper, ctx sdk.Context) {
-	k.SetConsensusVersions(ctx)
-}
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index ff910521f558..8916b302f973 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -29,7 +29,7 @@ type Keeper struct {
 	storeKey           sdk.StoreKey
 	cdc                codec.BinaryMarshaler
 	upgradeHandlers    map[string]types.UpgradeHandler
-	ModuleManager      *module.Manager
+	versionManager     module.VersionManager
 }
 
 // NewKeeper constructs an upgrade Keeper
@@ -50,45 +50,35 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
-func (k *Keeper) SetModuleManager(mm *module.Manager) *Keeper {
-	k.ModuleManager = mm
-	return k
+func (k *Keeper) SetVersionManager(vm module.VersionManager) {
+	k.versionManager = vm
 }
 
 // SetConsensusVersions saves the consensus versions retrieved from module.Manager
 func (k Keeper) SetConsensusVersions(ctx sdk.Context) {
-	modules := k.ModuleManager.GetConsensusVersions()
+	modules := k.versionManager.GetConsensusVersions()
+	store := ctx.KVStore(k.storeKey)
+	migrationStore := prefix.NewStore(store, []byte{types.MigrationMapByte})
 	for modName, ver := range modules {
-		k.setConsensusVersion(ctx, ver, modName)
+		nameBytes := []byte(modName)
+		verBytes := make([]byte, 8)
+		binary.LittleEndian.PutUint64(verBytes, ver)
+		migrationStore.Set(nameBytes, verBytes)
 	}
 }
 
-// SetConsensusVersion sets the given module's consensus version to the given version
-func (k Keeper) setConsensusVersion(ctx sdk.Context, version uint64, moduleName string) {
-	store := k.getConsensusVersionPrefixStore(ctx, moduleName)
-	versionBytes := make([]byte, 8)
-	binary.LittleEndian.PutUint64(versionBytes, version)
-	store.Set(types.MigrationMapKeyModule(moduleName), versionBytes)
-}
-
-// getConsensusVersionStore gets the store of the given module name
-func (k Keeper) getConsensusVersionPrefixStore(ctx sdk.Context, moduleName string) prefix.Store {
-	store := ctx.KVStore(k.storeKey)
-	return prefix.NewStore(store, types.MigrationMapKeyModule(moduleName))
-}
-
 // GetConsensusVersions gets a MigrationMap from state
 func (k Keeper) GetConsensusVersions(ctx sdk.Context) module.MigrationMap {
 	store := ctx.KVStore(k.storeKey)
 	it := sdk.KVStorePrefixIterator(store, []byte{types.MigrationMapByte})
 
-	migmap := make(map[string]uint64)
-
+	migmap := make(module.MigrationMap)
 	defer it.Close()
 	for ; it.Valid(); it.Next() {
-		moduleName := string(it.Key())
+		moduleBytes := it.Key()
+		name := string(moduleBytes[1:])
 		moduleVersion := binary.LittleEndian.Uint64(it.Value())
-		migmap[moduleName] = moduleVersion
+		migmap[name] = moduleVersion
 	}
 
 	return migmap
@@ -237,7 +227,12 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic("ApplyUpgrade should never be called without first checking HasHandler")
 	}
 
-	handler(ctx, plan, k.GetConsensusVersions(ctx))
+	err := handler(ctx, plan, k.GetConsensusVersions(ctx))
+	if err != nil {
+		panic(err)
+	}
+
+	k.SetConsensusVersions(ctx)
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
 	// This will prevent resubmission of upgrade msg after upgrade is already completed.
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index 3809eb35229d..bb7a080d8343 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -12,6 +12,7 @@ import (
 	store "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
+	"github.com/cosmos/cosmos-sdk/x/bank"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/keeper"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
@@ -153,7 +154,7 @@ func (s *KeeperTestSuite) TestScheduleUpgrade() {
 		s.Run(tc.name, func() {
 			// reset suite
 			s.SetupTest()
-
+			s.app.UpgradeKeeper.SetVersionManager(s.app)
 			// setup test case
 			tc.setup()
 
@@ -212,6 +213,34 @@ func (s *KeeperTestSuite) TestSetUpgradedClient() {
 
 }
 
+// Mock version manager for TestMigrations
+type MockVersionManager struct{}
+
+func (m MockVersionManager) GetConsensusVersions() module.MigrationMap {
+	migmap := make(module.MigrationMap)
+	migmap["bank"] = 1
+	return migmap
+}
+
+// Tests that the underlying state of x/upgrade is set correctly after
+// an upgrade.
+func (s *KeeperTestSuite) TestMigrations() {
+	mockVM := MockVersionManager{}
+	s.app.UpgradeKeeper.SetVersionManager(mockVM)
+	s.app.UpgradeKeeper.SetConsensusVersions(s.ctx)
+	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.MigrationMap) error { return nil })
+	dummyPlan := types.Plan{
+		Name: "dummy",
+		Info: "some text here",
+		Time: s.ctx.BlockTime().Add(time.Hour),
+	}
+
+	s.app.UpgradeKeeper.SetVersionManager(s.app)
+	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, dummyPlan)
+	migmap := s.app.UpgradeKeeper.GetConsensusVersions(s.ctx)
+	s.Require().Equal(bank.AppModule{}.ConsensusVersion(), migmap["bank"])
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	suite.Run(t, new(KeeperTestSuite))
 }
diff --git a/x/upgrade/module.go b/x/upgrade/module.go
index be5baec3d496..154f006784a7 100644
--- a/x/upgrade/module.go
+++ b/x/upgrade/module.go
@@ -102,7 +102,6 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 
 // InitGenesis is ignored, no sense in serializing future upgrades
 func (am AppModule) InitGenesis(ctx sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
-	InitChainer(am.keeper, ctx)
 	return []abci.ValidatorUpdate{}
 }
 
diff --git a/x/upgrade/module_test.go b/x/upgrade/module_test.go
index 292fb0688b52..491bacc6da95 100644
--- a/x/upgrade/module_test.go
+++ b/x/upgrade/module_test.go
@@ -5,14 +5,12 @@ import (
 
 	"github.com/stretchr/testify/require"
 	abcitypes "github.com/tendermint/tendermint/abci/types"
-	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 
 	"github.com/cosmos/cosmos-sdk/simapp"
 )
 
-func TestItSavesConsensusVersionToState(t *testing.T) {
+func TestItGetsModuleVersions(t *testing.T) {
 	app := simapp.Setup(false)
-	ctx := app.BaseApp.NewContext(false, tmproto.Header{})
 
 	app.InitChain(
 		abcitypes.RequestInitChain{
@@ -21,6 +19,6 @@ func TestItSavesConsensusVersionToState(t *testing.T) {
 		},
 	)
 
-	migmap := app.UpgradeKeeper.GetConsensusVersions(ctx)
+	migmap := app.GetConsensusVersions()
 	require.NotEmpty(t, migmap)
 }
diff --git a/x/upgrade/types/keys.go b/x/upgrade/types/keys.go
index fdffe1ef891a..593c8aeaa892 100644
--- a/x/upgrade/types/keys.go
+++ b/x/upgrade/types/keys.go
@@ -54,13 +54,3 @@ func UpgradedClientKey(height int64) []byte {
 func UpgradedConsStateKey(height int64) []byte {
 	return []byte(fmt.Sprintf("%s/%d/%s", KeyUpgradedIBCState, height, KeyUpgradedConsState))
 }
-
-func MigrationMapKey() []byte {
-	return []byte{MigrationMapByte}
-}
-
-func MigrationMapKeyModule(moduleName string) []byte {
-	nameBytes := []byte(moduleName)
-	key := append(MigrationMapKey(), nameBytes...)
-	return key
-}

From 5c1ae59f465f6f361039d45d01969ab9bd95aff4 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Thu, 4 Mar 2021 09:09:10 -0800
Subject: [PATCH 03/26] Changed MigrationMap identifier to VersionMap removed
 module_test

---
 simapp/app.go                       |  8 ++++----
 simapp/app_test.go                  |  2 +-
 types/module/module.go              | 16 ++++++++--------
 x/upgrade/abci_test.go              |  6 +++---
 x/upgrade/keeper/grpc_query_test.go |  2 +-
 x/upgrade/keeper/keeper.go          | 19 +++++++++++--------
 x/upgrade/keeper/keeper_test.go     | 20 ++++++++++----------
 x/upgrade/module.go                 |  2 +-
 x/upgrade/module_test.go            | 24 ------------------------
 x/upgrade/types/handler.go          |  2 +-
 x/upgrade/types/keys.go             |  4 ++--
 11 files changed, 42 insertions(+), 63 deletions(-)
 delete mode 100644 x/upgrade/module_test.go

diff --git a/simapp/app.go b/simapp/app.go
index 58b800fee261..9e9ee95baf42 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -619,7 +619,7 @@ func (app *SimApp) RegisterTendermintService(clientCtx client.Context) {
 // Example:
 //   cfg := module.NewConfigurator(...)
 //   app.UpgradeKeeper.SetUpgradeHandler("store-migration", func(ctx sdk.Context, plan upgradetypes.Plan) {
-//       err := app.RunMigrations(ctx, module.MigrationMap{
+//       err := app.RunMigrations(ctx, module.VersionMap{
 //           "bank": 1,     // Migrate x/bank from v1 to current x/bank's ConsensusVersion
 //           "staking": 8,  // Migrate x/staking from v8 to current x/staking's ConsensusVersion
 //      })
@@ -627,12 +627,12 @@ func (app *SimApp) RegisterTendermintService(clientCtx client.Context) {
 //           panic(err)
 //      }
 //   })
-func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.MigrationMap) error {
+func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.VersionMap) error {
 	return app.mm.RunMigrations(ctx, app.configurator, migrateFromVersions)
 }
 
-// Returns a map (MigrationMap) of key module name and value module consensus version
-func (app *SimApp) GetConsensusVersions() module.MigrationMap {
+// Returns a map (VersionMap) of key module name and value module consensus version
+func (app *SimApp) GetConsensusVersions() module.VersionMap {
 	return app.mm.GetConsensusVersions()
 }
 
diff --git a/simapp/app_test.go b/simapp/app_test.go
index 8f596f5ecc0d..95cdce2b0fcc 100644
--- a/simapp/app_test.go
+++ b/simapp/app_test.go
@@ -161,7 +161,7 @@ func TestRunMigrations(t *testing.T) {
 			// their latest ConsensusVersion.
 			err = app.RunMigrations(
 				app.NewContext(true, tmproto.Header{Height: app.LastBlockHeight()}),
-				module.MigrationMap{
+				module.VersionMap{
 					"bank":         1,
 					"auth":         auth.AppModule{}.ConsensusVersion(),
 					"authz":        authz.AppModule{}.ConsensusVersion(),
diff --git a/types/module/module.go b/types/module/module.go
index e821723b2677..00e754fdb503 100644
--- a/types/module/module.go
+++ b/types/module/module.go
@@ -333,16 +333,16 @@ func (m *Manager) ExportGenesis(ctx sdk.Context, cdc codec.JSONMarshaler) map[st
 // MigrationHandler is the migration function that each module registers.
 type MigrationHandler func(sdk.Context) error
 
-// MigrationMap is a map of moduleName -> version, where version denotes the
+// VersionMap is a map of moduleName -> version, where version denotes the
 // version from which we should perform the migration for each module.
-type MigrationMap map[string]uint64
+type VersionMap map[string]uint64
 
 type VersionManager interface {
-	GetConsensusVersions() MigrationMap
+	GetConsensusVersions() VersionMap
 }
 
 // RunMigrations performs in-place store migrations for all modules.
-func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, migrateFromVersions MigrationMap) error {
+func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, migrateFromVersions VersionMap) error {
 	c, ok := cfg.(configurator)
 	if !ok {
 		return sdkerrors.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", configurator{}, cfg)
@@ -401,13 +401,13 @@ func (m *Manager) EndBlock(ctx sdk.Context, req abci.RequestEndBlock) abci.Respo
 }
 
 // GetConsensusVersions gets consensus version from all modules
-func (m *Manager) GetConsensusVersions() MigrationMap {
-	migmap := make(MigrationMap)
+func (m *Manager) GetConsensusVersions() VersionMap {
+	vermap := make(VersionMap)
 	for _, v := range m.Modules {
 		version := v.ConsensusVersion()
 		name := v.Name()
-		migmap[name] = version
+		vermap[name] = version
 	}
 
-	return migmap
+	return vermap
 }
diff --git a/x/upgrade/abci_test.go b/x/upgrade/abci_test.go
index e7ba99c0e409..514480dfde12 100644
--- a/x/upgrade/abci_test.go
+++ b/x/upgrade/abci_test.go
@@ -128,7 +128,7 @@ func VerifyDoUpgrade(t *testing.T) {
 	})
 
 	t.Log("Verify that the upgrade can be successfully applied with a handler")
-	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan, migrationMap module.MigrationMap) error { return nil })
+	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan, versionMap module.VersionMap) error { return nil })
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -144,7 +144,7 @@ func VerifyDoUpgradeWithCtx(t *testing.T, newCtx sdk.Context, proposalName strin
 	})
 
 	t.Log("Verify that the upgrade can be successfully applied with a handler")
-	s.keeper.SetUpgradeHandler(proposalName, func(ctx sdk.Context, plan types.Plan, migrationMap module.MigrationMap) error { return nil })
+	s.keeper.SetUpgradeHandler(proposalName, func(ctx sdk.Context, plan types.Plan, versionMap module.VersionMap) error { return nil })
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -156,7 +156,7 @@ func TestHaltIfTooNew(t *testing.T) {
 	s := setupTest(10, map[int64]bool{})
 	t.Log("Verify that we don't panic with registered plan not in database at all")
 	var called int
-	s.keeper.SetUpgradeHandler("future", func(ctx sdk.Context, plan types.Plan, migrationMap module.MigrationMap) error { called++; return nil })
+	s.keeper.SetUpgradeHandler("future", func(ctx sdk.Context, plan types.Plan, versionMap module.VersionMap) error { called++; return nil })
 
 	newCtx := s.ctx.WithBlockHeight(s.ctx.BlockHeight() + 1).WithBlockTime(time.Now())
 	req := abci.RequestBeginBlock{Header: newCtx.BlockHeader()}
diff --git a/x/upgrade/keeper/grpc_query_test.go b/x/upgrade/keeper/grpc_query_test.go
index a71f015d8721..a817d6222346 100644
--- a/x/upgrade/keeper/grpc_query_test.go
+++ b/x/upgrade/keeper/grpc_query_test.go
@@ -111,7 +111,7 @@ func (suite *UpgradeTestSuite) TestAppliedCurrentPlan() {
 				suite.app.UpgradeKeeper.ScheduleUpgrade(suite.ctx, plan)
 
 				suite.ctx = suite.ctx.WithBlockHeight(expHeight)
-				suite.app.UpgradeKeeper.SetUpgradeHandler(planName, func(_ sdk.Context, _ types.Plan, _ module.MigrationMap) error { return nil })
+				suite.app.UpgradeKeeper.SetUpgradeHandler(planName, func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
 				suite.app.UpgradeKeeper.ApplyUpgrade(suite.ctx, plan)
 
 				req = &types.QueryAppliedPlanRequest{Name: planName}
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index 8916b302f973..fe322f8e213c 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -56,32 +56,35 @@ func (k *Keeper) SetVersionManager(vm module.VersionManager) {
 
 // SetConsensusVersions saves the consensus versions retrieved from module.Manager
 func (k Keeper) SetConsensusVersions(ctx sdk.Context) {
+	if k.versionManager == nil {
+		panic("Upgrade Keeper VersionManager was nil")
+	}
 	modules := k.versionManager.GetConsensusVersions()
 	store := ctx.KVStore(k.storeKey)
-	migrationStore := prefix.NewStore(store, []byte{types.MigrationMapByte})
+	versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
 	for modName, ver := range modules {
 		nameBytes := []byte(modName)
 		verBytes := make([]byte, 8)
 		binary.LittleEndian.PutUint64(verBytes, ver)
-		migrationStore.Set(nameBytes, verBytes)
+		versionStore.Set(nameBytes, verBytes)
 	}
 }
 
-// GetConsensusVersions gets a MigrationMap from state
-func (k Keeper) GetConsensusVersions(ctx sdk.Context) module.MigrationMap {
+// GetConsensusVersions gets a VersionMap from state
+func (k Keeper) GetConsensusVersions(ctx sdk.Context) module.VersionMap {
 	store := ctx.KVStore(k.storeKey)
-	it := sdk.KVStorePrefixIterator(store, []byte{types.MigrationMapByte})
+	it := sdk.KVStorePrefixIterator(store, []byte{types.VersionMapByte})
 
-	migmap := make(module.MigrationMap)
+	vermap := make(module.VersionMap)
 	defer it.Close()
 	for ; it.Valid(); it.Next() {
 		moduleBytes := it.Key()
 		name := string(moduleBytes[1:])
 		moduleVersion := binary.LittleEndian.Uint64(it.Value())
-		migmap[name] = moduleVersion
+		vermap[name] = moduleVersion
 	}
 
-	return migmap
+	return vermap
 }
 
 // ScheduleUpgrade schedules an upgrade based on the specified plan.
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index bb7a080d8343..04d0c7fa1658 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -137,7 +137,7 @@ func (s *KeeperTestSuite) TestScheduleUpgrade() {
 				Height: 123450000,
 			},
 			setup: func() {
-				s.app.UpgradeKeeper.SetUpgradeHandler("all-good", func(_ sdk.Context, _ types.Plan, _ module.MigrationMap) error { return nil })
+				s.app.UpgradeKeeper.SetUpgradeHandler("all-good", func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
 				s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, types.Plan{
 					Name:   "all-good",
 					Info:   "some text here",
@@ -214,21 +214,21 @@ func (s *KeeperTestSuite) TestSetUpgradedClient() {
 }
 
 // Mock version manager for TestMigrations
-type MockVersionManager struct{}
+type mockVersionManager struct{}
 
-func (m MockVersionManager) GetConsensusVersions() module.MigrationMap {
-	migmap := make(module.MigrationMap)
-	migmap["bank"] = 1
-	return migmap
+func (m mockVersionManager) GetConsensusVersions() module.VersionMap {
+	vermap := make(module.VersionMap)
+	vermap["bank"] = 1
+	return vermap
 }
 
 // Tests that the underlying state of x/upgrade is set correctly after
 // an upgrade.
 func (s *KeeperTestSuite) TestMigrations() {
-	mockVM := MockVersionManager{}
+	mockVM := mockVersionManager{}
 	s.app.UpgradeKeeper.SetVersionManager(mockVM)
 	s.app.UpgradeKeeper.SetConsensusVersions(s.ctx)
-	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.MigrationMap) error { return nil })
+	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
 	dummyPlan := types.Plan{
 		Name: "dummy",
 		Info: "some text here",
@@ -237,8 +237,8 @@ func (s *KeeperTestSuite) TestMigrations() {
 
 	s.app.UpgradeKeeper.SetVersionManager(s.app)
 	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, dummyPlan)
-	migmap := s.app.UpgradeKeeper.GetConsensusVersions(s.ctx)
-	s.Require().Equal(bank.AppModule{}.ConsensusVersion(), migmap["bank"])
+	vermap := s.app.UpgradeKeeper.GetConsensusVersions(s.ctx)
+	s.Require().Equal(bank.AppModule{}.ConsensusVersion(), vermap["bank"])
 }
 
 func TestKeeperTestSuite(t *testing.T) {
diff --git a/x/upgrade/module.go b/x/upgrade/module.go
index 154f006784a7..23311c22d5f3 100644
--- a/x/upgrade/module.go
+++ b/x/upgrade/module.go
@@ -101,7 +101,7 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 }
 
 // InitGenesis is ignored, no sense in serializing future upgrades
-func (am AppModule) InitGenesis(ctx sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
+func (am AppModule) InitGenesis(_ sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
 	return []abci.ValidatorUpdate{}
 }
 
diff --git a/x/upgrade/module_test.go b/x/upgrade/module_test.go
deleted file mode 100644
index 491bacc6da95..000000000000
--- a/x/upgrade/module_test.go
+++ /dev/null
@@ -1,24 +0,0 @@
-package upgrade_test
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/require"
-	abcitypes "github.com/tendermint/tendermint/abci/types"
-
-	"github.com/cosmos/cosmos-sdk/simapp"
-)
-
-func TestItGetsModuleVersions(t *testing.T) {
-	app := simapp.Setup(false)
-
-	app.InitChain(
-		abcitypes.RequestInitChain{
-			AppStateBytes: []byte("{}"),
-			ChainId:       "test-chain-id",
-		},
-	)
-
-	migmap := app.GetConsensusVersions()
-	require.NotEmpty(t, migmap)
-}
diff --git a/x/upgrade/types/handler.go b/x/upgrade/types/handler.go
index cc1cd78a495e..4344475c32ea 100644
--- a/x/upgrade/types/handler.go
+++ b/x/upgrade/types/handler.go
@@ -6,4 +6,4 @@ import (
 )
 
 // UpgradeHandler specifies the type of function that is called when an upgrade is applied
-type UpgradeHandler func(ctx sdk.Context, plan Plan, migrationMap module.MigrationMap) error
+type UpgradeHandler func(ctx sdk.Context, plan Plan, versionMap module.VersionMap) error
diff --git a/x/upgrade/types/keys.go b/x/upgrade/types/keys.go
index 593c8aeaa892..2505bd5ce451 100644
--- a/x/upgrade/types/keys.go
+++ b/x/upgrade/types/keys.go
@@ -22,8 +22,8 @@ const (
 	// DoneByte is a prefix for to look up completed upgrade plan by name
 	DoneByte = 0x1
 
-	// MigrationMapByte is a prefix to look up module names (key) and versions (value)
-	MigrationMapByte = 0x2
+	// VersionMapByte is a prefix to look up module names (key) and versions (value)
+	VersionMapByte = 0x2
 
 	// KeyUpgradedIBCState is the key under which upgraded ibc state is stored in the upgrade store
 	KeyUpgradedIBCState = "upgradedIBCState"

From 4c48f3503fd6d3b9e1edf8c2b5e3c801263fa21b Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Thu, 4 Mar 2021 10:23:27 -0800
Subject: [PATCH 04/26] updated docs

---
 x/upgrade/spec/01_concepts.md | 2 +-
 x/upgrade/spec/02_state.md    | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/x/upgrade/spec/01_concepts.md b/x/upgrade/spec/01_concepts.md
index 54147f8b5ebd..226e6fa8a7d4 100644
--- a/x/upgrade/spec/01_concepts.md
+++ b/x/upgrade/spec/01_concepts.md
@@ -48,7 +48,7 @@ and not defined on a per-module basis. Registering a `Handler` is done via
 `Keeper#SetUpgradeHandler` in the application.
 
 ```go
-type UpgradeHandler func(Context, Plan)
+type UpgradeHandler func(Context, Plan, VersionMap) error
 ```
 
 During each `EndBlock` execution, the `x/upgrade` module checks if there exists a
diff --git a/x/upgrade/spec/02_state.md b/x/upgrade/spec/02_state.md
index f069b4f96797..a225cbc1c56d 100644
--- a/x/upgrade/spec/02_state.md
+++ b/x/upgrade/spec/02_state.md
@@ -5,7 +5,10 @@ order: 2
 # State
 
 The internal state of the `x/upgrade` module is relatively minimal and simple. The
-state only contains the currently active upgrade `Plan` (if one exists) by key
-`0x0` and if a `Plan` is marked as "done" by key `0x1`.
+state contains the currently active upgrade `Plan` (if one exists) by key
+`0x0` and if a `Plan` is marked as "done" by key `0x1`. Additionally, the state
+contains versions of all app modules in the application and their respective
+consensus versions. The versions are stored as little endian `uint64`, and can be 
+accessed with prefix `0x2` appended by the corresponding module name of type `string`. 
 
 The `x/upgrade` module contains no genesis state.

From 6db21e364c11bc8b503c31f3e59eee4d2d0c4ae3 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Thu, 4 Mar 2021 10:23:54 -0800
Subject: [PATCH 05/26] forgot this

---
 x/upgrade/spec/02_state.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/x/upgrade/spec/02_state.md b/x/upgrade/spec/02_state.md
index a225cbc1c56d..be08cdb72a07 100644
--- a/x/upgrade/spec/02_state.md
+++ b/x/upgrade/spec/02_state.md
@@ -7,7 +7,7 @@ order: 2
 The internal state of the `x/upgrade` module is relatively minimal and simple. The
 state contains the currently active upgrade `Plan` (if one exists) by key
 `0x0` and if a `Plan` is marked as "done" by key `0x1`. Additionally, the state
-contains versions of all app modules in the application and their respective
+contains the consensus versions of all app modules in the application and their respective
 consensus versions. The versions are stored as little endian `uint64`, and can be 
 accessed with prefix `0x2` appended by the corresponding module name of type `string`. 
 

From c9b435a62ff161bc66975d4a1368edbe1d718d26 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Thu, 4 Mar 2021 13:33:07 -0800
Subject: [PATCH 06/26] added line to changelog for this PR

---
 CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d8e4e2007d0b..d6861763fd7e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -53,6 +53,8 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [\#8629](https://github.com/cosmos/cosmos-sdk/pull/8629) Deprecated `SetFullFundraiserPath` from `Config` in favor of `SetPurpose` and `SetCoinType`.
 * (x/upgrade) [\#8673](https://github.com/cosmos/cosmos-sdk/pull/8673) Remove IBC logic from x/upgrade. Deprecates IBC fields in an Upgrade Plan. IBC upgrade logic moved to 02-client and an IBC UpgradeProposal is added. 
 * (x/bank) [\#8517](https://github.com/cosmos/cosmos-sdk/pull/8517) `SupplyI` interface and `Supply` are removed and uses `sdk.Coins` for supply tracking
+* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) `UpgradeHandler` requires new argument `VersionManager` which retreives module consensus versions from upgrade's
+new store. Upgrade keeper requires a `VersionManager` to be set after `module.Manager` has been initialized. `ApplyUpgrade` method requires `SetConsensusVersions` to be called after handling an upgrade. 
 
 ### State Machine Breaking
 

From 1c5b13245e11d1d1b53979268b971ec173d68efe Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Mon, 8 Mar 2021 09:09:57 -0800
Subject: [PATCH 07/26] Change set consensus version function to match adr 041
 spec

---
 x/upgrade/keeper/keeper.go      | 7 ++++---
 x/upgrade/keeper/keeper_test.go | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index fe322f8e213c..0e4ad57c93be 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -54,8 +54,8 @@ func (k *Keeper) SetVersionManager(vm module.VersionManager) {
 	k.versionManager = vm
 }
 
-// SetConsensusVersions saves the consensus versions retrieved from module.Manager
-func (k Keeper) SetConsensusVersions(ctx sdk.Context) {
+// SetCurrentConsensusVersions saves the consensus versions retrieved from module.Manager
+func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) {
 	if k.versionManager == nil {
 		panic("Upgrade Keeper VersionManager was nil")
 	}
@@ -79,6 +79,7 @@ func (k Keeper) GetConsensusVersions(ctx sdk.Context) module.VersionMap {
 	defer it.Close()
 	for ; it.Valid(); it.Next() {
 		moduleBytes := it.Key()
+		// first byte is prefix key, so we remove it here
 		name := string(moduleBytes[1:])
 		moduleVersion := binary.LittleEndian.Uint64(it.Value())
 		vermap[name] = moduleVersion
@@ -235,7 +236,7 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic(err)
 	}
 
-	k.SetConsensusVersions(ctx)
+	k.SetCurrentConsensusVersions(ctx)
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
 	// This will prevent resubmission of upgrade msg after upgrade is already completed.
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index 04d0c7fa1658..021962537d8f 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -227,7 +227,7 @@ func (m mockVersionManager) GetConsensusVersions() module.VersionMap {
 func (s *KeeperTestSuite) TestMigrations() {
 	mockVM := mockVersionManager{}
 	s.app.UpgradeKeeper.SetVersionManager(mockVM)
-	s.app.UpgradeKeeper.SetConsensusVersions(s.ctx)
+	s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx)
 	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
 	dummyPlan := types.Plan{
 		Name: "dummy",

From 304abfa9a8575be66e9b475214e738e761a2c99d Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Mon, 8 Mar 2021 09:10:07 -0800
Subject: [PATCH 08/26] add documentation

---
 CHANGELOG.md               |  2 +-
 x/upgrade/spec/02_state.md | 13 +++++++++----
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d6861763fd7e..b3c0a823164e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -54,7 +54,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (x/upgrade) [\#8673](https://github.com/cosmos/cosmos-sdk/pull/8673) Remove IBC logic from x/upgrade. Deprecates IBC fields in an Upgrade Plan. IBC upgrade logic moved to 02-client and an IBC UpgradeProposal is added. 
 * (x/bank) [\#8517](https://github.com/cosmos/cosmos-sdk/pull/8517) `SupplyI` interface and `Supply` are removed and uses `sdk.Coins` for supply tracking
 * (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) `UpgradeHandler` requires new argument `VersionManager` which retreives module consensus versions from upgrade's
-new store. Upgrade keeper requires a `VersionManager` to be set after `module.Manager` has been initialized. `ApplyUpgrade` method requires `SetConsensusVersions` to be called after handling an upgrade. 
+new store.
 
 ### State Machine Breaking
 
diff --git a/x/upgrade/spec/02_state.md b/x/upgrade/spec/02_state.md
index be08cdb72a07..14ef7794c3ae 100644
--- a/x/upgrade/spec/02_state.md
+++ b/x/upgrade/spec/02_state.md
@@ -7,8 +7,13 @@ order: 2
 The internal state of the `x/upgrade` module is relatively minimal and simple. The
 state contains the currently active upgrade `Plan` (if one exists) by key
 `0x0` and if a `Plan` is marked as "done" by key `0x1`. Additionally, the state
-contains the consensus versions of all app modules in the application and their respective
-consensus versions. The versions are stored as little endian `uint64`, and can be 
-accessed with prefix `0x2` appended by the corresponding module name of type `string`. 
+contains the consensus versions of all app modules in the application. The versions 
+are stored as little endian `uint64`, and can be accessed with prefix `0x2` appended 
+by the corresponding module name of type `string`. 
 
-The `x/upgrade` module contains no genesis state.
+- Plan: `0x0 -> Plan`
+- Done: `0x1 | byte(plan name)  -> BigEndian(Block Height)`
+- ConsensusVersion: `0x2 | byte(module name)  -> LittleEndian(Module Consensus Version)`
+
+
+The `x/upgrade` module contains no genesis state.
\ No newline at end of file

From e53423ed54366612204f7eed429038d6ec57ebd8 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+technicallyty@users.noreply.github.com>
Date: Mon, 8 Mar 2021 09:25:06 -0800
Subject: [PATCH 09/26] remove newline from changelog

unnecessary newline removed
---
 CHANGELOG.md | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9dc00889a0bd..c0b806de5d15 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -58,8 +58,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [\#8629](https://github.com/cosmos/cosmos-sdk/pull/8629) Deprecated `SetFullFundraiserPath` from `Config` in favor of `SetPurpose` and `SetCoinType`.
 * (x/upgrade) [\#8673](https://github.com/cosmos/cosmos-sdk/pull/8673) Remove IBC logic from x/upgrade. Deprecates IBC fields in an Upgrade Plan. IBC upgrade logic moved to 02-client and an IBC UpgradeProposal is added.
 * (x/bank) [\#8517](https://github.com/cosmos/cosmos-sdk/pull/8517) `SupplyI` interface and `Supply` are removed and uses `sdk.Coins` for supply tracking
-* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) `UpgradeHandler` requires new argument `VersionManager` which retreives module consensus versions from upgrade's
-new store.
+* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) `UpgradeHandler` requires new argument `VersionManager` which retreives module consensus versions from upgrade's new store.
 
 ### State Machine Breaking
 

From c36086fa44013f20a8dd00aba65309c54542de99 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Mon, 8 Mar 2021 11:41:32 -0800
Subject: [PATCH 10/26] updated example in simapp for RunMigrations,
 SetCurrentConsensusVersions now returns an error

---
 simapp/app.go                   | 14 +++-----------
 x/upgrade/keeper/keeper.go      | 10 +++++++---
 x/upgrade/keeper/keeper_test.go |  4 +++-
 3 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/simapp/app.go b/simapp/app.go
index 6c9ade83d570..f8db78d9ac03 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -322,8 +322,6 @@ func NewSimApp(
 	)
 	// give upgrade keeper the module manager
 	app.UpgradeKeeper.SetVersionManager(app.mm)
-	// pass the updated keeper to the module manager
-	app.mm.Modules[upgradetypes.ModuleName] = upgrade.NewAppModule(app.UpgradeKeeper)
 
 	// During begin block slashing happens after distr.BeginBlocker so that
 	// there is nothing left over in the validator fee pool, so as to keep the
@@ -341,7 +339,7 @@ func NewSimApp(
 	// so that other modules that want to create or claim capabilities afterwards in InitChain
 	// can do so safely.
 	app.mm.SetOrderInitGenesis(
-		upgradetypes.ModuleName, capabilitytypes.ModuleName, authtypes.ModuleName, banktypes.ModuleName, distrtypes.ModuleName, stakingtypes.ModuleName,
+		capabilitytypes.ModuleName, authtypes.ModuleName, banktypes.ModuleName, distrtypes.ModuleName, stakingtypes.ModuleName,
 		slashingtypes.ModuleName, govtypes.ModuleName, minttypes.ModuleName, crisistypes.ModuleName,
 		genutiltypes.ModuleName, evidencetypes.ModuleName, authztypes.ModuleName,
 		feegranttypes.ModuleName,
@@ -555,14 +553,8 @@ func (app *SimApp) RegisterTendermintService(clientCtx client.Context) {
 //
 // Example:
 //   cfg := module.NewConfigurator(...)
-//   app.UpgradeKeeper.SetUpgradeHandler("store-migration", func(ctx sdk.Context, plan upgradetypes.Plan) {
-//       err := app.RunMigrations(ctx, module.VersionMap{
-//           "bank": 1,     // Migrate x/bank from v1 to current x/bank's ConsensusVersion
-//           "staking": 8,  // Migrate x/staking from v8 to current x/staking's ConsensusVersion
-//      })
-//      if err != nil {
-//           panic(err)
-//      }
+//   app.UpgradeKeeper.SetUpgradeHandler("store-migration", func(ctx sdk.Context, plan upgradetypes.Plan, versionMap module.VersionMap) {
+//       return app.RunMigrations(ctx, versionMap)
 //   })
 func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.VersionMap) error {
 	return app.mm.RunMigrations(ctx, app.configurator, migrateFromVersions)
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index 0e4ad57c93be..dc34c0f52397 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -55,9 +55,9 @@ func (k *Keeper) SetVersionManager(vm module.VersionManager) {
 }
 
 // SetCurrentConsensusVersions saves the consensus versions retrieved from module.Manager
-func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) {
+func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) error {
 	if k.versionManager == nil {
-		panic("Upgrade Keeper VersionManager was nil")
+		return sdkerrors.Wrap(sdkerrors.ErrLogic, "upgrade keeper VersionManager was nil, please call SetVersionManager on the keeper first")
 	}
 	modules := k.versionManager.GetConsensusVersions()
 	store := ctx.KVStore(k.storeKey)
@@ -68,6 +68,7 @@ func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) {
 		binary.LittleEndian.PutUint64(verBytes, ver)
 		versionStore.Set(nameBytes, verBytes)
 	}
+	return nil
 }
 
 // GetConsensusVersions gets a VersionMap from state
@@ -236,7 +237,10 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic(err)
 	}
 
-	k.SetCurrentConsensusVersions(ctx)
+	err = k.SetCurrentConsensusVersions(ctx)
+	if err != nil {
+		panic(err)
+	}
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
 	// This will prevent resubmission of upgrade msg after upgrade is already completed.
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index 021962537d8f..f45f1dad2593 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -227,7 +227,9 @@ func (m mockVersionManager) GetConsensusVersions() module.VersionMap {
 func (s *KeeperTestSuite) TestMigrations() {
 	mockVM := mockVersionManager{}
 	s.app.UpgradeKeeper.SetVersionManager(mockVM)
-	s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx)
+	if err := s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx); err != nil {
+		s.T().Fatalf("error setting consensus versions: %v", err)
+	}
 	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
 	dummyPlan := types.Plan{
 		Name: "dummy",

From 758c67d41b7c6d0658a84f5200bc59878c7458e8 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Mon, 8 Mar 2021 12:26:53 -0800
Subject: [PATCH 11/26] switch TestMigrations to use Require instead of t.Fatal

---
 x/upgrade/keeper/keeper_test.go | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index f45f1dad2593..41f9ab713bde 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -227,9 +227,8 @@ func (m mockVersionManager) GetConsensusVersions() module.VersionMap {
 func (s *KeeperTestSuite) TestMigrations() {
 	mockVM := mockVersionManager{}
 	s.app.UpgradeKeeper.SetVersionManager(mockVM)
-	if err := s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx); err != nil {
-		s.T().Fatalf("error setting consensus versions: %v", err)
-	}
+	err := s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx)
+	s.Require().Nil(err)
 	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
 	dummyPlan := types.Plan{
 		Name: "dummy",

From 0dac9c228b8bf882dd3e80b38c28c296ee4fa2a2 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+technicallyty@users.noreply.github.com>
Date: Wed, 10 Mar 2021 10:27:18 -0800
Subject: [PATCH 12/26] Update CHANGELOG.md

Co-authored-by: Aaron Craelius <aaron@regen.network>
---
 CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c0b806de5d15..05b20f9ce3c4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -58,7 +58,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [\#8629](https://github.com/cosmos/cosmos-sdk/pull/8629) Deprecated `SetFullFundraiserPath` from `Config` in favor of `SetPurpose` and `SetCoinType`.
 * (x/upgrade) [\#8673](https://github.com/cosmos/cosmos-sdk/pull/8673) Remove IBC logic from x/upgrade. Deprecates IBC fields in an Upgrade Plan. IBC upgrade logic moved to 02-client and an IBC UpgradeProposal is added.
 * (x/bank) [\#8517](https://github.com/cosmos/cosmos-sdk/pull/8517) `SupplyI` interface and `Supply` are removed and uses `sdk.Coins` for supply tracking
-* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) `UpgradeHandler` requires new argument `VersionManager` which retreives module consensus versions from upgrade's new store.
+* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) `UpgradeHandler` includes a new argument `VersionManager` which retreives module consensus versions from upgrade's new store.
 
 ### State Machine Breaking
 

From 8f526ef98f1fbfda957448f4136e3a6f72a99725 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Wed, 10 Mar 2021 11:20:54 -0800
Subject: [PATCH 13/26] docs for SetVersionManager

---
 x/upgrade/keeper/keeper.go | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index dc34c0f52397..fe5ccd9fcaa5 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -50,6 +50,10 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
+// SetVersionManager sets a VersionManager for the keeper to
+// gain access to modules and their consensus versions.
+// This MUST be set in app.go, or versions will not be committed
+// to state.
 func (k *Keeper) SetVersionManager(vm module.VersionManager) {
 	k.versionManager = vm
 }

From 027740dea5e3770070111f84e68e7f71b6531fde Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Thu, 11 Mar 2021 11:31:54 -0800
Subject: [PATCH 14/26] -init genesis method added -removed panics/fails from
 setting consensus versions

---
 x/upgrade/abci.go               |  4 ++++
 x/upgrade/keeper/keeper.go      | 33 +++++++++++++++------------------
 x/upgrade/keeper/keeper_test.go |  3 +--
 x/upgrade/module.go             |  3 ++-
 x/upgrade/module_test.go        | 24 ++++++++++++++++++++++++
 5 files changed, 46 insertions(+), 21 deletions(-)
 create mode 100644 x/upgrade/module_test.go

diff --git a/x/upgrade/abci.go b/x/upgrade/abci.go
index d346decb023c..6e85abb1e7b5 100644
--- a/x/upgrade/abci.go
+++ b/x/upgrade/abci.go
@@ -69,6 +69,10 @@ func BeginBlocker(k keeper.Keeper, ctx sdk.Context, _ abci.RequestBeginBlock) {
 	}
 }
 
+func InitChainer(k keeper.Keeper, ctx sdk.Context) {
+	k.SetCurrentConsensusVersions(ctx)
+}
+
 // BuildUpgradeNeededMsg prints the message that notifies that an upgrade is needed.
 func BuildUpgradeNeededMsg(plan types.Plan) string {
 	return fmt.Sprintf("UPGRADE \"%s\" NEEDED at %s: %s", plan.Name, plan.DueAt(), plan.Info)
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index fe5ccd9fcaa5..dd5d7596f744 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -51,7 +51,7 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 }
 
 // SetVersionManager sets a VersionManager for the keeper to
-// gain access to modules and their consensus versions.
+// gain access to app modules and their consensus versions.
 // This MUST be set in app.go, or versions will not be committed
 // to state.
 func (k *Keeper) SetVersionManager(vm module.VersionManager) {
@@ -59,20 +59,20 @@ func (k *Keeper) SetVersionManager(vm module.VersionManager) {
 }
 
 // SetCurrentConsensusVersions saves the consensus versions retrieved from module.Manager
-func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) error {
-	if k.versionManager == nil {
-		return sdkerrors.Wrap(sdkerrors.ErrLogic, "upgrade keeper VersionManager was nil, please call SetVersionManager on the keeper first")
-	}
-	modules := k.versionManager.GetConsensusVersions()
-	store := ctx.KVStore(k.storeKey)
-	versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
-	for modName, ver := range modules {
-		nameBytes := []byte(modName)
-		verBytes := make([]byte, 8)
-		binary.LittleEndian.PutUint64(verBytes, ver)
-		versionStore.Set(nameBytes, verBytes)
+// if versionManager is not set from app.go, consensus versions will NOT be saved
+// to state.
+func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) {
+	if k.versionManager != nil {
+		modules := k.versionManager.GetConsensusVersions()
+		store := ctx.KVStore(k.storeKey)
+		versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
+		for modName, ver := range modules {
+			nameBytes := []byte(modName)
+			verBytes := make([]byte, 8)
+			binary.LittleEndian.PutUint64(verBytes, ver)
+			versionStore.Set(nameBytes, verBytes)
+		}
 	}
-	return nil
 }
 
 // GetConsensusVersions gets a VersionMap from state
@@ -241,10 +241,7 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic(err)
 	}
 
-	err = k.SetCurrentConsensusVersions(ctx)
-	if err != nil {
-		panic(err)
-	}
+	k.SetCurrentConsensusVersions(ctx)
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
 	// This will prevent resubmission of upgrade msg after upgrade is already completed.
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index 41f9ab713bde..021962537d8f 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -227,8 +227,7 @@ func (m mockVersionManager) GetConsensusVersions() module.VersionMap {
 func (s *KeeperTestSuite) TestMigrations() {
 	mockVM := mockVersionManager{}
 	s.app.UpgradeKeeper.SetVersionManager(mockVM)
-	err := s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx)
-	s.Require().Nil(err)
+	s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx)
 	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
 	dummyPlan := types.Plan{
 		Name: "dummy",
diff --git a/x/upgrade/module.go b/x/upgrade/module.go
index 23311c22d5f3..be5baec3d496 100644
--- a/x/upgrade/module.go
+++ b/x/upgrade/module.go
@@ -101,7 +101,8 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 }
 
 // InitGenesis is ignored, no sense in serializing future upgrades
-func (am AppModule) InitGenesis(_ sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
+func (am AppModule) InitGenesis(ctx sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
+	InitChainer(am.keeper, ctx)
 	return []abci.ValidatorUpdate{}
 }
 
diff --git a/x/upgrade/module_test.go b/x/upgrade/module_test.go
new file mode 100644
index 000000000000..f23f51fa90c5
--- /dev/null
+++ b/x/upgrade/module_test.go
@@ -0,0 +1,24 @@
+package upgrade_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	abcitypes "github.com/tendermint/tendermint/abci/types"
+
+	"github.com/cosmos/cosmos-sdk/simapp"
+)
+
+func TestInitChainer(t *testing.T) {
+	app := simapp.Setup(false)
+
+	app.InitChain(
+		abcitypes.RequestInitChain{
+			AppStateBytes: []byte("{}"),
+			ChainId:       "test-chain-id",
+		},
+	)
+
+	versionMap := app.GetConsensusVersions()
+	require.NotEmpty(t, versionMap)
+}

From fc38e5a005bca1ad999e125731e7954fb1bf94dc Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Fri, 12 Mar 2021 13:35:03 -0800
Subject: [PATCH 15/26] update identifiers to be more go-like

---
 simapp/app.go                   |  6 +++---
 types/module/module.go          | 10 +++++-----
 x/upgrade/abci_test.go          |  6 +++---
 x/upgrade/keeper/keeper.go      |  8 ++++----
 x/upgrade/keeper/keeper_test.go |  6 ++++--
 x/upgrade/module_test.go        |  2 +-
 6 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/simapp/app.go b/simapp/app.go
index 81cd35b32d28..88a9cf7874fd 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -544,9 +544,9 @@ func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.Ver
 	return app.mm.RunMigrations(ctx, app.configurator, migrateFromVersions)
 }
 
-// Returns a map (VersionMap) of key module name and value module consensus version
-func (app *SimApp) GetConsensusVersions() module.VersionMap {
-	return app.mm.GetConsensusVersions()
+// GetVersionMap gets a map of key module name and value module consensus version
+func (app *SimApp) GetVersionMap() module.VersionMap {
+	return app.mm.GetVersionMap()
 }
 
 // RegisterSwaggerAPI registers swagger route with API Server
diff --git a/types/module/module.go b/types/module/module.go
index 00e754fdb503..4a47bbf92e23 100644
--- a/types/module/module.go
+++ b/types/module/module.go
@@ -338,18 +338,18 @@ type MigrationHandler func(sdk.Context) error
 type VersionMap map[string]uint64
 
 type VersionManager interface {
-	GetConsensusVersions() VersionMap
+	GetVersionMap() VersionMap
 }
 
 // RunMigrations performs in-place store migrations for all modules.
-func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, migrateFromVersions VersionMap) error {
+func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, vm VersionMap) error {
 	c, ok := cfg.(configurator)
 	if !ok {
 		return sdkerrors.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", configurator{}, cfg)
 	}
 
 	for moduleName, module := range m.Modules {
-		err := c.runModuleMigrations(ctx, moduleName, migrateFromVersions[moduleName], module.ConsensusVersion())
+		err := c.runModuleMigrations(ctx, moduleName, vm[moduleName], module.ConsensusVersion())
 		if err != nil {
 			return err
 		}
@@ -400,8 +400,8 @@ func (m *Manager) EndBlock(ctx sdk.Context, req abci.RequestEndBlock) abci.Respo
 	}
 }
 
-// GetConsensusVersions gets consensus version from all modules
-func (m *Manager) GetConsensusVersions() VersionMap {
+// GetVersionMap gets consensus version from all modules
+func (m *Manager) GetVersionMap() VersionMap {
 	vermap := make(VersionMap)
 	for _, v := range m.Modules {
 		version := v.ConsensusVersion()
diff --git a/x/upgrade/abci_test.go b/x/upgrade/abci_test.go
index 514480dfde12..6e309717481a 100644
--- a/x/upgrade/abci_test.go
+++ b/x/upgrade/abci_test.go
@@ -128,7 +128,7 @@ func VerifyDoUpgrade(t *testing.T) {
 	})
 
 	t.Log("Verify that the upgrade can be successfully applied with a handler")
-	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan, versionMap module.VersionMap) error { return nil })
+	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) error { return nil })
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -144,7 +144,7 @@ func VerifyDoUpgradeWithCtx(t *testing.T, newCtx sdk.Context, proposalName strin
 	})
 
 	t.Log("Verify that the upgrade can be successfully applied with a handler")
-	s.keeper.SetUpgradeHandler(proposalName, func(ctx sdk.Context, plan types.Plan, versionMap module.VersionMap) error { return nil })
+	s.keeper.SetUpgradeHandler(proposalName, func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) error { return nil })
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -156,7 +156,7 @@ func TestHaltIfTooNew(t *testing.T) {
 	s := setupTest(10, map[int64]bool{})
 	t.Log("Verify that we don't panic with registered plan not in database at all")
 	var called int
-	s.keeper.SetUpgradeHandler("future", func(ctx sdk.Context, plan types.Plan, versionMap module.VersionMap) error { called++; return nil })
+	s.keeper.SetUpgradeHandler("future", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) error { called++; return nil })
 
 	newCtx := s.ctx.WithBlockHeight(s.ctx.BlockHeight() + 1).WithBlockTime(time.Now())
 	req := abci.RequestBeginBlock{Header: newCtx.BlockHeader()}
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index dd5d7596f744..4a75ac6b0cfb 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -63,7 +63,7 @@ func (k *Keeper) SetVersionManager(vm module.VersionManager) {
 // to state.
 func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) {
 	if k.versionManager != nil {
-		modules := k.versionManager.GetConsensusVersions()
+		modules := k.versionManager.GetVersionMap()
 		store := ctx.KVStore(k.storeKey)
 		versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
 		for modName, ver := range modules {
@@ -75,8 +75,8 @@ func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) {
 	}
 }
 
-// GetConsensusVersions gets a VersionMap from state
-func (k Keeper) GetConsensusVersions(ctx sdk.Context) module.VersionMap {
+// GetVersionMap gets a VersionMap from state
+func (k Keeper) GetVersionMap(ctx sdk.Context) module.VersionMap {
 	store := ctx.KVStore(k.storeKey)
 	it := sdk.KVStorePrefixIterator(store, []byte{types.VersionMapByte})
 
@@ -236,7 +236,7 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic("ApplyUpgrade should never be called without first checking HasHandler")
 	}
 
-	err := handler(ctx, plan, k.GetConsensusVersions(ctx))
+	err := handler(ctx, plan, k.GetVersionMap(ctx))
 	if err != nil {
 		panic(err)
 	}
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index 021962537d8f..788d42ee9265 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -216,7 +216,7 @@ func (s *KeeperTestSuite) TestSetUpgradedClient() {
 // Mock version manager for TestMigrations
 type mockVersionManager struct{}
 
-func (m mockVersionManager) GetConsensusVersions() module.VersionMap {
+func (m mockVersionManager) GetVersionMap() module.VersionMap {
 	vermap := make(module.VersionMap)
 	vermap["bank"] = 1
 	return vermap
@@ -228,6 +228,7 @@ func (s *KeeperTestSuite) TestMigrations() {
 	mockVM := mockVersionManager{}
 	s.app.UpgradeKeeper.SetVersionManager(mockVM)
 	s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx)
+	vermapBefore := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
 	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
 	dummyPlan := types.Plan{
 		Name: "dummy",
@@ -237,8 +238,9 @@ func (s *KeeperTestSuite) TestMigrations() {
 
 	s.app.UpgradeKeeper.SetVersionManager(s.app)
 	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, dummyPlan)
-	vermap := s.app.UpgradeKeeper.GetConsensusVersions(s.ctx)
+	vermap := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
 	s.Require().Equal(bank.AppModule{}.ConsensusVersion(), vermap["bank"])
+	s.Require().Greater(vermap["bank"], vermapBefore["bank"])
 }
 
 func TestKeeperTestSuite(t *testing.T) {
diff --git a/x/upgrade/module_test.go b/x/upgrade/module_test.go
index f23f51fa90c5..5b162ac10e52 100644
--- a/x/upgrade/module_test.go
+++ b/x/upgrade/module_test.go
@@ -19,6 +19,6 @@ func TestInitChainer(t *testing.T) {
 		},
 	)
 
-	versionMap := app.GetConsensusVersions()
+	versionMap := app.GetVersionMap()
 	require.NotEmpty(t, versionMap)
 }

From 1fca959460017eaba3083cbc983304e9636ebe1a Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Mon, 15 Mar 2021 11:33:22 -0700
Subject: [PATCH 16/26] update docs and UpgradeHandler fnc sig

---
 x/upgrade/keeper/keeper.go | 3 ++-
 x/upgrade/types/handler.go | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index 4a75ac6b0cfb..47558c43060a 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -75,7 +75,8 @@ func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) {
 	}
 }
 
-// GetVersionMap gets a VersionMap from state
+// GetConsensusVersions returns a map of key module name and value module consensus version
+// as defined in ADR-041.
 func (k Keeper) GetVersionMap(ctx sdk.Context) module.VersionMap {
 	store := ctx.KVStore(k.storeKey)
 	it := sdk.KVStorePrefixIterator(store, []byte{types.VersionMapByte})
diff --git a/x/upgrade/types/handler.go b/x/upgrade/types/handler.go
index 4344475c32ea..0fa9212cec93 100644
--- a/x/upgrade/types/handler.go
+++ b/x/upgrade/types/handler.go
@@ -6,4 +6,4 @@ import (
 )
 
 // UpgradeHandler specifies the type of function that is called when an upgrade is applied
-type UpgradeHandler func(ctx sdk.Context, plan Plan, versionMap module.VersionMap) error
+type UpgradeHandler func(ctx sdk.Context, plan Plan, vm module.VersionMap) error

From a99c538664150285bb03fa911fa283fbae969ba0 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Mon, 15 Mar 2021 14:45:14 -0700
Subject: [PATCH 17/26] Upgrade Keeper now takes a VersionMap instead of a
 VersionManager interface

---
 .../adr-041-in-place-store-migrations.md      | 28 +++++--------
 docs/run-node/cosmovisor.md                   |  2 +-
 simapp/app.go                                 | 13 +++---
 types/module/module.go                        |  4 --
 x/upgrade/abci_test.go                        | 13 ++++--
 x/upgrade/keeper/grpc_query_test.go           |  4 +-
 x/upgrade/keeper/keeper.go                    | 38 +++++++++---------
 x/upgrade/keeper/keeper_test.go               | 40 +++++++++----------
 x/upgrade/module_test.go                      |  5 ++-
 x/upgrade/spec/01_concepts.md                 |  2 +-
 x/upgrade/spec/02_state.md                    |  2 +-
 x/upgrade/types/handler.go                    |  2 +-
 12 files changed, 74 insertions(+), 79 deletions(-)

diff --git a/docs/architecture/adr-041-in-place-store-migrations.md b/docs/architecture/adr-041-in-place-store-migrations.md
index 072a121f5d66..b10c3a6982a9 100644
--- a/docs/architecture/adr-041-in-place-store-migrations.md
+++ b/docs/architecture/adr-041-in-place-store-migrations.md
@@ -82,48 +82,42 @@ Each module's migration functions are specific to the module's store evolutions,
 We introduce a new prefix store in `x/upgrade`'s store. This store will track each module's current version, it can be modelized as a `map[string]uint64` of module name to module ConsensusVersion, and will be used when running the migrations (see next section for details). The key prefix used is `0x1`, and the key/value format is:
 
 ```
-0x2 | {bytes(module_name)} => LittleEndian(module_consensus_version)
+0x2 | {bytes(module_name)} => BigEndian(module_consensus_version)
 ```
 
-s
-We add a new private field `versionManager` of type `VersionManager` to `x/upgrade`'s keeper, where `VersionManager` is:
+We add a new private field `versionMap` of type `VersionMap` to `x/upgrade`'s keeper, where `VersionManager` is:
 
 ```go
-type VersionManager interface {
-    GetConsensusVersions() VersionMap
-}
-
 // Map of module name => new module Consensus Version.
 type VersionMap map[string]uint64
 ```
 
-This `versionManager` field can be modified via the `SetVersionManager` field, and will allow the upgrade keeper to know the current versions of loaded modules. `SetVersionManager` MUST be called as early as possible in the app initialization; in the SDK's `simapp`, it is called in the `NewSimApp` constructor function.
+This `versionMap` field can be modified via the `SetInitialVersionMap` field, and will allow the upgrade keeper to know the current versions of loaded modules. `SetInitialVersionMap` MUST be called as early as possible in the app initialization; in the SDK's `simapp`, it is called in the `NewSimApp` constructor function.
 
-The UpgradeHandler signature needs to be updated to take a `VersionMap`, as well as return an error:
+The UpgradeHandler signature needs to be updated to take a `VersionMap`, as well as return a new `VersionMap` and an error:
 
 ```diff
 - type UpgradeHandler func(ctx sdk.Context, plan Plan)
-+ type UpgradeHandler func(ctx sdk.Context, plan Plan, versionMap VersionMap) error
++ type UpgradeHandler func(ctx sdk.Context, plan Plan, versionMap VersionMap) (VersionMap, error)
 ```
 
-To apply an upgrade, we query the `VersionMap` from the `x/upgrade` store and pass it into the handler. The handler runs the actual migration functions (see next section), and if successful, the current ConsensusVersions of all loaded modules will be stored into state.
+To apply an upgrade, we query the `VersionMap` from the `x/upgrade` store and pass it into the handler. The handler runs the actual migration functions (see next section), and if successful, returns the updated ConsensusVersions of the modules to be stored in state.
 
 ```diff
 func (k UpgradeKeeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
     // --snip--
 -   handler(ctx, plan)
-+   err := handler(ctx, plan, k.GetConsensusVersions()) // k.GetConsensusVersions() fetches the VersionMap stored in state.
++   updatedVM, err := handler(ctx, plan, k.GetConsensusVersions(ctx)) // k.GetConsensusVersions() fetches the VersionMap stored in state.
 +   if err != nil {
 +       return err
 +   }
 +
-+   // Get the current ConsensusVersions of the loaded modules (retrieved from
-+   // `k.versionManager`), and save them to state.
-+   k.SetCurrentConsensusVersions()
++   // Set the updated consensus versions to state
++   k.setConsensusVersions(ctx, updatedVM)
 }
 ```
 
-An gRPC query endpoint to query the `VersionMap` stored in `x/upgrade`'s state will also be added, so that app developers can double-check the `VersionMap` before the upgrade handler runs.
+A gRPC query endpoint to query the `VersionMap` stored in `x/upgrade`'s state will also be added, so that app developers can double-check the `VersionMap` before the upgrade handler runs.
 
 ### Running Migrations
 
@@ -139,7 +133,7 @@ If a required migration is missing (e.g. if it has not been registered in the `C
 In practice, the `RunMigrations` method should be called from inside an `UpgradeHandler`.
 
 ```go
-app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, versionMap VersionMap) error {
+app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, versionMap module.VersionMap)  (module.VersionMap, error) {
     return app.mm.RunMigrations(ctx, versionMap)
 })
 ```
diff --git a/docs/run-node/cosmovisor.md b/docs/run-node/cosmovisor.md
index 99cb81e0cc45..cf7575b9a31f 100644
--- a/docs/run-node/cosmovisor.md
+++ b/docs/run-node/cosmovisor.md
@@ -147,7 +147,7 @@ In `simapp/app.go`, find the line containing the upgrade Keeper initialisation,
 After that line, add the following snippet:
 
  ```
- app.UpgradeKeeper.SetUpgradeHandler("test1", func(ctx sdk.Context, plan upgradetypes.Plan) {
+ app.UpgradeKeeper.SetUpgradeHandler("test1", func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap) {
 		// Add some coins to a random account
 		addr, err := sdk.AccAddressFromBech32("cosmos18cgkqduwuh253twzmhedesw3l7v3fm37sppt58")
 		if err != nil {
diff --git a/simapp/app.go b/simapp/app.go
index 88a9cf7874fd..7c86a3b107f8 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -315,8 +315,10 @@ func NewSimApp(
 		params.NewAppModule(app.ParamsKeeper),
 		authz.NewAppModule(appCodec, app.AuthzKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
 	)
-	// give upgrade keeper the module manager
-	app.UpgradeKeeper.SetVersionManager(app.mm)
+
+	// Pass the version map to the upgrade keeper
+	app.UpgradeKeeper.SetInitialVersionMap(app.mm.GetVersionMap())
+	app.mm.Modules[upgradetypes.ModuleName] = upgrade.NewAppModule(app.UpgradeKeeper)
 
 	// During begin block slashing happens after distr.BeginBlocker so that
 	// there is nothing left over in the validator fee pool, so as to keep the
@@ -334,7 +336,7 @@ func NewSimApp(
 	// so that other modules that want to create or claim capabilities afterwards in InitChain
 	// can do so safely.
 	app.mm.SetOrderInitGenesis(
-		capabilitytypes.ModuleName, authtypes.ModuleName, banktypes.ModuleName, distrtypes.ModuleName, stakingtypes.ModuleName,
+		upgradetypes.ModuleName, capabilitytypes.ModuleName, authtypes.ModuleName, banktypes.ModuleName, distrtypes.ModuleName, stakingtypes.ModuleName,
 		slashingtypes.ModuleName, govtypes.ModuleName, minttypes.ModuleName, crisistypes.ModuleName,
 		genutiltypes.ModuleName, evidencetypes.ModuleName, authztypes.ModuleName,
 		feegranttypes.ModuleName,
@@ -544,11 +546,6 @@ func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.Ver
 	return app.mm.RunMigrations(ctx, app.configurator, migrateFromVersions)
 }
 
-// GetVersionMap gets a map of key module name and value module consensus version
-func (app *SimApp) GetVersionMap() module.VersionMap {
-	return app.mm.GetVersionMap()
-}
-
 // RegisterSwaggerAPI registers swagger route with API Server
 func RegisterSwaggerAPI(ctx client.Context, rtr *mux.Router) {
 	statikFS, err := fs.New()
diff --git a/types/module/module.go b/types/module/module.go
index 4a47bbf92e23..6e3b1d670db2 100644
--- a/types/module/module.go
+++ b/types/module/module.go
@@ -337,10 +337,6 @@ type MigrationHandler func(sdk.Context) error
 // version from which we should perform the migration for each module.
 type VersionMap map[string]uint64
 
-type VersionManager interface {
-	GetVersionMap() VersionMap
-}
-
 // RunMigrations performs in-place store migrations for all modules.
 func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, vm VersionMap) error {
 	c, ok := cfg.(configurator)
diff --git a/x/upgrade/abci_test.go b/x/upgrade/abci_test.go
index 075c637a498f..9735ae619638 100644
--- a/x/upgrade/abci_test.go
+++ b/x/upgrade/abci_test.go
@@ -105,7 +105,9 @@ func VerifyDoUpgrade(t *testing.T) {
 	})
 
 	t.Log("Verify that the upgrade can be successfully applied with a handler")
-	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) error { return nil })
+	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		return vm, nil
+	})
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -121,7 +123,9 @@ func VerifyDoUpgradeWithCtx(t *testing.T, newCtx sdk.Context, proposalName strin
 	})
 
 	t.Log("Verify that the upgrade can be successfully applied with a handler")
-	s.keeper.SetUpgradeHandler(proposalName, func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) error { return nil })
+	s.keeper.SetUpgradeHandler(proposalName, func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		return vm, nil
+	})
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -133,7 +137,10 @@ func TestHaltIfTooNew(t *testing.T) {
 	s := setupTest(10, map[int64]bool{})
 	t.Log("Verify that we don't panic with registered plan not in database at all")
 	var called int
-	s.keeper.SetUpgradeHandler("future", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) error { called++; return nil })
+	s.keeper.SetUpgradeHandler("future", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		called++
+		return vm, nil
+	})
 
 	newCtx := s.ctx.WithBlockHeight(s.ctx.BlockHeight() + 1).WithBlockTime(time.Now())
 	req := abci.RequestBeginBlock{Header: newCtx.BlockHeader()}
diff --git a/x/upgrade/keeper/grpc_query_test.go b/x/upgrade/keeper/grpc_query_test.go
index a817d6222346..b44fc1e67c00 100644
--- a/x/upgrade/keeper/grpc_query_test.go
+++ b/x/upgrade/keeper/grpc_query_test.go
@@ -111,7 +111,9 @@ func (suite *UpgradeTestSuite) TestAppliedCurrentPlan() {
 				suite.app.UpgradeKeeper.ScheduleUpgrade(suite.ctx, plan)
 
 				suite.ctx = suite.ctx.WithBlockHeight(expHeight)
-				suite.app.UpgradeKeeper.SetUpgradeHandler(planName, func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
+				suite.app.UpgradeKeeper.SetUpgradeHandler(planName, func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+					return vm, nil
+				})
 				suite.app.UpgradeKeeper.ApplyUpgrade(suite.ctx, plan)
 
 				req = &types.QueryAppliedPlanRequest{Name: planName}
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index ef24800af632..571a8e902989 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -29,7 +29,7 @@ type Keeper struct {
 	storeKey           sdk.StoreKey
 	cdc                codec.BinaryMarshaler
 	upgradeHandlers    map[string]types.UpgradeHandler
-	versionManager     module.VersionManager
+	versionMap         module.VersionMap
 }
 
 // NewKeeper constructs an upgrade Keeper
@@ -40,6 +40,7 @@ func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey sdk.StoreKey, cdc cod
 		storeKey:           storeKey,
 		cdc:                cdc,
 		upgradeHandlers:    map[string]types.UpgradeHandler{},
+		versionMap:         module.VersionMap{},
 	}
 }
 
@@ -50,26 +51,26 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
-// SetVersionManager sets a VersionManager for the keeper to
-// gain access to app modules and their consensus versions.
-// This MUST be set in app.go, or versions will not be committed
-// to state.
-func (k *Keeper) SetVersionManager(vm module.VersionManager) {
-	k.versionManager = vm
+// SetInitialVersionMap sets an initial version map on the keeper
+func (k *Keeper) SetInitialVersionMap(vm module.VersionMap) {
+	k.versionMap = vm
 }
 
-// SetCurrentConsensusVersions saves the consensus versions retrieved from module.Manager
-// if versionManager is not set from app.go, consensus versions will NOT be saved
+// SetCurrentConsensusVersions saves the consensus versions.
+// If versionMap is not set from app.go, consensus versions will NOT be saved
 // to state.
 func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) {
-	if k.versionManager != nil {
-		modules := k.versionManager.GetVersionMap()
+	k.setConsensusVersions(ctx, k.versionMap)
+}
+
+func (k Keeper) setConsensusVersions(ctx sdk.Context, vm module.VersionMap) {
+	if len(vm) > 0 {
 		store := ctx.KVStore(k.storeKey)
 		versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
-		for modName, ver := range modules {
+		for modName, ver := range vm {
 			nameBytes := []byte(modName)
 			verBytes := make([]byte, 8)
-			binary.LittleEndian.PutUint64(verBytes, ver)
+			binary.BigEndian.PutUint64(verBytes, ver)
 			versionStore.Set(nameBytes, verBytes)
 		}
 	}
@@ -81,17 +82,17 @@ func (k Keeper) GetVersionMap(ctx sdk.Context) module.VersionMap {
 	store := ctx.KVStore(k.storeKey)
 	it := sdk.KVStorePrefixIterator(store, []byte{types.VersionMapByte})
 
-	vermap := make(module.VersionMap)
+	vm := make(module.VersionMap)
 	defer it.Close()
 	for ; it.Valid(); it.Next() {
 		moduleBytes := it.Key()
 		// first byte is prefix key, so we remove it here
 		name := string(moduleBytes[1:])
-		moduleVersion := binary.LittleEndian.Uint64(it.Value())
-		vermap[name] = moduleVersion
+		moduleVersion := binary.BigEndian.Uint64(it.Value())
+		vm[name] = moduleVersion
 	}
 
-	return vermap
+	return vm
 }
 
 // ScheduleUpgrade schedules an upgrade based on the specified plan.
@@ -233,11 +234,12 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic("ApplyUpgrade should never be called without first checking HasHandler")
 	}
 
-	err := handler(ctx, plan, k.GetVersionMap(ctx))
+	updatedVM, err := handler(ctx, plan, k.GetVersionMap(ctx))
 	if err != nil {
 		panic(err)
 	}
 
+	k.versionMap = updatedVM
 	k.SetCurrentConsensusVersions(ctx)
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index c6a02fed5fa3..ed4559a029be 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -117,7 +117,9 @@ func (s *KeeperTestSuite) TestScheduleUpgrade() {
 				Height: 123450000,
 			},
 			setup: func() {
-				s.app.UpgradeKeeper.SetUpgradeHandler("all-good", func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
+				s.app.UpgradeKeeper.SetUpgradeHandler("all-good", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+					return vm, nil
+				})
 				s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, types.Plan{
 					Name:   "all-good",
 					Info:   "some text here",
@@ -134,7 +136,7 @@ func (s *KeeperTestSuite) TestScheduleUpgrade() {
 		s.Run(tc.name, func() {
 			// reset suite
 			s.SetupTest()
-			s.app.UpgradeKeeper.SetVersionManager(s.app)
+			s.app.UpgradeKeeper.SetInitialVersionMap(module.VersionMap{})
 			// setup test case
 			tc.setup()
 
@@ -193,34 +195,28 @@ func (s *KeeperTestSuite) TestSetUpgradedClient() {
 
 }
 
-// Mock version manager for TestMigrations
-type mockVersionManager struct{}
-
-func (m mockVersionManager) GetVersionMap() module.VersionMap {
-	vermap := make(module.VersionMap)
-	vermap["bank"] = 1
-	return vermap
-}
-
 // Tests that the underlying state of x/upgrade is set correctly after
 // an upgrade.
 func (s *KeeperTestSuite) TestMigrations() {
-	mockVM := mockVersionManager{}
-	s.app.UpgradeKeeper.SetVersionManager(mockVM)
+	initialVM := module.VersionMap{"bank": uint64(1)}
+	s.app.UpgradeKeeper.SetInitialVersionMap(initialVM)
 	s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx)
-	vermapBefore := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
-	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
+	vmBefore := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
+	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		// simulate upgrading the bank module
+		vm["bank"] = vm["bank"] + 1
+		return vm, nil
+	})
 	dummyPlan := types.Plan{
-		Name: "dummy",
-		Info: "some text here",
-		Time: s.ctx.BlockTime().Add(time.Hour),
+		Name:   "dummy",
+		Info:   "some text here",
+		Height: 123450000,
 	}
 
-	s.app.UpgradeKeeper.SetVersionManager(s.app)
 	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, dummyPlan)
-	vermap := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
-	s.Require().Equal(bank.AppModule{}.ConsensusVersion(), vermap["bank"])
-	s.Require().Greater(vermap["bank"], vermapBefore["bank"])
+	vm := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
+	s.Require().Equal(bank.AppModule{}.ConsensusVersion(), vm["bank"])
+	s.Require().Greater(vm["bank"], vmBefore["bank"])
 }
 
 func TestKeeperTestSuite(t *testing.T) {
diff --git a/x/upgrade/module_test.go b/x/upgrade/module_test.go
index 5b162ac10e52..5b4bf1e1a631 100644
--- a/x/upgrade/module_test.go
+++ b/x/upgrade/module_test.go
@@ -5,13 +5,14 @@ import (
 
 	"github.com/stretchr/testify/require"
 	abcitypes "github.com/tendermint/tendermint/abci/types"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 
 	"github.com/cosmos/cosmos-sdk/simapp"
 )
 
 func TestInitChainer(t *testing.T) {
 	app := simapp.Setup(false)
-
+	ctx := app.BaseApp.NewContext(false, tmproto.Header{})
 	app.InitChain(
 		abcitypes.RequestInitChain{
 			AppStateBytes: []byte("{}"),
@@ -19,6 +20,6 @@ func TestInitChainer(t *testing.T) {
 		},
 	)
 
-	versionMap := app.GetVersionMap()
+	versionMap := app.UpgradeKeeper.GetVersionMap(ctx)
 	require.NotEmpty(t, versionMap)
 }
diff --git a/x/upgrade/spec/01_concepts.md b/x/upgrade/spec/01_concepts.md
index ec8344352be4..3985631303eb 100644
--- a/x/upgrade/spec/01_concepts.md
+++ b/x/upgrade/spec/01_concepts.md
@@ -47,7 +47,7 @@ and not defined on a per-module basis. Registering a `Handler` is done via
 `Keeper#SetUpgradeHandler` in the application.
 
 ```go
-type UpgradeHandler func(Context, Plan, VersionMap) error
+type UpgradeHandler func(Context, Plan, VersionMap) (VersionMap, error)
 ```
 
 During each `EndBlock` execution, the `x/upgrade` module checks if there exists a
diff --git a/x/upgrade/spec/02_state.md b/x/upgrade/spec/02_state.md
index 14ef7794c3ae..97692eb5c60c 100644
--- a/x/upgrade/spec/02_state.md
+++ b/x/upgrade/spec/02_state.md
@@ -13,7 +13,7 @@ by the corresponding module name of type `string`.
 
 - Plan: `0x0 -> Plan`
 - Done: `0x1 | byte(plan name)  -> BigEndian(Block Height)`
-- ConsensusVersion: `0x2 | byte(module name)  -> LittleEndian(Module Consensus Version)`
+- ConsensusVersion: `0x2 | byte(module name)  -> BigEndian(Module Consensus Version)`
 
 
 The `x/upgrade` module contains no genesis state.
\ No newline at end of file
diff --git a/x/upgrade/types/handler.go b/x/upgrade/types/handler.go
index 0fa9212cec93..0dccaefebde8 100644
--- a/x/upgrade/types/handler.go
+++ b/x/upgrade/types/handler.go
@@ -6,4 +6,4 @@ import (
 )
 
 // UpgradeHandler specifies the type of function that is called when an upgrade is applied
-type UpgradeHandler func(ctx sdk.Context, plan Plan, vm module.VersionMap) error
+type UpgradeHandler func(ctx sdk.Context, plan Plan, vm module.VersionMap) (module.VersionMap, error)

From e6424ceb8689adbf29b1222bd3a049ac38a0cafe Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Tue, 16 Mar 2021 10:49:07 -0700
Subject: [PATCH 18/26] upgrade keeper transition to Version Map

---
 simapp/app.go              |  3 +--
 x/upgrade/abci_test.go     |  2 +-
 x/upgrade/keeper/keeper.go | 12 ++++++------
 x/upgrade/module.go        | 10 +++++-----
 4 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/simapp/app.go b/simapp/app.go
index 7c86a3b107f8..849167087469 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -310,7 +310,7 @@ func NewSimApp(
 		slashing.NewAppModule(appCodec, app.SlashingKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper),
 		distr.NewAppModule(appCodec, app.DistrKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper),
 		staking.NewAppModule(appCodec, app.StakingKeeper, app.AccountKeeper, app.BankKeeper),
-		upgrade.NewAppModule(app.UpgradeKeeper),
+		upgrade.NewAppModule(&app.UpgradeKeeper),
 		evidence.NewAppModule(app.EvidenceKeeper),
 		params.NewAppModule(app.ParamsKeeper),
 		authz.NewAppModule(appCodec, app.AuthzKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
@@ -318,7 +318,6 @@ func NewSimApp(
 
 	// Pass the version map to the upgrade keeper
 	app.UpgradeKeeper.SetInitialVersionMap(app.mm.GetVersionMap())
-	app.mm.Modules[upgradetypes.ModuleName] = upgrade.NewAppModule(app.UpgradeKeeper)
 
 	// During begin block slashing happens after distr.BeginBlocker so that
 	// there is nothing left over in the validator fee pool, so as to keep the
diff --git a/x/upgrade/abci_test.go b/x/upgrade/abci_test.go
index 9735ae619638..226589e94e9e 100644
--- a/x/upgrade/abci_test.go
+++ b/x/upgrade/abci_test.go
@@ -54,7 +54,7 @@ func setupTest(height int64, skip map[int64]bool) TestSuite {
 	s.keeper = app.UpgradeKeeper
 	s.ctx = app.BaseApp.NewContext(false, tmproto.Header{Height: height, Time: time.Now()})
 
-	s.module = upgrade.NewAppModule(s.keeper)
+	s.module = upgrade.NewAppModule(&s.keeper)
 	s.querier = s.module.LegacyQuerierHandler(app.LegacyAmino())
 	s.handler = upgrade.NewSoftwareUpgradeProposalHandler(s.keeper)
 	return s
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index 571a8e902989..f619ddac0338 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -51,18 +51,19 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
-// SetInitialVersionMap sets an initial version map on the keeper
+// SetInitialVersionMap sets an initial version map on the keeper.
+// This must be set from app.go after module Manager is initialized.
 func (k *Keeper) SetInitialVersionMap(vm module.VersionMap) {
 	k.versionMap = vm
 }
 
-// SetCurrentConsensusVersions saves the consensus versions.
-// If versionMap is not set from app.go, consensus versions will NOT be saved
-// to state.
+// SetCurrentConsensusVersions saves the current version map
+// to state from InitGenesis
 func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) {
 	k.setConsensusVersions(ctx, k.versionMap)
 }
 
+// setConsensusVersions saves a given version map to state
 func (k Keeper) setConsensusVersions(ctx sdk.Context, vm module.VersionMap) {
 	if len(vm) > 0 {
 		store := ctx.KVStore(k.storeKey)
@@ -239,8 +240,7 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic(err)
 	}
 
-	k.versionMap = updatedVM
-	k.SetCurrentConsensusVersions(ctx)
+	k.setConsensusVersions(ctx, updatedVM)
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
 	// This will prevent resubmission of upgrade msg after upgrade is already completed.
diff --git a/x/upgrade/module.go b/x/upgrade/module.go
index be5baec3d496..ce07fc68a69d 100644
--- a/x/upgrade/module.go
+++ b/x/upgrade/module.go
@@ -69,11 +69,11 @@ func (b AppModuleBasic) RegisterInterfaces(registry codectypes.InterfaceRegistry
 // AppModule implements the sdk.AppModule interface
 type AppModule struct {
 	AppModuleBasic
-	keeper keeper.Keeper
+	keeper *keeper.Keeper
 }
 
 // NewAppModule creates a new AppModule object
-func NewAppModule(keeper keeper.Keeper) AppModule {
+func NewAppModule(keeper *keeper.Keeper) AppModule {
 	return AppModule{
 		AppModuleBasic: AppModuleBasic{},
 		keeper:         keeper,
@@ -91,7 +91,7 @@ func (AppModule) QuerierRoute() string { return types.QuerierKey }
 
 // LegacyQuerierHandler registers a query handler to respond to the module-specific queries
 func (am AppModule) LegacyQuerierHandler(legacyQuerierCdc *codec.LegacyAmino) sdk.Querier {
-	return keeper.NewQuerier(am.keeper, legacyQuerierCdc)
+	return keeper.NewQuerier(*am.keeper, legacyQuerierCdc)
 }
 
 // RegisterServices registers a GRPC query service to respond to the
@@ -102,7 +102,7 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 
 // InitGenesis is ignored, no sense in serializing future upgrades
 func (am AppModule) InitGenesis(ctx sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
-	InitChainer(am.keeper, ctx)
+	InitChainer(*am.keeper, ctx)
 	return []abci.ValidatorUpdate{}
 }
 
@@ -128,7 +128,7 @@ func (AppModule) ConsensusVersion() uint64 { return 1 }
 //
 // CONTRACT: this is registered in BeginBlocker *before* all other modules' BeginBlock functions
 func (am AppModule) BeginBlock(ctx sdk.Context, req abci.RequestBeginBlock) {
-	BeginBlocker(am.keeper, ctx, req)
+	BeginBlocker(*am.keeper, ctx, req)
 }
 
 // EndBlock does nothing

From 799b646a97f46c6e2d1c8fbf727806b11c91d866 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Tue, 16 Mar 2021 11:41:47 -0700
Subject: [PATCH 19/26] cleanup, added versionmap return to RunMigrations

---
 CHANGELOG.md                                         |  2 +-
 .../adr-041-in-place-store-migrations.md             |  8 ++++----
 docs/run-node/cosmovisor.md                          |  2 +-
 simapp/app.go                                        |  2 +-
 simapp/app_test.go                                   |  2 +-
 types/module/module.go                               | 12 +++++++-----
 x/upgrade/keeper/keeper.go                           |  1 -
 x/upgrade/spec/02_state.md                           |  2 +-
 8 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 90795db4cc41..c2c9edf2f0b1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -61,7 +61,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [\#8629](https://github.com/cosmos/cosmos-sdk/pull/8629) Deprecated `SetFullFundraiserPath` from `Config` in favor of `SetPurpose` and `SetCoinType`.
 * (x/upgrade) [\#8673](https://github.com/cosmos/cosmos-sdk/pull/8673) Remove IBC logic from x/upgrade. Deprecates IBC fields in an Upgrade Plan. IBC upgrade logic moved to 02-client and an IBC UpgradeProposal is added.
 * (x/bank) [\#8517](https://github.com/cosmos/cosmos-sdk/pull/8517) `SupplyI` interface and `Supply` are removed and uses `sdk.Coins` for supply tracking
-* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) `UpgradeHandler` includes a new argument `VersionManager` which retreives module consensus versions from upgrade's new store.
+* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) `UpgradeHandler` includes a new argument `VersionMap` which helps facilitate in-place migrations.
 * (x/auth) [\#8129](https://github.com/cosmos/cosmos-sdk/pull/8828) Updated `SigVerifiableTx.GetPubKeys` method signature to return error.
 
 
diff --git a/docs/architecture/adr-041-in-place-store-migrations.md b/docs/architecture/adr-041-in-place-store-migrations.md
index b10c3a6982a9..474996eb201d 100644
--- a/docs/architecture/adr-041-in-place-store-migrations.md
+++ b/docs/architecture/adr-041-in-place-store-migrations.md
@@ -85,23 +85,23 @@ We introduce a new prefix store in `x/upgrade`'s store. This store will track ea
 0x2 | {bytes(module_name)} => BigEndian(module_consensus_version)
 ```
 
-We add a new private field `versionMap` of type `VersionMap` to `x/upgrade`'s keeper, where `VersionManager` is:
+We add a new private field `versionMap` of type `VersionMap` to `x/upgrade`'s keeper, where `VersionMap` is:
 
 ```go
 // Map of module name => new module Consensus Version.
 type VersionMap map[string]uint64
 ```
 
-This `versionMap` field can be modified via the `SetInitialVersionMap` field, and will allow the upgrade keeper to know the current versions of loaded modules. `SetInitialVersionMap` MUST be called as early as possible in the app initialization; in the SDK's `simapp`, it is called in the `NewSimApp` constructor function.
+This `versionMap` field can be modified via the `SetInitialVersionMap` field, and will allow the upgrade keeper to know the current versions of loaded modules. `SetInitialVersionMap` MUST be called AFTER the `module.NewManager()` is called; in the SDK's `simapp`, it is called in the `NewSimApp` constructor function.
 
-The UpgradeHandler signature needs to be updated to take a `VersionMap`, as well as return a new `VersionMap` and an error:
+The UpgradeHandler signature needs to be updated to take a `VersionMap`, as well as return an upgraded `VersionMap` and an error:
 
 ```diff
 - type UpgradeHandler func(ctx sdk.Context, plan Plan)
 + type UpgradeHandler func(ctx sdk.Context, plan Plan, versionMap VersionMap) (VersionMap, error)
 ```
 
-To apply an upgrade, we query the `VersionMap` from the `x/upgrade` store and pass it into the handler. The handler runs the actual migration functions (see next section), and if successful, returns the updated ConsensusVersions of the modules to be stored in state.
+To apply an upgrade, we query the `VersionMap` from the `x/upgrade` store and pass it into the handler. The handler runs the actual migration functions (see next section), and if successful, returns an updated `VersionMap` to be stored in state.
 
 ```diff
 func (k UpgradeKeeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
diff --git a/docs/run-node/cosmovisor.md b/docs/run-node/cosmovisor.md
index cf7575b9a31f..99cb81e0cc45 100644
--- a/docs/run-node/cosmovisor.md
+++ b/docs/run-node/cosmovisor.md
@@ -147,7 +147,7 @@ In `simapp/app.go`, find the line containing the upgrade Keeper initialisation,
 After that line, add the following snippet:
 
  ```
- app.UpgradeKeeper.SetUpgradeHandler("test1", func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap) {
+ app.UpgradeKeeper.SetUpgradeHandler("test1", func(ctx sdk.Context, plan upgradetypes.Plan) {
 		// Add some coins to a random account
 		addr, err := sdk.AccAddressFromBech32("cosmos18cgkqduwuh253twzmhedesw3l7v3fm37sppt58")
 		if err != nil {
diff --git a/simapp/app.go b/simapp/app.go
index 849167087469..d26f288d1d98 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -541,7 +541,7 @@ func (app *SimApp) RegisterTendermintService(clientCtx client.Context) {
 //   app.UpgradeKeeper.SetUpgradeHandler("store-migration", func(ctx sdk.Context, plan upgradetypes.Plan, versionMap module.VersionMap) {
 //       return app.RunMigrations(ctx, versionMap)
 //   })
-func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.VersionMap) error {
+func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.VersionMap) (module.VersionMap, error) {
 	return app.mm.RunMigrations(ctx, app.configurator, migrateFromVersions)
 }
 
diff --git a/simapp/app_test.go b/simapp/app_test.go
index dc1bdaa43603..2048cd683f8b 100644
--- a/simapp/app_test.go
+++ b/simapp/app_test.go
@@ -160,7 +160,7 @@ func TestRunMigrations(t *testing.T) {
 			// Run migrations only for bank. That's why we put the initial
 			// version for bank as 1, and for all other modules, we put as
 			// their latest ConsensusVersion.
-			err = app.RunMigrations(
+			_, err = app.RunMigrations(
 				app.NewContext(true, tmproto.Header{Height: app.LastBlockHeight()}),
 				module.VersionMap{
 					"bank":         1,
diff --git a/types/module/module.go b/types/module/module.go
index 6e3b1d670db2..ba6a6a47a6ce 100644
--- a/types/module/module.go
+++ b/types/module/module.go
@@ -338,20 +338,22 @@ type MigrationHandler func(sdk.Context) error
 type VersionMap map[string]uint64
 
 // RunMigrations performs in-place store migrations for all modules.
-func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, vm VersionMap) error {
+func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, fromVM VersionMap) (VersionMap, error) {
 	c, ok := cfg.(configurator)
 	if !ok {
-		return sdkerrors.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", configurator{}, cfg)
+		return nil, sdkerrors.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", configurator{}, cfg)
 	}
 
+	updatedVM := make(VersionMap)
 	for moduleName, module := range m.Modules {
-		err := c.runModuleMigrations(ctx, moduleName, vm[moduleName], module.ConsensusVersion())
+		err := c.runModuleMigrations(ctx, moduleName, fromVM[moduleName], module.ConsensusVersion())
+		updatedVM[moduleName] = module.ConsensusVersion()
 		if err != nil {
-			return err
+			return nil, err
 		}
 	}
 
-	return nil
+	return updatedVM, nil
 }
 
 // BeginBlock performs begin block functionality for all modules. It creates a
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index f619ddac0338..342aa13520a3 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -40,7 +40,6 @@ func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey sdk.StoreKey, cdc cod
 		storeKey:           storeKey,
 		cdc:                cdc,
 		upgradeHandlers:    map[string]types.UpgradeHandler{},
-		versionMap:         module.VersionMap{},
 	}
 }
 
diff --git a/x/upgrade/spec/02_state.md b/x/upgrade/spec/02_state.md
index 97692eb5c60c..05ad215a4054 100644
--- a/x/upgrade/spec/02_state.md
+++ b/x/upgrade/spec/02_state.md
@@ -16,4 +16,4 @@ by the corresponding module name of type `string`.
 - ConsensusVersion: `0x2 | byte(module name)  -> BigEndian(Module Consensus Version)`
 
 
-The `x/upgrade` module contains no genesis state.
\ No newline at end of file
+The `x/upgrade` module contains no genesis state.

From 0f9447cc02a0b17325aaa238de5f8ad4f808f8e8 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Tue, 16 Mar 2021 11:44:26 -0700
Subject: [PATCH 20/26] quick fix

---
 docs/architecture/adr-041-in-place-store-migrations.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/architecture/adr-041-in-place-store-migrations.md b/docs/architecture/adr-041-in-place-store-migrations.md
index 474996eb201d..9e14932d5603 100644
--- a/docs/architecture/adr-041-in-place-store-migrations.md
+++ b/docs/architecture/adr-041-in-place-store-migrations.md
@@ -133,8 +133,8 @@ If a required migration is missing (e.g. if it has not been registered in the `C
 In practice, the `RunMigrations` method should be called from inside an `UpgradeHandler`.
 
 ```go
-app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, versionMap module.VersionMap)  (module.VersionMap, error) {
-    return app.mm.RunMigrations(ctx, versionMap)
+app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap)  (module.VersionMap, error) {
+    return app.mm.RunMigrations(ctx, vm)
 })
 ```
 

From b839a59ad737617a8481be80858f413df37c3fe5 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+technicallyty@users.noreply.github.com>
Date: Wed, 17 Mar 2021 10:30:23 -0700
Subject: [PATCH 21/26] Update
 docs/architecture/adr-041-in-place-store-migrations.md

Co-authored-by: Amaury <1293565+amaurym@users.noreply.github.com>
---
 docs/architecture/adr-041-in-place-store-migrations.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/architecture/adr-041-in-place-store-migrations.md b/docs/architecture/adr-041-in-place-store-migrations.md
index 9e14932d5603..ecd41e480bfc 100644
--- a/docs/architecture/adr-041-in-place-store-migrations.md
+++ b/docs/architecture/adr-041-in-place-store-migrations.md
@@ -92,7 +92,7 @@ We add a new private field `versionMap` of type `VersionMap` to `x/upgrade`'s ke
 type VersionMap map[string]uint64
 ```
 
-This `versionMap` field can be modified via the `SetInitialVersionMap` field, and will allow the upgrade keeper to know the current versions of loaded modules. `SetInitialVersionMap` MUST be called AFTER the `module.NewManager()` is called; in the SDK's `simapp`, it is called in the `NewSimApp` constructor function.
+This `versionMap` field can be modified via the `SetVersionMap` method, and will allow the upgrade keeper to store in state the current versions of loaded modules. To initialize the x/upgrade store correctly, `SetVersionMap` MUST be called during the app's `InitChain`.
 
 The UpgradeHandler signature needs to be updated to take a `VersionMap`, as well as return an upgraded `VersionMap` and an error:
 

From 75f8ae5dc0f3ac8011709ea29573b24debc08f13 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Wed, 17 Mar 2021 11:02:43 -0700
Subject: [PATCH 22/26] remove support for versionmap field on upgrade keeper

---
 CHANGELOG.md                    |  1 +
 simapp/app.go                   |  6 ++----
 simapp/app_test.go              | 24 ++++++++++++++++++++++++
 x/upgrade/abci.go               |  4 ----
 x/upgrade/abci_test.go          |  2 +-
 x/upgrade/keeper/keeper.go      | 19 +++----------------
 x/upgrade/keeper/keeper_test.go |  7 ++-----
 x/upgrade/module.go             |  9 ++++-----
 x/upgrade/module_test.go        | 25 -------------------------
 9 files changed, 37 insertions(+), 60 deletions(-)
 delete mode 100644 x/upgrade/module_test.go

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c2c9edf2f0b1..551ab5614ebf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -80,6 +80,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (x/auth) [\#8522](https://github.com/cosmos/cosmos-sdk/pull/8522) Allow to query all stored accounts
 * (crypto/types) [\#8600](https://github.com/cosmos/cosmos-sdk/pull/8600) `CompactBitArray`: optimize the `NumTrueBitsBefore` method and add an `Equal` method.
 * (grpc) [\#8815](https://github.com/cosmos/cosmos-sdk/pull/8815) Add orderBy parameter to `TxsByEvents` endpoint.
+* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) Add tracking module versions as per ADR-041
 
 ### Bug Fixes
 
diff --git a/simapp/app.go b/simapp/app.go
index d26f288d1d98..c6032e710a62 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -310,15 +310,12 @@ func NewSimApp(
 		slashing.NewAppModule(appCodec, app.SlashingKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper),
 		distr.NewAppModule(appCodec, app.DistrKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper),
 		staking.NewAppModule(appCodec, app.StakingKeeper, app.AccountKeeper, app.BankKeeper),
-		upgrade.NewAppModule(&app.UpgradeKeeper),
+		upgrade.NewAppModule(app.UpgradeKeeper),
 		evidence.NewAppModule(app.EvidenceKeeper),
 		params.NewAppModule(app.ParamsKeeper),
 		authz.NewAppModule(appCodec, app.AuthzKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
 	)
 
-	// Pass the version map to the upgrade keeper
-	app.UpgradeKeeper.SetInitialVersionMap(app.mm.GetVersionMap())
-
 	// During begin block slashing happens after distr.BeginBlocker so that
 	// there is nothing left over in the validator fee pool, so as to keep the
 	// CanWithdrawInvariant invariant.
@@ -424,6 +421,7 @@ func (app *SimApp) InitChainer(ctx sdk.Context, req abci.RequestInitChain) abci.
 	if err := json.Unmarshal(req.AppStateBytes, &genesisState); err != nil {
 		panic(err)
 	}
+	app.UpgradeKeeper.SetConsensusVersions(ctx, app.mm.GetVersionMap())
 	return app.mm.InitGenesis(ctx, app.appCodec, genesisState)
 }
 
diff --git a/simapp/app_test.go b/simapp/app_test.go
index 2048cd683f8b..2915338d821b 100644
--- a/simapp/app_test.go
+++ b/simapp/app_test.go
@@ -190,3 +190,27 @@ func TestRunMigrations(t *testing.T) {
 		})
 	}
 }
+
+func TestUpgradeStateOnGenesis(t *testing.T) {
+	encCfg := MakeTestEncodingConfig()
+	db := dbm.NewMemDB()
+	app := NewSimApp(log.NewTMLogger(log.NewSyncWriter(os.Stdout)), db, nil, true, map[int64]bool{}, DefaultNodeHome, 0, encCfg, EmptyAppOptions{})
+	genesisState := NewDefaultGenesisState(encCfg.Marshaler)
+	stateBytes, err := json.MarshalIndent(genesisState, "", "  ")
+	require.NoError(t, err)
+
+	// Initialize the chain
+	app.InitChain(
+		abci.RequestInitChain{
+			Validators:    []abci.ValidatorUpdate{},
+			AppStateBytes: stateBytes,
+		},
+	)
+
+	// make sure the upgrade keeper has version map in state
+	ctx := app.NewContext(false, tmproto.Header{})
+	vm := app.UpgradeKeeper.GetVersionMap(ctx)
+	for v, i := range app.mm.Modules {
+		require.Equal(t, vm[v], i.ConsensusVersion())
+	}
+}
diff --git a/x/upgrade/abci.go b/x/upgrade/abci.go
index 6e85abb1e7b5..d346decb023c 100644
--- a/x/upgrade/abci.go
+++ b/x/upgrade/abci.go
@@ -69,10 +69,6 @@ func BeginBlocker(k keeper.Keeper, ctx sdk.Context, _ abci.RequestBeginBlock) {
 	}
 }
 
-func InitChainer(k keeper.Keeper, ctx sdk.Context) {
-	k.SetCurrentConsensusVersions(ctx)
-}
-
 // BuildUpgradeNeededMsg prints the message that notifies that an upgrade is needed.
 func BuildUpgradeNeededMsg(plan types.Plan) string {
 	return fmt.Sprintf("UPGRADE \"%s\" NEEDED at %s: %s", plan.Name, plan.DueAt(), plan.Info)
diff --git a/x/upgrade/abci_test.go b/x/upgrade/abci_test.go
index 226589e94e9e..9735ae619638 100644
--- a/x/upgrade/abci_test.go
+++ b/x/upgrade/abci_test.go
@@ -54,7 +54,7 @@ func setupTest(height int64, skip map[int64]bool) TestSuite {
 	s.keeper = app.UpgradeKeeper
 	s.ctx = app.BaseApp.NewContext(false, tmproto.Header{Height: height, Time: time.Now()})
 
-	s.module = upgrade.NewAppModule(&s.keeper)
+	s.module = upgrade.NewAppModule(s.keeper)
 	s.querier = s.module.LegacyQuerierHandler(app.LegacyAmino())
 	s.handler = upgrade.NewSoftwareUpgradeProposalHandler(s.keeper)
 	return s
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index 342aa13520a3..3e7bf7d81519 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -29,7 +29,6 @@ type Keeper struct {
 	storeKey           sdk.StoreKey
 	cdc                codec.BinaryMarshaler
 	upgradeHandlers    map[string]types.UpgradeHandler
-	versionMap         module.VersionMap
 }
 
 // NewKeeper constructs an upgrade Keeper
@@ -50,20 +49,8 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
-// SetInitialVersionMap sets an initial version map on the keeper.
-// This must be set from app.go after module Manager is initialized.
-func (k *Keeper) SetInitialVersionMap(vm module.VersionMap) {
-	k.versionMap = vm
-}
-
-// SetCurrentConsensusVersions saves the current version map
-// to state from InitGenesis
-func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) {
-	k.setConsensusVersions(ctx, k.versionMap)
-}
-
-// setConsensusVersions saves a given version map to state
-func (k Keeper) setConsensusVersions(ctx sdk.Context, vm module.VersionMap) {
+// SetConsensusVersions saves a given version map to state
+func (k Keeper) SetConsensusVersions(ctx sdk.Context, vm module.VersionMap) {
 	if len(vm) > 0 {
 		store := ctx.KVStore(k.storeKey)
 		versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
@@ -239,7 +226,7 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic(err)
 	}
 
-	k.setConsensusVersions(ctx, updatedVM)
+	k.SetConsensusVersions(ctx, updatedVM)
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
 	// This will prevent resubmission of upgrade msg after upgrade is already completed.
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index ed4559a029be..349e6f889d64 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -12,7 +12,6 @@ import (
 	store "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
-	"github.com/cosmos/cosmos-sdk/x/bank"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/keeper"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
@@ -136,7 +135,6 @@ func (s *KeeperTestSuite) TestScheduleUpgrade() {
 		s.Run(tc.name, func() {
 			// reset suite
 			s.SetupTest()
-			s.app.UpgradeKeeper.SetInitialVersionMap(module.VersionMap{})
 			// setup test case
 			tc.setup()
 
@@ -199,8 +197,7 @@ func (s *KeeperTestSuite) TestSetUpgradedClient() {
 // an upgrade.
 func (s *KeeperTestSuite) TestMigrations() {
 	initialVM := module.VersionMap{"bank": uint64(1)}
-	s.app.UpgradeKeeper.SetInitialVersionMap(initialVM)
-	s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx)
+	s.app.UpgradeKeeper.SetConsensusVersions(s.ctx, initialVM)
 	vmBefore := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
 	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		// simulate upgrading the bank module
@@ -215,7 +212,7 @@ func (s *KeeperTestSuite) TestMigrations() {
 
 	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, dummyPlan)
 	vm := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
-	s.Require().Equal(bank.AppModule{}.ConsensusVersion(), vm["bank"])
+	s.Require().Equal(vmBefore["bank"]+1, vm["bank"])
 	s.Require().Greater(vm["bank"], vmBefore["bank"])
 }
 
diff --git a/x/upgrade/module.go b/x/upgrade/module.go
index ce07fc68a69d..154f006784a7 100644
--- a/x/upgrade/module.go
+++ b/x/upgrade/module.go
@@ -69,11 +69,11 @@ func (b AppModuleBasic) RegisterInterfaces(registry codectypes.InterfaceRegistry
 // AppModule implements the sdk.AppModule interface
 type AppModule struct {
 	AppModuleBasic
-	keeper *keeper.Keeper
+	keeper keeper.Keeper
 }
 
 // NewAppModule creates a new AppModule object
-func NewAppModule(keeper *keeper.Keeper) AppModule {
+func NewAppModule(keeper keeper.Keeper) AppModule {
 	return AppModule{
 		AppModuleBasic: AppModuleBasic{},
 		keeper:         keeper,
@@ -91,7 +91,7 @@ func (AppModule) QuerierRoute() string { return types.QuerierKey }
 
 // LegacyQuerierHandler registers a query handler to respond to the module-specific queries
 func (am AppModule) LegacyQuerierHandler(legacyQuerierCdc *codec.LegacyAmino) sdk.Querier {
-	return keeper.NewQuerier(*am.keeper, legacyQuerierCdc)
+	return keeper.NewQuerier(am.keeper, legacyQuerierCdc)
 }
 
 // RegisterServices registers a GRPC query service to respond to the
@@ -102,7 +102,6 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 
 // InitGenesis is ignored, no sense in serializing future upgrades
 func (am AppModule) InitGenesis(ctx sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
-	InitChainer(*am.keeper, ctx)
 	return []abci.ValidatorUpdate{}
 }
 
@@ -128,7 +127,7 @@ func (AppModule) ConsensusVersion() uint64 { return 1 }
 //
 // CONTRACT: this is registered in BeginBlocker *before* all other modules' BeginBlock functions
 func (am AppModule) BeginBlock(ctx sdk.Context, req abci.RequestBeginBlock) {
-	BeginBlocker(*am.keeper, ctx, req)
+	BeginBlocker(am.keeper, ctx, req)
 }
 
 // EndBlock does nothing
diff --git a/x/upgrade/module_test.go b/x/upgrade/module_test.go
deleted file mode 100644
index 5b4bf1e1a631..000000000000
--- a/x/upgrade/module_test.go
+++ /dev/null
@@ -1,25 +0,0 @@
-package upgrade_test
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/require"
-	abcitypes "github.com/tendermint/tendermint/abci/types"
-	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
-
-	"github.com/cosmos/cosmos-sdk/simapp"
-)
-
-func TestInitChainer(t *testing.T) {
-	app := simapp.Setup(false)
-	ctx := app.BaseApp.NewContext(false, tmproto.Header{})
-	app.InitChain(
-		abcitypes.RequestInitChain{
-			AppStateBytes: []byte("{}"),
-			ChainId:       "test-chain-id",
-		},
-	)
-
-	versionMap := app.UpgradeKeeper.GetVersionMap(ctx)
-	require.NotEmpty(t, versionMap)
-}

From 22ab55e47712b484a91dbe8d4532b803d216e984 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Wed, 17 Mar 2021 11:13:38 -0700
Subject: [PATCH 23/26] cleanup

---
 .../adr-041-in-place-store-migrations.md           | 14 +++-----------
 simapp/app.go                                      |  8 ++++----
 x/upgrade/keeper/keeper.go                         |  8 ++++----
 x/upgrade/keeper/keeper_test.go                    |  3 ++-
 x/upgrade/module.go                                |  2 +-
 x/upgrade/spec/02_state.md                         |  2 +-
 6 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/docs/architecture/adr-041-in-place-store-migrations.md b/docs/architecture/adr-041-in-place-store-migrations.md
index 4b2129efaec3..d6f3496ff585 100644
--- a/docs/architecture/adr-041-in-place-store-migrations.md
+++ b/docs/architecture/adr-041-in-place-store-migrations.md
@@ -84,15 +84,7 @@ We introduce a new prefix store in `x/upgrade`'s store. This store will track ea
 ```
 0x2 | {bytes(module_name)} => BigEndian(module_consensus_version)
 ```
-
-We add a new private field `versionMap` of type `VersionMap` to `x/upgrade`'s keeper, where `VersionMap` is:
-
-```go
-// Map of module name => new module Consensus Version.
-type VersionMap map[string]uint64
-```
-
-This `versionMap` field can be modified via the `SetInitialVersionMap` field, and will allow the upgrade keeper to know the current versions of loaded modules. `SetInitialVersionMap` MUST be called AFTER the `module.NewManager()` is called; in the SDK's `simapp`, it is called in the `NewSimApp` constructor function.
+The initial state of the store is set from `app.go`'s `InitChainer` method.
 
 The UpgradeHandler signature needs to be updated to take a `VersionMap`, as well as return an upgraded `VersionMap` and an error:
 
@@ -107,13 +99,13 @@ To apply an upgrade, we query the `VersionMap` from the `x/upgrade` store and pa
 func (k UpgradeKeeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
     // --snip--
 -   handler(ctx, plan)
-+   updatedVM, err := handler(ctx, plan, k.GetConsensusVersions(ctx)) // k.GetConsensusVersions() fetches the VersionMap stored in state.
++   updatedVM, err := handler(ctx, plan, k.GetVersionMap(ctx)) // k.GetVersionMap() fetches the VersionMap stored in state.
 +   if err != nil {
 +       return err
 +   }
 +
 +   // Set the updated consensus versions to state
-+   k.setConsensusVersions(ctx, updatedVM)
++   k.SetVersionMap(ctx, updatedVM)
 }
 ```
 
diff --git a/simapp/app.go b/simapp/app.go
index c6032e710a62..b34eb8a93550 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -332,7 +332,7 @@ func NewSimApp(
 	// so that other modules that want to create or claim capabilities afterwards in InitChain
 	// can do so safely.
 	app.mm.SetOrderInitGenesis(
-		upgradetypes.ModuleName, capabilitytypes.ModuleName, authtypes.ModuleName, banktypes.ModuleName, distrtypes.ModuleName, stakingtypes.ModuleName,
+		capabilitytypes.ModuleName, authtypes.ModuleName, banktypes.ModuleName, distrtypes.ModuleName, stakingtypes.ModuleName,
 		slashingtypes.ModuleName, govtypes.ModuleName, minttypes.ModuleName, crisistypes.ModuleName,
 		genutiltypes.ModuleName, evidencetypes.ModuleName, authztypes.ModuleName,
 		feegranttypes.ModuleName,
@@ -421,7 +421,7 @@ func (app *SimApp) InitChainer(ctx sdk.Context, req abci.RequestInitChain) abci.
 	if err := json.Unmarshal(req.AppStateBytes, &genesisState); err != nil {
 		panic(err)
 	}
-	app.UpgradeKeeper.SetConsensusVersions(ctx, app.mm.GetVersionMap())
+	app.UpgradeKeeper.SetVersionMap(ctx, app.mm.GetVersionMap())
 	return app.mm.InitGenesis(ctx, app.appCodec, genesisState)
 }
 
@@ -536,8 +536,8 @@ func (app *SimApp) RegisterTendermintService(clientCtx client.Context) {
 //
 // Example:
 //   cfg := module.NewConfigurator(...)
-//   app.UpgradeKeeper.SetUpgradeHandler("store-migration", func(ctx sdk.Context, plan upgradetypes.Plan, versionMap module.VersionMap) {
-//       return app.RunMigrations(ctx, versionMap)
+//   app.UpgradeKeeper.SetUpgradeHandler("store-migration", func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap) {
+//       return app.RunMigrations(ctx, vm)
 //   })
 func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.VersionMap) (module.VersionMap, error) {
 	return app.mm.RunMigrations(ctx, app.configurator, migrateFromVersions)
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index 3e7bf7d81519..70e87dff5229 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -49,8 +49,8 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
-// SetConsensusVersions saves a given version map to state
-func (k Keeper) SetConsensusVersions(ctx sdk.Context, vm module.VersionMap) {
+// SetVersionMap saves a given version map to state
+func (k Keeper) SetVersionMap(ctx sdk.Context, vm module.VersionMap) {
 	if len(vm) > 0 {
 		store := ctx.KVStore(k.storeKey)
 		versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
@@ -63,7 +63,7 @@ func (k Keeper) SetConsensusVersions(ctx sdk.Context, vm module.VersionMap) {
 	}
 }
 
-// GetConsensusVersions returns a map of key module name and value module consensus version
+// GetVersionMap returns a map of key module name and value module consensus version
 // as defined in ADR-041.
 func (k Keeper) GetVersionMap(ctx sdk.Context) module.VersionMap {
 	store := ctx.KVStore(k.storeKey)
@@ -226,7 +226,7 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic(err)
 	}
 
-	k.SetConsensusVersions(ctx, updatedVM)
+	k.SetVersionMap(ctx, updatedVM)
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
 	// This will prevent resubmission of upgrade msg after upgrade is already completed.
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index 349e6f889d64..5ff33faee2ac 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -135,6 +135,7 @@ func (s *KeeperTestSuite) TestScheduleUpgrade() {
 		s.Run(tc.name, func() {
 			// reset suite
 			s.SetupTest()
+
 			// setup test case
 			tc.setup()
 
@@ -197,7 +198,7 @@ func (s *KeeperTestSuite) TestSetUpgradedClient() {
 // an upgrade.
 func (s *KeeperTestSuite) TestMigrations() {
 	initialVM := module.VersionMap{"bank": uint64(1)}
-	s.app.UpgradeKeeper.SetConsensusVersions(s.ctx, initialVM)
+	s.app.UpgradeKeeper.SetVersionMap(s.ctx, initialVM)
 	vmBefore := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
 	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		// simulate upgrading the bank module
diff --git a/x/upgrade/module.go b/x/upgrade/module.go
index 154f006784a7..23311c22d5f3 100644
--- a/x/upgrade/module.go
+++ b/x/upgrade/module.go
@@ -101,7 +101,7 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 }
 
 // InitGenesis is ignored, no sense in serializing future upgrades
-func (am AppModule) InitGenesis(ctx sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
+func (am AppModule) InitGenesis(_ sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
 	return []abci.ValidatorUpdate{}
 }
 
diff --git a/x/upgrade/spec/02_state.md b/x/upgrade/spec/02_state.md
index 05ad215a4054..796b99686d20 100644
--- a/x/upgrade/spec/02_state.md
+++ b/x/upgrade/spec/02_state.md
@@ -8,7 +8,7 @@ The internal state of the `x/upgrade` module is relatively minimal and simple. T
 state contains the currently active upgrade `Plan` (if one exists) by key
 `0x0` and if a `Plan` is marked as "done" by key `0x1`. Additionally, the state
 contains the consensus versions of all app modules in the application. The versions 
-are stored as little endian `uint64`, and can be accessed with prefix `0x2` appended 
+are stored as big endian `uint64`, and can be accessed with prefix `0x2` appended 
 by the corresponding module name of type `string`. 
 
 - Plan: `0x0 -> Plan`

From 7f5b6aa7a16e3b7824a2310dc34216c6d2410ef2 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Thu, 18 Mar 2021 15:20:56 -0700
Subject: [PATCH 24/26] rename get/set version map keeper functions

---
 simapp/app.go                   |  2 +-
 simapp/app_test.go              |  2 +-
 x/upgrade/keeper/keeper.go      | 12 ++++++------
 x/upgrade/keeper/keeper_test.go |  6 +++---
 4 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/simapp/app.go b/simapp/app.go
index b34eb8a93550..97a71d4015e4 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -421,7 +421,7 @@ func (app *SimApp) InitChainer(ctx sdk.Context, req abci.RequestInitChain) abci.
 	if err := json.Unmarshal(req.AppStateBytes, &genesisState); err != nil {
 		panic(err)
 	}
-	app.UpgradeKeeper.SetVersionMap(ctx, app.mm.GetVersionMap())
+	app.UpgradeKeeper.SetModuleVersionMap(ctx, app.mm.GetVersionMap())
 	return app.mm.InitGenesis(ctx, app.appCodec, genesisState)
 }
 
diff --git a/simapp/app_test.go b/simapp/app_test.go
index 2915338d821b..c4aa9885b5d8 100644
--- a/simapp/app_test.go
+++ b/simapp/app_test.go
@@ -209,7 +209,7 @@ func TestUpgradeStateOnGenesis(t *testing.T) {
 
 	// make sure the upgrade keeper has version map in state
 	ctx := app.NewContext(false, tmproto.Header{})
-	vm := app.UpgradeKeeper.GetVersionMap(ctx)
+	vm := app.UpgradeKeeper.GetModuleVersionMap(ctx)
 	for v, i := range app.mm.Modules {
 		require.Equal(t, vm[v], i.ConsensusVersion())
 	}
diff --git a/x/upgrade/keeper/keeper.go b/x/upgrade/keeper/keeper.go
index 70e87dff5229..b171a219b59b 100644
--- a/x/upgrade/keeper/keeper.go
+++ b/x/upgrade/keeper/keeper.go
@@ -49,8 +49,8 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
-// SetVersionMap saves a given version map to state
-func (k Keeper) SetVersionMap(ctx sdk.Context, vm module.VersionMap) {
+// SetModuleVersionMap saves a given version map to state
+func (k Keeper) SetModuleVersionMap(ctx sdk.Context, vm module.VersionMap) {
 	if len(vm) > 0 {
 		store := ctx.KVStore(k.storeKey)
 		versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
@@ -63,9 +63,9 @@ func (k Keeper) SetVersionMap(ctx sdk.Context, vm module.VersionMap) {
 	}
 }
 
-// GetVersionMap returns a map of key module name and value module consensus version
+// GetModuleVersionMap returns a map of key module name and value module consensus version
 // as defined in ADR-041.
-func (k Keeper) GetVersionMap(ctx sdk.Context) module.VersionMap {
+func (k Keeper) GetModuleVersionMap(ctx sdk.Context) module.VersionMap {
 	store := ctx.KVStore(k.storeKey)
 	it := sdk.KVStorePrefixIterator(store, []byte{types.VersionMapByte})
 
@@ -221,12 +221,12 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic("ApplyUpgrade should never be called without first checking HasHandler")
 	}
 
-	updatedVM, err := handler(ctx, plan, k.GetVersionMap(ctx))
+	updatedVM, err := handler(ctx, plan, k.GetModuleVersionMap(ctx))
 	if err != nil {
 		panic(err)
 	}
 
-	k.SetVersionMap(ctx, updatedVM)
+	k.SetModuleVersionMap(ctx, updatedVM)
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
 	// This will prevent resubmission of upgrade msg after upgrade is already completed.
diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index 5ff33faee2ac..5f7d63beec44 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -198,8 +198,8 @@ func (s *KeeperTestSuite) TestSetUpgradedClient() {
 // an upgrade.
 func (s *KeeperTestSuite) TestMigrations() {
 	initialVM := module.VersionMap{"bank": uint64(1)}
-	s.app.UpgradeKeeper.SetVersionMap(s.ctx, initialVM)
-	vmBefore := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
+	s.app.UpgradeKeeper.SetModuleVersionMap(s.ctx, initialVM)
+	vmBefore := s.app.UpgradeKeeper.GetModuleVersionMap(s.ctx)
 	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		// simulate upgrading the bank module
 		vm["bank"] = vm["bank"] + 1
@@ -212,7 +212,7 @@ func (s *KeeperTestSuite) TestMigrations() {
 	}
 
 	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, dummyPlan)
-	vm := s.app.UpgradeKeeper.GetVersionMap(s.ctx)
+	vm := s.app.UpgradeKeeper.GetModuleVersionMap(s.ctx)
 	s.Require().Equal(vmBefore["bank"]+1, vm["bank"])
 	s.Require().Greater(vm["bank"], vmBefore["bank"])
 }

From 85474264631af363d1363a17c58d209d43fbd3d8 Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Thu, 18 Mar 2021 15:28:30 -0700
Subject: [PATCH 25/26] update adr doc to match name changes

---
 docs/architecture/adr-041-in-place-store-migrations.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/architecture/adr-041-in-place-store-migrations.md b/docs/architecture/adr-041-in-place-store-migrations.md
index d6f3496ff585..179219b94cc9 100644
--- a/docs/architecture/adr-041-in-place-store-migrations.md
+++ b/docs/architecture/adr-041-in-place-store-migrations.md
@@ -99,13 +99,13 @@ To apply an upgrade, we query the `VersionMap` from the `x/upgrade` store and pa
 func (k UpgradeKeeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
     // --snip--
 -   handler(ctx, plan)
-+   updatedVM, err := handler(ctx, plan, k.GetVersionMap(ctx)) // k.GetVersionMap() fetches the VersionMap stored in state.
++   updatedVM, err := handler(ctx, plan, k.GetModuleVersionMap(ctx)) // k.GetModuleVersionMap() fetches the VersionMap stored in state.
 +   if err != nil {
 +       return err
 +   }
 +
 +   // Set the updated consensus versions to state
-+   k.SetVersionMap(ctx, updatedVM)
++   k.SetModuleVersionMap(ctx, updatedVM)
 }
 ```
 

From c3072abf4bdc0b251f0fe6b7be956d566592f46a Mon Sep 17 00:00:00 2001
From: technicallyty <48813565+tytech3@users.noreply.github.com>
Date: Fri, 19 Mar 2021 06:59:39 -0700
Subject: [PATCH 26/26] remove redudant line

---
 x/upgrade/keeper/keeper_test.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/x/upgrade/keeper/keeper_test.go b/x/upgrade/keeper/keeper_test.go
index 5f7d63beec44..9a5c38f64974 100644
--- a/x/upgrade/keeper/keeper_test.go
+++ b/x/upgrade/keeper/keeper_test.go
@@ -214,7 +214,6 @@ func (s *KeeperTestSuite) TestMigrations() {
 	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, dummyPlan)
 	vm := s.app.UpgradeKeeper.GetModuleVersionMap(s.ctx)
 	s.Require().Equal(vmBefore["bank"]+1, vm["bank"])
-	s.Require().Greater(vm["bank"], vmBefore["bank"])
 }
 
 func TestKeeperTestSuite(t *testing.T) {