Count should not be branch size. Its number of messages processed

imt/tree.go

@@ -19,32 +19,34 @@ type Tree struct {
 }
 
 // Insert inserts node and returns the branches
-func (t *Tree) InsertAndReturnBranches(node []byte) ([32][]byte, error) {
+func (t *Tree) InsertAndReturnBranches(node []byte) ([32][]byte, uint32, error) {
+
 	if t.count >= MaxLeaves {
-		return [32][]byte{}, errors.New("merkle tree full")
+		return [32][]byte{}, t.count, errors.New("merkle tree full")
 	}
 
 	if len(node) != 32 {
-		return [32][]byte{}, errors.New("must be 32-bytes")
+		return [32][]byte{}, t.count, errors.New("must be 32-bytes")
 	}
 
 	t.count += 1
 	size := t.count
 	for i := 0; i < TreeDepth; i++ {
 		if (size & 1) == 1 {
 			t.Branch[i] = node
-			return t.Branch, nil
+			return t.Branch, t.count, nil
 		}
 		temp := append(t.Branch[i][:], node...)
 		node = crypto.Keccak256(temp)
 		size /= 2
 	}
 
-	return [32][]byte{}, errors.New("unreachable")
+	return [32][]byte{}, t.count, errors.New("unreachable")
+
 }
 
 func (t *Tree) Insert(node []byte) error {
-	_, errors := t.InsertAndReturnBranches(node)
+	_, _, errors := t.InsertAndReturnBranches(node)
 	return errors
 }
 
@@ -53,6 +55,11 @@ func (t *Tree) Count() uint32 {
 	return t.count
 }
 
+// Count returns the number of inserts performed on the Tree
+func (t *Tree) SetCount(count uint32) {
+	t.count = count
+}
+
 // Print dumps the tree (for debugging)
 func (t *Tree) Print() {
 	for i := 0; i < TreeDepth; i++ {

x/mailbox/keeper/genesis.go

@@ -16,6 +16,8 @@ func (k *Keeper) InitGenesis(ctx sdk.Context, gs types.GenesisState) error {
 
 	copy(k.Tree.Branch[:], gs.Tree.Branch)
 
+	k.Tree.SetCount(gs.Tree.Count)
+
 	//Delivered Messages.
 	for _, msgDelivered := range gs.DeliveredMessages {
 		k.Delivered[msgDelivered.Id] = true
@@ -32,7 +34,7 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) types.GenesisState {
 		DeliveredMessages: ExportDeliveredMessages(ctx.KVStore(k.storeKey)),
 		Tree: types.Tree{
 			Branch: k.Tree.Branch[:],
-			Count:  uint32(32), //TODO:Should we change this to reflect only levels populated ?
+			Count:  k.Tree.Count(),
 		},
 		Domain: k.GetDomain(ctx),
 	}

x/mailbox/keeper/genesis_test.go

@@ -10,7 +10,7 @@ import (
 )
 
 func (suite *KeeperTestSuite) TestGenesis() {
-	idMap := make([]string, 128)
+	idMap := make([]string, 128) //Will populate 8 levels.
 	for i := 0; i < 128; i++ {
 		sender := "cosmos14hj2tavq8fpesdwxxcu44rty3hh90vhujrvcmstl4zr3txmfvw9s4hmalr"
 		recipientBech32 := "cosmos10qa7yajp3fp869mdegtpap5zg056exja3chkw5"
@@ -26,10 +26,10 @@ func (suite *KeeperTestSuite) TestGenesis() {
 
 	// Exporting Genesis and logging the length of Branches
 	gs := suite.keeper.ExportGenesis(suite.ctx)
-	suite.Require().Equal(8, countPopulatedSlices(suite.keeper.Branch)) // 2^7 + 36 = 100  .. only 8 levels will be populated.
+	suite.Require().Equal(8, countPopulatedSlices(suite.keeper.Branch))
 
 	// Adding delivered message ids to the exported state
-	for i := 0; i < 100; i++ {
+	for i := 0; i < 128; i++ {
 		gs.DeliveredMessages = append(gs.DeliveredMessages, &types.MessageDelivered{
 			Id: idMap[i],
 		})
@@ -43,11 +43,10 @@ func (suite *KeeperTestSuite) TestGenesis() {
 	suite.Require().NoError(err)
 
 	suite.Require().Equal(8, countPopulatedSlices(suite.keeper.Tree.Branch))
-	suite.Require().Equal(100, len(suite.keeper.Delivered))
+	suite.Require().Equal(128, len(suite.keeper.Delivered))
 }
 
 func countPopulatedSlices(arr [32][]byte) int {
-	fmt.Printf("%+v", arr)
 	count := 0
 	for _, slice := range arr {
 		if len(slice) > 0 {

x/mailbox/keeper/msg_server.go

@@ -47,7 +47,7 @@ func (k *Keeper) Dispatch(goCtx context.Context, msg *types.MsgDispatch) (*types
 	message = append(message, version...)
 
 	//Nonce is the current branch length.
-	nonce := uint32(0) //TODO: It should be this --> uint32(len(k.Tree.Branch))
+	nonce := k.Tree.Count()
 	nonceBytes := make([]byte, 4)
 	binary.BigEndian.PutUint32(nonceBytes, nonce)
 	message = append(message, nonceBytes...)
@@ -100,7 +100,7 @@ func (k *Keeper) Dispatch(goCtx context.Context, msg *types.MsgDispatch) (*types
 
 	store.Set(types.MailboxIMTKey(), id)
 
-	branch, err := k.Tree.InsertAndReturnBranches(id)
+	branch, _, err := k.Tree.InsertAndReturnBranches(id)
 
 	if err == nil {
 		k.Branch = branch