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/02-client/keeper/keeper_test.go

@@ -253,7 +253,7 @@ func (suite KeeperTestSuite) TestGetAllGenesisClients() { //nolint:govet // this
 		expGenClients[i] = types.NewIdentifiedClientState(clientIDs[i], expClients[i])
 	}
 
-	genClients := suite.chainA.App.GetIBCKeeper().ClientKeeper.GetAllGenesisClients(suite.chainA.GetContext())
+	genClients := suite.chainA.App.GetIBCKeeper().ClientKeeper.GetAllGenesisClients(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ConnectionKeeper)
 
 	suite.Require().Equal(expGenClients.Sort(), genClients)
 }
@@ -385,7 +385,7 @@ func (suite KeeperTestSuite) TestGetAllConsensusStates() { //nolint:govet // thi
 		}),
 	}.Sort()
 
-	consStates := suite.chainA.App.GetIBCKeeper().ClientKeeper.GetAllConsensusStates(suite.chainA.GetContext())
+	consStates := suite.chainA.App.GetIBCKeeper().ClientKeeper.GetAllConsensusStates(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ConnectionKeeper)
 	suite.Require().Equal(expConsensusStates, consStates, "%s \n\n%s", expConsensusStates, consStates)
 }
 

modules/core/02-client/migrations/v7/genesis_test.go

@@ -33,7 +33,7 @@ func (suite *MigrationsV7TestSuite) TestMigrateGenesisSolomachine() {
 	solomachine := ibctesting.NewSolomachine(suite.T(), suite.chainA.Codec, "06-solomachine-0", "testing", 1)
 	solomachineMulti := ibctesting.NewSolomachine(suite.T(), suite.chainA.Codec, "06-solomachine-1", "testing", 4)
 
-	clientGenState := ibcclient.ExportGenesis(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ClientKeeper)
+	clientGenState := ibcclient.ExportGenesis(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ClientKeeper, suite.chainA.App.GetIBCKeeper().ConnectionKeeper)
 
 	// manually generate old proto buf definitions and set in genesis
 	// NOTE: we cannot use 'ExportGenesis' for the solo machines since we are
@@ -108,7 +108,7 @@ func (suite *MigrationsV7TestSuite) TestMigrateGenesisSolomachine() {
 	// NOTE: tendermint clients are not pruned in genesis so the test should not have expired tendermint clients
 	err := v7.MigrateStore(suite.chainA.GetContext(), suite.chainA.GetSimApp().GetKey(ibcexported.StoreKey), suite.chainA.App.AppCodec(), suite.chainA.GetSimApp().IBCKeeper.ClientKeeper)
 	suite.Require().NoError(err)
-	expectedClientGenState := ibcclient.ExportGenesis(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ClientKeeper)
+	expectedClientGenState := ibcclient.ExportGenesis(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ClientKeeper, suite.chainA.App.GetIBCKeeper().ConnectionKeeper)
 
 	cdc, ok := suite.chainA.App.AppCodec().(codec.ProtoCodecMarshaler)
 	suite.Require().True(ok)

modules/core/02-client/types/genesis_test.go

@@ -34,7 +34,7 @@ func (suite *TypesTestSuite) TestMarshalGenesisState() {
 	err := path.EndpointA.UpdateClient()
 	suite.Require().NoError(err)
 
-	genesis := client.ExportGenesis(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ClientKeeper)
+	genesis := client.ExportGenesis(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ClientKeeper, suite.chainA.App.GetIBCKeeper().ConnectionKeeper)
 
 	bz, err := cdc.MarshalJSON(&genesis)
 	suite.Require().NoError(err)

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

@@ -152,21 +152,43 @@ 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 {
-		paths, found := k.GetClientConnectionPaths(ctx, clientID)
-		if !found {
-			// continue when connection handshake is not initialized
+	for _, _clientId := range clientIds {
+		prefix := host.PrefixedClientStoreKey([]byte(_clientId))
+		k.clientKeeper.IterateClientStates(ctx, prefix, func(clientID string, cs exported.ClientState) bool {
+			paths, found := k.GetClientConnectionPaths(ctx, clientID)
+			if !found {
+				// continue when connection handshake is not initialized
+				return false
+			}
+			connPaths := types.NewConnectionPaths(clientID, paths)
+			allConnectionPaths = append(allConnectionPaths, connPaths)
 			return false
-		}
-		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),
 	}

modules/core/migrations/v7/genesis_test.go

@@ -60,7 +60,7 @@ func (suite *MigrationsV7TestSuite) TestMigrateGenesisSolomachine() {
 	solomachine := ibctesting.NewSolomachine(suite.T(), suite.chainA.Codec, "06-solomachine-0", "testing", 1)
 	solomachineMulti := ibctesting.NewSolomachine(suite.T(), suite.chainA.Codec, "06-solomachine-1", "testing", 4)
 
-	clientGenState := ibcclient.ExportGenesis(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ClientKeeper)
+	clientGenState := ibcclient.ExportGenesis(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ClientKeeper, suite.chainA.App.GetIBCKeeper().ConnectionKeeper)
 
 	// manually generate old proto buf definitions and set in genesis
 	// NOTE: we cannot use 'ExportGenesis' for the solo machines since we are
@@ -134,7 +134,7 @@ func (suite *MigrationsV7TestSuite) TestMigrateGenesisSolomachine() {
 	// NOTE: tendermint clients are not pruned in genesis so the test should not have expired tendermint clients
 	err := clientv7.MigrateStore(suite.chainA.GetContext(), suite.chainA.GetSimApp().GetKey(ibcexported.StoreKey), suite.chainA.App.AppCodec(), suite.chainA.GetSimApp().IBCKeeper.ClientKeeper)
 	suite.Require().NoError(err)
-	expectedClientGenState := ibcclient.ExportGenesis(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ClientKeeper)
+	expectedClientGenState := ibcclient.ExportGenesis(suite.chainA.GetContext(), suite.chainA.App.GetIBCKeeper().ClientKeeper, suite.chainA.App.GetIBCKeeper().ConnectionKeeper)
 
 	cdc := suite.chainA.App.AppCodec().(*codec.ProtoCodec)