From 6a71537f39f4d6da23f68b43dc5b44182ee440a6 Mon Sep 17 00:00:00 2001
From: vimystic <122659254+vimystic@users.noreply.github.com>
Date: Wed, 4 Oct 2023 11:49:17 -0600
Subject: [PATCH] Update based on feedback.

---
 imt/tree.go                    | 2 +-
 x/mailbox/keeper/msg_server.go | 9 ++++++---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/imt/tree.go b/imt/tree.go
index 2b4edef..594e440 100644
--- a/imt/tree.go
+++ b/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
 }
diff --git a/x/mailbox/keeper/msg_server.go b/x/mailbox/keeper/msg_server.go
index 2b9ccbb..bdeec49 100644
--- a/x/mailbox/keeper/msg_server.go
+++ b/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(