Speed up ExportGenesis

modules/core/02-client/genesis.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	connectionkeeper "github.com/cosmos/ibc-go/v7/modules/core/03-connection/keeper"
 
 	"github.com/cosmos/ibc-go/v7/modules/core/02-client/keeper"
 	"github.com/cosmos/ibc-go/v7/modules/core/02-client/types"
@@ -55,19 +56,19 @@ func InitGenesis(ctx sdk.Context, k keeper.Keeper, gs types.GenesisState) {
 }
 
 // ExportGenesis returns the ibc client submodule's exported genesis.
-func ExportGenesis(ctx sdk.Context, k keeper.Keeper) types.GenesisState {
-	genClients := k.GetAllGenesisClients(ctx)
-	clientsMetadata, err := k.GetAllClientMetadata(ctx, genClients)
+func ExportGenesis(ctx sdk.Context, clientKeeper keeper.Keeper, connectionKeeper connectionkeeper.Keeper) types.GenesisState {
+	genClients := clientKeeper.GetAllGenesisClients(ctx, connectionKeeper)
+	clientsMetadata, err := clientKeeper.GetAllClientMetadata(ctx, genClients)
 	if err != nil {
 		panic(err)
 	}
 	return types.GenesisState{
 		Clients:          genClients,
 		ClientsMetadata:  clientsMetadata,
-		ClientsConsensus: k.GetAllConsensusStates(ctx),
-		Params:           k.GetParams(ctx),
+		ClientsConsensus: clientKeeper.GetAllConsensusStates(ctx, connectionKeeper),
+		Params:           clientKeeper.GetParams(ctx),
 		// Warning: CreateLocalhost is deprecated
 		CreateLocalhost:    false,
-		NextClientSequence: k.GetNextClientSequence(ctx),
+		NextClientSequence: clientKeeper.GetNextClientSequence(ctx),
 	}
 }

modules/core/02-client/keeper/keeper.go

@@ -14,6 +14,7 @@ import (
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
 	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+	connectionkeeper "github.com/cosmos/ibc-go/v7/modules/core/03-connection/keeper"
 
 	"github.com/cosmos/ibc-go/v7/modules/core/02-client/types"
 	commitmenttypes "github.com/cosmos/ibc-go/v7/modules/core/23-commitment/types"
@@ -133,9 +134,9 @@ func (k Keeper) SetNextClientSequence(ctx sdk.Context, sequence uint64) {
 // IterateConsensusStates provides an iterator over all stored consensus states.
 // objects. For each State object, cb will be called. If the cb returns true,
 // the iterator will close and stop.
-func (k Keeper) IterateConsensusStates(ctx sdk.Context, cb func(clientID string, cs types.ConsensusStateWithHeight) bool) {
+func (k Keeper) IterateConsensusStates(ctx sdk.Context, prefix []byte, cb func(clientID string, cs types.ConsensusStateWithHeight) bool) {
 	store := ctx.KVStore(k.storeKey)
-	iterator := sdk.KVStorePrefixIterator(store, host.KeyClientStorePrefix)
+	iterator := sdk.KVStorePrefixIterator(store, prefix)
 
 	defer sdk.LogDeferred(ctx.Logger(), func() error { return iterator.Close() })
 	for ; iterator.Valid(); iterator.Next() {
@@ -157,12 +158,16 @@ func (k Keeper) IterateConsensusStates(ctx sdk.Context, cb func(clientID string,
 }
 
 // GetAllGenesisClients returns all the clients in state with their client ids returned as IdentifiedClientState
-func (k Keeper) GetAllGenesisClients(ctx sdk.Context) types.IdentifiedClientStates {
+func (k Keeper) GetAllGenesisClients(ctx sdk.Context, connectionKeeper connectionkeeper.Keeper) types.IdentifiedClientStates {
+	clientIds := connectionKeeper.GetAllClientIds(ctx)
+
 	var genClients types.IdentifiedClientStates
-	k.IterateClientStates(ctx, nil, func(clientID string, cs exported.ClientState) bool {
-		genClients = append(genClients, types.NewIdentifiedClientState(clientID, cs))
-		return false
-	})
+	for _, clientID := range clientIds {
+		k.IterateClientStates(ctx, []byte(clientID), func(clientID string, cs exported.ClientState) bool {
+			genClients = append(genClients, types.NewIdentifiedClientState(clientID, cs))
+			return false
+		})
+	}
 
 	return genClients.Sort()
 }
@@ -211,26 +216,30 @@ func (k Keeper) SetAllClientMetadata(ctx sdk.Context, genMetadata []types.Identi
 }
 
 // GetAllConsensusStates returns all stored client consensus states.
-func (k Keeper) GetAllConsensusStates(ctx sdk.Context) types.ClientsConsensusStates {
+func (k Keeper) GetAllConsensusStates(ctx sdk.Context, connectionKeeper connectionkeeper.Keeper) types.ClientsConsensusStates {
 	clientConsStates := make(types.ClientsConsensusStates, 0)
 	mapClientIDToConsStateIdx := make(map[string]int)
+	clientIds := connectionKeeper.GetAllClientIds(ctx)
+
+	for _, _clientId := range clientIds {
+		prefix := host.PrefixedClientStoreKey([]byte(_clientId))
+		k.IterateConsensusStates(ctx, prefix, func(clientID string, cs types.ConsensusStateWithHeight) bool {
+			idx, ok := mapClientIDToConsStateIdx[clientID]
+			if ok {
+				clientConsStates[idx].ConsensusStates = append(clientConsStates[idx].ConsensusStates, cs)
+				return false
+			}
 
-	k.IterateConsensusStates(ctx, func(clientID string, cs types.ConsensusStateWithHeight) bool {
-		idx, ok := mapClientIDToConsStateIdx[clientID]
-		if ok {
-			clientConsStates[idx].ConsensusStates = append(clientConsStates[idx].ConsensusStates, cs)
-			return false
-		}
-
-		clientConsState := types.ClientConsensusStates{
-			ClientId:        clientID,
-			ConsensusStates: []types.ConsensusStateWithHeight{cs},
-		}
+			clientConsState := types.ClientConsensusStates{
+				ClientId:        clientID,
+				ConsensusStates: []types.ConsensusStateWithHeight{cs},
+			}
 
-		clientConsStates = append(clientConsStates, clientConsState)
-		mapClientIDToConsStateIdx[clientID] = len(clientConsStates) - 1
-		return false
-	})
+			clientConsStates = append(clientConsStates, clientConsState)
+			mapClientIDToConsStateIdx[clientID] = len(clientConsStates) - 1
+			return false
+		})
+	}
 
 	return clientConsStates.Sort()
 }

modules/core/03-connection/keeper/keeper.go

@@ -152,21 +152,40 @@ func (k Keeper) SetNextConnectionSequence(ctx sdk.Context, sequence uint64) {
 // will ignore the clients that haven't initialized a connection handshake since
 // no paths are stored.
 func (k Keeper) GetAllClientConnectionPaths(ctx sdk.Context) []types.ConnectionPaths {
+	clientIds := k.GetAllClientIds(ctx)
+
 	var allConnectionPaths []types.ConnectionPaths
-	k.clientKeeper.IterateClientStates(ctx, nil, func(clientID string, cs exported.ClientState) bool {
+	for _, clientID := range clientIds {
 		paths, found := k.GetClientConnectionPaths(ctx, clientID)
 		if !found {
 			// continue when connection handshake is not initialized
-			return false
+			continue
 		}
 		connPaths := types.NewConnectionPaths(clientID, paths)
 		allConnectionPaths = append(allConnectionPaths, connPaths)
-		return false
-	})
+
+	}
 
 	return allConnectionPaths
 }
 
+func (k Keeper) GetAllClientIds(ctx sdk.Context) []string {
+	connections := k.GetAllConnections(ctx)
+	clientIds := make([]string, 0)
+	clientIdSet := make(map[string]bool) // To avoid duplicate client IDs
+
+	for _, connection := range connections {
+		if connection.ClientId == "polymer-0" {
+			continue
+		}
+		if !clientIdSet[connection.ClientId] {
+			clientIdSet[connection.ClientId] = true
+			clientIds = append(clientIds, connection.ClientId)
+		}
+	}
+	return clientIds
+}
+
 // IterateConnections provides an iterator over all ConnectionEnd objects.
 // For each ConnectionEnd, cb will be called. If the cb returns true, the
 // iterator will close and stop.

modules/core/genesis.go

@@ -21,7 +21,7 @@ func InitGenesis(ctx sdk.Context, k keeper.Keeper, gs *types.GenesisState) {
 // ExportGenesis returns the ibc exported genesis.
 func ExportGenesis(ctx sdk.Context, k keeper.Keeper) *types.GenesisState {
 	return &types.GenesisState{
-		ClientGenesis:     client.ExportGenesis(ctx, k.ClientKeeper),
+		ClientGenesis:     client.ExportGenesis(ctx, k.ClientKeeper, k.ConnectionKeeper),
 		ConnectionGenesis: connection.ExportGenesis(ctx, k.ConnectionKeeper),
 		ChannelGenesis:    channel.ExportGenesis(ctx, k.ChannelKeeper),
 	}