Update based on feedback.

imt/tree.go

@@ -51,7 +51,7 @@ func (t *Tree) Count() uint32 {
 	return t.count
 }
 
-// Count returns the number of inserts performed on the Tree
+// Set the count . (Used when restoring from genesis state)
 func (t *Tree) SetCount(count uint32) {
 	t.count = count
 }

x/mailbox/keeper/msg_server.go

@@ -46,7 +46,7 @@ func (k *Keeper) Dispatch(goCtx context.Context, msg *types.MsgDispatch) (*types
 	version := make([]byte, 1)
 	message = append(message, version...)
 
-	//Nonce is the current branch length.
+	// Nonce is the tree count.
 	nonce := k.Tree.Count()
 	nonceBytes := make([]byte, 4)
 	binary.BigEndian.PutUint32(nonceBytes, nonce)
@@ -95,15 +95,18 @@ func (k *Keeper) Dispatch(goCtx context.Context, msg *types.MsgDispatch) (*types
 	message = append(message, messageBytes...)
 	hyperlaneMsg := hexutil.Encode(message)
 
-	// Get the message ID. (i.e: Keccak256 sha)
+	// Get the message ID. (i.e: Keccak256 hash of the message)
 	id := common.Id(message)
 
-	store.Set(types.MailboxIMTKey(), id)
+	// Insert the message id into the tree
 	err := k.Tree.Insert(id)
 	if err != nil {
 		return nil, err
 	}
 
+	// Store that the current root
+	store.Set(types.MailboxIMTKey(), k.Tree.Root())
+
 	// Emit the events
 	ctx.EventManager().EmitEvents(sdk.Events{
 		sdk.NewEvent(