refactor: generic reading of merkle files

x/bundles/migration/migration.go

@@ -4,6 +4,7 @@ import (
 	"cosmossdk.io/log"
 	"cosmossdk.io/store/prefix"
 	storeTypes "cosmossdk.io/store/types"
+	"embed"
 	_ "embed"
 	"encoding/hex"
 	"fmt"
@@ -12,6 +13,8 @@ import (
 	"github.com/KYVENetwork/chain/x/bundles/types"
 	"github.com/cosmos/cosmos-sdk/runtime"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"strconv"
+	"strings"
 )
 
 var logger log.Logger
@@ -21,27 +24,41 @@ const (
 	WaitingBlockPeriod                     = 1
 )
 
-//go:embed files/merkle_roots_pool_1
-var merkleRootsPool1 []byte
-
-//go:embed files/merkle_roots_pool_2
-var merkleRootsPool2 []byte
+//go:embed files/*
+var merkelRoots embed.FS
 
 type BundlesMigrationEntry struct {
 	merkleRoots []byte
 	poolId      uint64
+	maxBundleId uint64
 }
 
 // bundlesMigration includes the poolId and max bundleId to determine which bundles are migrated
-var bundlesMigration = []BundlesMigrationEntry{
-	{
-		merkleRoots: merkleRootsPool1,
-		poolId:      1,
-	},
-	{
-		merkleRoots: merkleRootsPool2,
-		poolId:      2,
-	},
+var bundlesMigration []BundlesMigrationEntry
+
+func init() {
+	dir, err := merkelRoots.ReadDir("files")
+	if err != nil {
+		panic(err)
+	}
+
+	for _, file := range dir {
+		readFile, err := merkelRoots.ReadFile(fmt.Sprintf("files/%s", file.Name()))
+		if err != nil {
+			panic(err)
+		}
+
+		poolId, err := strconv.ParseUint(strings.ReplaceAll(file.Name(), "merkle_roots_pool_", ""), 10, 64)
+		if err != nil {
+			panic(err)
+		}
+
+		bundlesMigration = append(bundlesMigration, BundlesMigrationEntry{
+			merkleRoots: readFile,
+			poolId:      poolId,
+			maxBundleId: uint64(len(readFile)) / 32,
+		})
+	}
 }
 
 func MigrateBundlesModule(sdkCtx sdk.Context, bundlesKeeper bundleskeeper.Keeper, upgradeHeight int64) {
@@ -54,21 +71,13 @@ func MigrateBundlesModule(sdkCtx sdk.Context, bundlesKeeper bundleskeeper.Keeper
 		step := uint64(sdkCtx.BlockHeight()-upgradeHeight) - WaitingBlockPeriod
 		offset := step * BundlesMigrationStepSizePerPool
 
-		var maxBundleId uint64
-		switch bundlesMigrationEntry.poolId {
-		case 1:
-			maxBundleId = uint64(len(merkleRootsPool1)) / 32
-		case 2:
-			maxBundleId = uint64(len(merkleRootsPool2)) / 32
-		}
-
 		// Exit if all bundles have already been migrated
-		if offset > maxBundleId {
+		if offset > bundlesMigrationEntry.maxBundleId {
 			logger.Info("offset > maxBundleId > return")
 			return
 		}
 
-		if err := migrateFinalizedBundles(sdkCtx, bundlesKeeper, offset, bundlesMigrationEntry.poolId, maxBundleId); err != nil {
+		if err := migrateFinalizedBundles(sdkCtx, bundlesKeeper, offset, bundlesMigrationEntry); err != nil {
 			// TODO: Error handling
 			panic(err)
 		}
@@ -77,10 +86,10 @@ func MigrateBundlesModule(sdkCtx sdk.Context, bundlesKeeper bundleskeeper.Keeper
 
 // MigrateFinalizedBundles ...
 // maxBundleId -> inclusive
-func migrateFinalizedBundles(ctx sdk.Context, bundlesKeeper bundleskeeper.Keeper, offset uint64, poolId uint64, maxBundleId uint64) error {
+func migrateFinalizedBundles(ctx sdk.Context, bundlesKeeper bundleskeeper.Keeper, offset uint64, bundlesMigrationEntry BundlesMigrationEntry) error {
 	// Init Bundles Store
 	storeAdapter := runtime.KVStoreAdapter(bundlesKeeper.Migration_GetStoreService().OpenKVStore(ctx))
-	store := prefix.NewStore(storeAdapter, util.GetByteKey(types.FinalizedBundlePrefix, poolId))
+	store := prefix.NewStore(storeAdapter, util.GetByteKey(types.FinalizedBundlePrefix, bundlesMigrationEntry.poolId))
 
 	var iterator storeTypes.Iterator
 	iterator = store.Iterator(util.GetByteKey(offset), util.GetByteKey(offset+BundlesMigrationStepSizePerPool))
@@ -93,17 +102,12 @@ func migrateFinalizedBundles(ctx sdk.Context, bundlesKeeper bundleskeeper.Keeper
 			return err
 		}
 
-		if rawFinalizedBundle.Id >= maxBundleId {
+		if rawFinalizedBundle.Id >= bundlesMigrationEntry.maxBundleId {
 			return nil
 		}
 
-		var merkleRoot []byte
-		switch rawFinalizedBundle.PoolId {
-		case 1:
-			merkleRoot = merkleRootsPool1[rawFinalizedBundle.Id*32 : rawFinalizedBundle.Id*32+32]
-		case 2:
-			merkleRoot = merkleRootsPool2[rawFinalizedBundle.Id*32 : rawFinalizedBundle.Id*32+32]
-		}
+		merkleRoot := bundlesMigrationEntry.merkleRoots[rawFinalizedBundle.Id*32 : rawFinalizedBundle.Id*32+32]
+
 		rawFinalizedBundle.BundleSummary = fmt.Sprintf("{\"merkle_root\":\"%v\"}", hex.EncodeToString(merkleRoot))
 
 		migratedBundles = append(migratedBundles, rawFinalizedBundle)