Store count directly instead of iterating

x/mailbox/bindings/message_plugin.go

@@ -61,7 +61,7 @@ func MsgDispatch(k *mailboxkeeper.Keeper, ctx sdk.Context, contractAddr sdk.AccA
 		return nil, wasmvmtypes.InvalidRequest{Err: "msgdispatch null msg"}
 	}
 
-	msgServer := mailboxkeeper.NewMsgServerImpl(*k)
+	msgServer := mailboxkeeper.NewMsgServerImpl(k)
 
 	msgMsgDispatch := mailboxtypes.NewMsgDispatch(msg.Sender, msg.DestinationDomain, msg.RecipientAddress, msg.MessageBody)
 

x/mailbox/keeper/genesis.go

@@ -15,17 +15,9 @@ import (
 // InitGenesis initializes the hyperlane mailbox module's state from a provided genesis
 // state.
 func (k *Keeper) InitGenesis(ctx sdk.Context, gs types.GenesisState) error {
-	tempTree := make(map[uint32][]byte, gs.Tree.Count)
-	for _, treeEntry := range gs.Tree.TreeEntries {
-		tempTree[treeEntry.Index] = treeEntry.Message
-	}
-	var index uint32
-	for index = 0; index < gs.Tree.Count; index++ {
-		err := k.Tree.Insert(tempTree[index])
-		if err != nil {
-			panic("unreachable")
-		}
-	}
+
+	k.ImtCount = gs.Tree.Count
+
 	for _, msgDelivered := range gs.DeliveredMessages {
 		k.Delivered[msgDelivered.Id] = true
 	}
@@ -39,40 +31,42 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) types.GenesisState {
 	store := ctx.KVStore(k.storeKey)
 
 	var genesisState types.GenesisState
-	genesisState.Tree = ExportTreeGenesis(store)
+	genesisState.Tree.Count = k.ImtCount
+	genesisState.Tree.TreeEntries = ExportTreeEntries(store)
 	genesisState.DeliveredMessages = ExportDeliveredGenesis(store)
 	return genesisState
 }
 
-func ExportTreeGenesis(store sdk.KVStore) types.Tree {
+func ExportTreeEntries(store sdk.KVStore) []*types.TreeEntry {
 	iterator := sdk.KVStorePrefixIterator(store, []byte(types.KeyMailboxIMT))
 	defer iterator.Close()
 
-	var genesisTree types.Tree
-	count := uint32(0)
+	var treeEntries []*types.TreeEntry
+	prefix := []byte(fmt.Sprintf("%s/", types.KeyMailboxIMT))
+
 	for ; iterator.Valid(); iterator.Next() {
-		indexBytes := bytes.TrimPrefix(iterator.Key(), []byte(fmt.Sprintf("%s/", types.KeyMailboxIMT)))
+		indexBytes := bytes.TrimPrefix(iterator.Key(), prefix)
 		index, err := strconv.ParseUint(string(indexBytes), 10, 32)
 		if err != nil {
 			panic(err)
 		}
-		genesisTree.TreeEntries = append(genesisTree.TreeEntries, &types.TreeEntry{
+		treeEntries = append(treeEntries, &types.TreeEntry{
 			Index:   uint32(index),
 			Message: iterator.Value(),
 		})
-		count++
 	}
-	genesisTree.Count = count
-	return genesisTree
+	return treeEntries
 }
 
 func ExportDeliveredGenesis(store sdk.KVStore) []*types.MessageDelivered {
 	iterator := sdk.KVStorePrefixIterator(store, []byte(types.KeyMailboxDelivered))
 	defer iterator.Close()
 
 	var delivered []*types.MessageDelivered
+	prefix := []byte(fmt.Sprintf("%s/", types.KeyMailboxDelivered))
+
 	for ; iterator.Valid(); iterator.Next() {
-		idBytes := bytes.TrimPrefix(iterator.Key(), []byte(fmt.Sprintf("%s/", types.KeyMailboxDelivered)))
+		idBytes := bytes.TrimPrefix(iterator.Key(), prefix)
 		delivered = append(delivered, &types.MessageDelivered{
 			Id: string(idBytes),
 		})

x/mailbox/keeper/genesis_test.go

@@ -31,7 +31,7 @@ func (suite *KeeperTestSuite) TestGenesis() {
 
 	// Verify tree exported correctly
 	gs := suite.keeper.ExportGenesis(suite.ctx)
-	suite.Require().Equal(uint32(100), gs.Tree.Count)
+	suite.Require().Equal(uint32(100), suite.keeper.ImtCount)
 	count := 0
 	for i := 0; i < 100; i++ {
 		for j := 0; j < 100; j++ {
@@ -60,6 +60,6 @@ func (suite *KeeperTestSuite) TestGenesis() {
 	suite.Require().NoError(err)
 
 	// Check Tree and Delivered
-	suite.Require().Equal(uint32(100), suite.keeper.Tree.Count())
+	suite.Require().Equal(uint32(100), suite.keeper.ImtCount)
 	suite.Require().Equal(100, len(suite.keeper.Delivered))
 }

x/mailbox/keeper/keeper.go

@@ -34,6 +34,7 @@ type Keeper struct {
 
 	Tree      *imt.Tree
 	Delivered map[string]bool
+	ImtCount  uint32
 }
 
 func NewKeeper(cdc codec.BinaryCodec, key storetypes.StoreKey, cwKeeper *cosmwasm.Keeper, ismKeeper *ismkeeper.Keeper) Keeper {
@@ -51,6 +52,7 @@ func NewKeeper(cdc codec.BinaryCodec, key storetypes.StoreKey, cwKeeper *cosmwas
 		pcwKeeper:   cosmwasm.NewDefaultPermissionKeeper(cwKeeper),
 		Tree:        &imt.Tree{},
 		Delivered:   map[string]bool{},
+		ImtCount:    0,
 	}
 }
 

x/mailbox/keeper/keeper_test.go

@@ -55,6 +55,6 @@ func (suite *KeeperTestSuite) SetupTest() {
 	queryClient := types.NewQueryClient(queryHelper)
 
 	suite.queryClient = queryClient
-	suite.msgServer = keeper.NewMsgServerImpl(suite.keeper)
+	suite.msgServer = keeper.NewMsgServerImpl(&suite.keeper)
 	suite.encCfg = encCfg
 }

x/mailbox/keeper/msg_server.go

@@ -31,12 +31,12 @@ type ContractProcessMsg struct {
 }
 
 // NewMsgServerImpl return an implementation of the mailbox MsgServer interface for the provided keeper
-func NewMsgServerImpl(keeper Keeper) types.MsgServer {
+func NewMsgServerImpl(keeper *Keeper) types.MsgServer {
 	return keeper
 }
 
 // Dispatch defines a rpc handler method for MsgDispatch
-func (k Keeper) Dispatch(goCtx context.Context, msg *types.MsgDispatch) (*types.MsgDispatchResponse, error) {
+func (k *Keeper) Dispatch(goCtx context.Context, msg *types.MsgDispatch) (*types.MsgDispatchResponse, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
 	// TODO: NewMessage
@@ -99,6 +99,7 @@ func (k Keeper) Dispatch(goCtx context.Context, msg *types.MsgDispatch) (*types.
 
 	// Insert the message id into the tree
 	err := k.Tree.Insert(id)
+	k.ImtCount++
 	if err != nil {
 		return nil, err
 	}
@@ -134,7 +135,7 @@ func (k Keeper) Dispatch(goCtx context.Context, msg *types.MsgDispatch) (*types.
 }
 
 // Process defines a rpc handler method for MsgProcess
-func (k Keeper) Process(goCtx context.Context, msg *types.MsgProcess) (*types.MsgProcessResponse, error) {
+func (k *Keeper) Process(goCtx context.Context, msg *types.MsgProcess) (*types.MsgProcessResponse, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
 	messageBytes := hexutil.MustDecode(msg.Message)

x/mailbox/module.go

@@ -92,7 +92,7 @@ func NewAppModule(k keeper.Keeper) AppModule {
 
 // RegisterServices registers module services.
 func (am AppModule) RegisterServices(cfg module.Configurator) {
-	types.RegisterMsgServer(cfg.MsgServer(), am.keeper)
+	types.RegisterMsgServer(cfg.MsgServer(), &am.keeper)
 	types.RegisterQueryServer(cfg.QueryServer(), am.keeper)
 }