From ebb46a53a682f644e82f21aae1f03f278f8363fa Mon Sep 17 00:00:00 2001
From: Dreamer <745124335@qq.com>
Date: Tue, 30 May 2023 15:27:39 +0800
Subject: [PATCH] refactor: refactor upgrade register

---
 app/upgrade.go        | 28 ++++++++++++++--------------
 app/upgrades/types.go | 24 ++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/app/upgrade.go b/app/upgrade.go
index 83b573c89..8ccf810fb 100644
--- a/app/upgrade.go
+++ b/app/upgrade.go
@@ -14,13 +14,12 @@ import (
 )
 
 var (
-	plans = []upgrades.Upgrade{
-		v110.Upgrade,
-		v120.Upgrade,
-		v130.Upgrade,
-		v140.Upgrade,
-		v200.Upgrade,
-	}
+	router = upgrades.NewUpgradeRouter().
+		Register(v110.Upgrade).
+		Register(v120.Upgrade).
+		Register(v130.Upgrade).
+		Register(v140.Upgrade).
+		Register(v200.Upgrade)
 )
 
 // RegisterUpgradePlans register a handler of upgrade plan
@@ -57,17 +56,18 @@ func (app *IrisApp) setupUpgradeStoreLoaders() {
 		return
 	}
 
-	for _, upgrade := range plans {
-		if upgradeInfo.Name == upgrade.UpgradeName {
-			app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, upgrade.StoreUpgrades))
-		}
-	}
+	app.SetStoreLoader(
+		upgradetypes.UpgradeStoreLoader(
+			upgradeInfo.Height,
+			router.UpgradeInfo(upgradeInfo.Name).StoreUpgrades,
+		),
+	)
 }
 
 func (app *IrisApp) setupUpgradeHandlers() {
-	for _, upgrade := range plans {
+	for upgradeName, upgrade := range router.Routers() {
 		app.UpgradeKeeper.SetUpgradeHandler(
-			upgrade.UpgradeName,
+			upgradeName,
 			upgrade.UpgradeHandlerConstructor(
 				app.mm,
 				app.configurator,
diff --git a/app/upgrades/types.go b/app/upgrades/types.go
index 63ae0b38d..060444813 100644
--- a/app/upgrades/types.go
+++ b/app/upgrades/types.go
@@ -56,3 +56,27 @@ type AppKeepers struct {
 	TokenKeeper     tokenkeeper.Keeper
 	ReaderWriter    ConsensusParamsReaderWriter
 }
+
+type upgradeRouter struct {
+	mu map[string]Upgrade
+}
+
+func NewUpgradeRouter() *upgradeRouter {
+	return &upgradeRouter{make(map[string]Upgrade)}
+}
+
+func (r *upgradeRouter) Register(u Upgrade) *upgradeRouter {
+	if _, has := r.mu[u.UpgradeName]; has {
+		panic(u.UpgradeName + " already registered")
+	}
+	r.mu[u.UpgradeName] = u
+	return r
+}
+
+func (r *upgradeRouter) Routers() map[string]Upgrade {
+	return r.mu
+}
+
+func (r *upgradeRouter) UpgradeInfo(planName string) Upgrade {
+	return r.mu[planName]
+}