Use depinject to provide VersionMap (#14301)

Co-authored-by: Amaury M <1293565+amaurym@users.noreply.github.com>

runtime/module.go

@@ -33,6 +33,7 @@ func init() {
 			ProvideTransientStoreKey,
 			ProvideMemoryStoreKey,
 			ProvideDeliverTx,
+			ProvideVersionMap,
 		),
 		appmodule.Invoke(SetupAppBuilder),
 	)
@@ -141,3 +142,9 @@ func ProvideDeliverTx(appBuilder *AppBuilder) func(abci.RequestDeliverTx) abci.R
 		return appBuilder.app.BaseApp.DeliverTx(tx)
 	}
 }
+
+func ProvideVersionMap(appBuilder *AppBuilder) func() module.VersionMap {
+	return func() module.VersionMap {
+		return appBuilder.app.ModuleManager.GetVersionMap()
+	}
+}

simapp/app.go

@@ -339,7 +339,7 @@ func NewSimApp(
 	}
 	homePath := cast.ToString(appOpts.Get(flags.FlagHome))
 	// set the governance module account as the authority for conducting upgrades
-	app.UpgradeKeeper = upgradekeeper.NewKeeper(skipUpgradeHeights, keys[upgradetypes.StoreKey], appCodec, homePath, app.BaseApp, authtypes.NewModuleAddress(govtypes.ModuleName).String())
+	app.UpgradeKeeper = upgradekeeper.NewKeeper(skipUpgradeHeights, keys[upgradetypes.StoreKey], appCodec, homePath, app.BaseApp, authtypes.NewModuleAddress(govtypes.ModuleName).String(), app.ModuleManager.GetVersionMap())
 
 	// Register the proposal types
 	// Deprecated: Avoid adding new handlers, instead use the new proposal flow

x/upgrade/keeper/keeper.go

@@ -35,7 +35,7 @@ type Keeper struct {
 	versionSetter      xp.ProtocolVersionSetter        // implements setting the protocol version field on BaseApp
 	downgradeVerified  bool                            // tells if we've already sanity checked that this binary version isn't being used against an old state.
 	authority          string                          // the address capable of executing and cancelling an upgrade. Usually the gov module account
-	initVersionMap     module.VersionMap               // the module version map at init genesis
+	appVersionMap      module.VersionMap               // the current module version map of the app
 }
 
 // NewKeeper constructs an upgrade Keeper which requires the following arguments:
@@ -44,7 +44,7 @@ type Keeper struct {
 // cdc - the app-wide binary codec
 // homePath - root directory of the application's config
 // vs - the interface implemented by baseapp which allows setting baseapp's protocol version field
-func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey storetypes.StoreKey, cdc codec.BinaryCodec, homePath string, vs xp.ProtocolVersionSetter, authority string) *Keeper {
+func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey storetypes.StoreKey, cdc codec.BinaryCodec, homePath string, vs xp.ProtocolVersionSetter, authority string, versionMap module.VersionMap) *Keeper {
 	return &Keeper{
 		homePath:           homePath,
 		skipUpgradeHeights: skipUpgradeHeights,
@@ -53,6 +53,7 @@ func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey storetypes.StoreKey,
 		upgradeHandlers:    map[string]types.UpgradeHandler{},
 		versionSetter:      vs,
 		authority:          authority,
+		appVersionMap:      versionMap,
 	}
 }
 
@@ -66,16 +67,10 @@ func (k *Keeper) GetVersionSetter() xp.ProtocolVersionSetter {
 	return k.versionSetter
 }
 
-// SetInitVersionMap sets the initial version map.
-// This is only used in app wiring and should not be used in any other context.
-func (k *Keeper) SetInitVersionMap(vm module.VersionMap) {
-	k.initVersionMap = vm
-}
-
-// GetInitVersionMap gets the initial version map
+// GetAppVersionMap gets the current version map of the app
 // This is only used in upgrade InitGenesis and should not be used in any other context.
-func (k *Keeper) GetInitVersionMap() module.VersionMap {
-	return k.initVersionMap
+func (k *Keeper) GetAppVersionMap() module.VersionMap {
+	return k.appVersionMap
 }
 
 // SetUpgradeHandler sets an UpgradeHandler for the upgrade specified by name. This handler will be called when the upgrade

x/upgrade/module.go

@@ -120,18 +120,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.JSONCodec, _ json.RawMessage) []abci.ValidatorUpdate {
-	// set version map automatically if available
-	if versionMap := am.keeper.GetInitVersionMap(); versionMap != nil {
-		// chains can still use a custom init chainer for setting the version map
-		// this means that we need to combine the manually wired modules version map with app wiring enabled modules version map
-		for name, version := range am.keeper.GetModuleVersionMap(ctx) {
-			if _, ok := versionMap[name]; !ok {
-				versionMap[name] = version
-			}
-		}
-
-		am.keeper.SetModuleVersionMap(ctx, versionMap)
-	}
+	am.keeper.SetModuleVersionMap(ctx, am.keeper.GetAppVersionMap())
 
 	return []abci.ValidatorUpdate{}
 }
@@ -168,16 +157,16 @@ func (am AppModule) BeginBlock(ctx sdk.Context, req abci.RequestBeginBlock) {
 func init() {
 	appmodule.Register(&modulev1.Module{},
 		appmodule.Provide(ProvideModule),
-		appmodule.Invoke(PopulateVersionMap),
 	)
 }
 
 type UpgradeInputs struct {
 	depinject.In
 
-	Config *modulev1.Module
-	Key    *store.KVStoreKey
-	Cdc    codec.Codec
+	Config     *modulev1.Module
+	Key        *store.KVStoreKey
+	Cdc        codec.Codec
+	VersionMap module.VersionMap
 
 	AppOpts servertypes.AppOptions `optional:"true"`
 }
@@ -212,7 +201,7 @@ func ProvideModule(in UpgradeInputs) UpgradeOutputs {
 	}
 
 	// set the governance module account as the authority for conducting upgrades
-	k := keeper.NewKeeper(skipUpgradeHeights, in.Key, in.Cdc, homePath, nil, authority.String())
+	k := keeper.NewKeeper(skipUpgradeHeights, in.Key, in.Cdc, homePath, nil, authority.String(), in.VersionMap)
 	baseappOpt := func(app *baseapp.BaseApp) {
 		k.SetVersionSetter(app)
 	}
@@ -221,11 +210,3 @@ func ProvideModule(in UpgradeInputs) UpgradeOutputs {
 
 	return UpgradeOutputs{UpgradeKeeper: k, Module: m, GovHandler: gh, BaseAppOption: baseappOpt}
 }
-
-func PopulateVersionMap(upgradeKeeper *keeper.Keeper, modules map[string]appmodule.AppModule) {
-	if upgradeKeeper == nil {
-		return
-	}
-
-	upgradeKeeper.SetInitVersionMap(module.NewManagerFromMap(modules).GetVersionMap())
-}