Properly export and rollback accounts that were created after the target slot

pkg/protocol/engine/accounts/accountsledger/manager.go

@@ -220,7 +220,7 @@ func (m *Manager) account(accountID iotago.AccountID, targetSlot iotago.SlotInde
 		loadedAccount = accounts.NewAccountData(accountID, accounts.WithCredits(accounts.NewBlockIssuanceCredits(0, targetSlot)))
 	}
 
-	wasDestroyed, err := m.rollbackAccountTo(loadedAccount, targetSlot)
+	_, wasDestroyed, err := m.rollbackAccountTo(loadedAccount, targetSlot)
 	if err != nil {
 		return nil, false, err
 	}
@@ -250,7 +250,7 @@ func (m *Manager) PastAccounts(accountIDs iotago.AccountIDs, targetSlot iotago.S
 		if !exists {
 			loadedAccount = accounts.NewAccountData(accountID, accounts.WithCredits(accounts.NewBlockIssuanceCredits(0, targetSlot)))
 		}
-		wasDestroyed, err := m.rollbackAccountTo(loadedAccount, targetSlot)
+		_, wasDestroyed, err := m.rollbackAccountTo(loadedAccount, targetSlot)
 		if err != nil {
 			continue
 		}
@@ -290,7 +290,7 @@ func (m *Manager) Rollback(targetSlot iotago.SlotIndex) error {
 				accountData = accounts.NewAccountData(accountID)
 			}
 
-			if _, err := m.rollbackAccountTo(accountData, targetSlot); err != nil {
+			if _, _, err := m.rollbackAccountTo(accountData, targetSlot); err != nil {
 				internalErr = ierrors.Wrapf(err, "unable to rollback account %s to target slot %d", accountID, targetSlot)
 
 				return false
@@ -314,7 +314,11 @@ func (m *Manager) Rollback(targetSlot iotago.SlotIndex) error {
 		}
 	}
 
-	return m.accountsTree.Commit()
+	if err := m.accountsTree.Commit(); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 // AddAccount adds a new account to the Account tree, allotting to it the balance on the given output.
@@ -368,17 +372,17 @@ func (m *Manager) Reset() {
 	m.latestSupportedVersionSignals.Clear()
 }
 
-func (m *Manager) rollbackAccountTo(accountData *accounts.AccountData, targetSlot iotago.SlotIndex) (wasDestroyed bool, err error) {
+func (m *Manager) rollbackAccountTo(accountData *accounts.AccountData, targetSlot iotago.SlotIndex) (wasCreatedAfterTargetSlot bool, wasDestroyed bool, err error) {
 	// to reach targetSlot, we need to rollback diffs from the current latestCommittedSlot down to targetSlot + 1
 	for diffSlot := m.latestCommittedSlot; diffSlot > targetSlot; diffSlot-- {
 		diffStore, err := m.slotDiff(diffSlot)
 		if err != nil {
-			return false, ierrors.Errorf("can't retrieve account, could not find diff store for slot %d", diffSlot)
+			return false, false, ierrors.Errorf("can't retrieve account, could not find diff store for slot %d", diffSlot)
 		}
 
 		found, err := diffStore.Has(accountData.ID)
 		if err != nil {
-			return false, ierrors.Wrapf(err, "can't retrieve account, could not check if diff store for slot %d has account %s", diffSlot, accountData.ID)
+			return false, false, ierrors.Wrapf(err, "can't retrieve account, could not check if diff store for slot %d has account %s", diffSlot, accountData.ID)
 		}
 
 		// no changes for this account in this slot
@@ -388,7 +392,7 @@ func (m *Manager) rollbackAccountTo(accountData *accounts.AccountData, targetSlo
 
 		diffChange, destroyed, err := diffStore.Load(accountData.ID)
 		if err != nil {
-			return false, ierrors.Wrapf(err, "can't retrieve account, could not load diff for account %s in slot %d", accountData.ID, diffSlot)
+			return false, false, ierrors.Wrapf(err, "can't retrieve account, could not load diff for account %s in slot %d", accountData.ID, diffSlot)
 		}
 
 		// update the account data with the diff
@@ -397,34 +401,39 @@ func (m *Manager) rollbackAccountTo(accountData *accounts.AccountData, targetSlo
 		if diffChange.PreviousExpirySlot != diffChange.NewExpirySlot {
 			accountData.ExpirySlot = diffChange.PreviousExpirySlot
 		}
-		// update the outputID only if the account got actually transitioned, not if it was only an allotment target
-		if diffChange.PreviousOutputID != iotago.EmptyOutputID {
-			accountData.OutputID = diffChange.PreviousOutputID
+
+		if diffChange.PreviousOutputID == iotago.EmptyOutputID {
+			// Account was created in this slot, so we need to remove it
+			return true, false, nil
 		}
+
+		// update the output ID of the account if it was changed
+		accountData.OutputID = diffChange.PreviousOutputID
+
 		accountData.AddBlockIssuerKeys(diffChange.BlockIssuerKeysRemoved...)
 		accountData.RemoveBlockIssuerKey(diffChange.BlockIssuerKeysAdded...)
 
 		validatorStake, err := safemath.SafeSub(int64(accountData.ValidatorStake), diffChange.ValidatorStakeChange)
 		if err != nil {
-			return false, ierrors.Wrapf(err, "can't retrieve account, validator stake underflow for account %s in slot %d: %d - %d", accountData.ID, diffSlot, accountData.ValidatorStake, diffChange.ValidatorStakeChange)
+			return false, false, ierrors.Wrapf(err, "can't retrieve account, validator stake underflow for account %s in slot %d: %d - %d", accountData.ID, diffSlot, accountData.ValidatorStake, diffChange.ValidatorStakeChange)
 		}
 		accountData.ValidatorStake = iotago.BaseToken(validatorStake)
 
 		delegationStake, err := safemath.SafeSub(int64(accountData.DelegationStake), diffChange.DelegationStakeChange)
 		if err != nil {
-			return false, ierrors.Wrapf(err, "can't retrieve account, delegation stake underflow for account %s in slot %d: %d - %d", accountData.ID, diffSlot, accountData.DelegationStake, diffChange.DelegationStakeChange)
+			return false, false, ierrors.Wrapf(err, "can't retrieve account, delegation stake underflow for account %s in slot %d: %d - %d", accountData.ID, diffSlot, accountData.DelegationStake, diffChange.DelegationStakeChange)
 		}
 		accountData.DelegationStake = iotago.BaseToken(delegationStake)
 
 		stakeEpochEnd, err := safemath.SafeSub(int64(accountData.StakeEndEpoch), diffChange.StakeEndEpochChange)
 		if err != nil {
-			return false, ierrors.Wrapf(err, "can't retrieve account, stake end epoch underflow for account %s in slot %d: %d - %d", accountData.ID, diffSlot, accountData.StakeEndEpoch, diffChange.StakeEndEpochChange)
+			return false, false, ierrors.Wrapf(err, "can't retrieve account, stake end epoch underflow for account %s in slot %d: %d - %d", accountData.ID, diffSlot, accountData.StakeEndEpoch, diffChange.StakeEndEpochChange)
 		}
 		accountData.StakeEndEpoch = iotago.EpochIndex(stakeEpochEnd)
 
 		fixedCost, err := safemath.SafeSub(int64(accountData.FixedCost), diffChange.FixedCostChange)
 		if err != nil {
-			return false, ierrors.Wrapf(err, "can't retrieve account, fixed cost underflow for account %s in slot %d: %d - %d", accountData.ID, diffSlot, accountData.FixedCost, diffChange.FixedCostChange)
+			return false, false, ierrors.Wrapf(err, "can't retrieve account, fixed cost underflow for account %s in slot %d: %d - %d", accountData.ID, diffSlot, accountData.FixedCost, diffChange.FixedCostChange)
 		}
 		accountData.FixedCost = iotago.Mana(fixedCost)
 		if diffChange.PrevLatestSupportedVersionAndHash != diffChange.NewLatestSupportedVersionAndHash {
@@ -435,7 +444,7 @@ func (m *Manager) rollbackAccountTo(accountData *accounts.AccountData, targetSlo
 		wasDestroyed = wasDestroyed || destroyed
 	}
 
-	return wasDestroyed, nil
+	return false, wasDestroyed, nil
 }
 
 func (m *Manager) preserveDestroyedAccountData(accountID iotago.AccountID) (accountDiff *model.AccountDiff, err error) {

pkg/protocol/engine/accounts/accountsledger/snapshot.go

@@ -36,6 +36,10 @@ func (m *Manager) Import(reader io.ReadSeeker) error {
 		return ierrors.Wrap(err, "unable to import slot diffs")
 	}
 
+	if err := m.accountsTree.Commit(); err != nil {
+		return ierrors.Wrap(err, "unable to commit account tree")
+	}
+
 	return nil
 }
 
@@ -73,14 +77,19 @@ func (m *Manager) exportAccountTree(writer io.WriteSeeker, targetIndex iotago.Sl
 	var accountCount int
 
 	if err := m.accountsTree.Stream(func(accountID iotago.AccountID, accountData *accounts.AccountData) error {
-		if _, err := m.rollbackAccountTo(accountData, targetIndex); err != nil {
+		wasCreatedAfterTargetSlot, _, err := m.rollbackAccountTo(accountData, targetIndex)
+		if err != nil {
 			return ierrors.Wrapf(err, "unable to rollback account %s", accountID)
 		}
 
+		// Account was created after the target slot, so we don't need to export it.
+		if wasCreatedAfterTargetSlot {
+			return nil
+		}
+
 		if err := stream.WriteObject(writer, accountData, (*accounts.AccountData).Bytes); err != nil {
 			return ierrors.Wrapf(err, "unable to write account %s", accountID)
 		}
-
 		accountCount++
 
 		return nil
@@ -105,7 +114,6 @@ func (m *Manager) recreateDestroyedAccounts(writer io.WriteSeeker, targetSlot io
 			accountData := accounts.NewAccountData(accountID)
 
 			destroyedAccounts[accountID] = accountData
-			recreatedAccountsCount++
 
 			return true
 		})
@@ -115,15 +123,20 @@ func (m *Manager) recreateDestroyedAccounts(writer io.WriteSeeker, targetSlot io
 	}
 
 	for accountID, accountData := range destroyedAccounts {
-		if wasDestroyed, err := m.rollbackAccountTo(accountData, targetSlot); err != nil {
+		if wasCreatedAfterTargetSlot, wasDestroyed, err := m.rollbackAccountTo(accountData, targetSlot); err != nil {
 			return 0, ierrors.Wrapf(err, "unable to rollback account %s to target slot %d", accountID, targetSlot)
+		} else if wasCreatedAfterTargetSlot {
+			// Account was created after the target slot, so we don't need to export it.
+			continue
 		} else if !wasDestroyed {
 			return 0, ierrors.Errorf("account %s was not destroyed", accountID)
 		}
 
 		if err := stream.WriteObject(writer, accountData, (*accounts.AccountData).Bytes); err != nil {
 			return 0, ierrors.Wrapf(err, "unable to write account %s", accountID)
 		}
+
+		recreatedAccountsCount++
 	}
 
 	return recreatedAccountsCount, nil

pkg/protocol/engine/accounts/accountsledger/snapshot_test.go

@@ -15,6 +15,11 @@ func TestManager_Import_Export(t *testing.T) {
 	ts := NewTestSuite(t)
 	latestSupportedVersionHash1 := tpkg.Rand32ByteArray()
 	latestSupportedVersionHash2 := tpkg.Rand32ByteArray()
+
+	accountTreeRoots := []iotago.Identifier{}
+
+	accountTreeRoots = append(accountTreeRoots, ts.Instance.AccountsTreeRoot())
+
 	ts.ApplySlotActions(1, 5, map[string]*AccountActions{
 		"A": {
 			TotalAllotments: 10,
@@ -44,6 +49,8 @@ func TestManager_Import_Export(t *testing.T) {
 		},
 	})
 
+	accountTreeRoots = append(accountTreeRoots, ts.Instance.AccountsTreeRoot())
+
 	ts.AssertAccountLedgerUntil(1, map[string]*AccountState{
 		"A": {
 			BICUpdatedTime:  1,
@@ -90,6 +97,8 @@ func TestManager_Import_Export(t *testing.T) {
 		},
 	})
 
+	accountTreeRoots = append(accountTreeRoots, ts.Instance.AccountsTreeRoot())
+
 	ts.AssertAccountLedgerUntil(2, map[string]*AccountState{
 		"A": {
 			BICUpdatedTime:  2,
@@ -145,6 +154,8 @@ func TestManager_Import_Export(t *testing.T) {
 		},
 	})
 
+	accountTreeRoots = append(accountTreeRoots, ts.Instance.AccountsTreeRoot())
+
 	ts.AssertAccountLedgerUntil(3, map[string]*AccountState{
 		"A": {
 			Destroyed: true,
@@ -178,31 +189,31 @@ func TestManager_Import_Export(t *testing.T) {
 
 	// Export and import the account ledger into new manager for the latest slot.
 	{
-		writer := stream.NewByteBuffer()
-
-		err := ts.Instance.Export(writer, iotago.SlotIndex(3))
-		require.NoError(t, err)
-
-		ts.Instance = ts.initAccountLedger()
-		err = ts.Instance.Import(writer.Reader())
-		require.NoError(t, err)
-		ts.Instance.SetLatestCommittedSlot(3)
-
-		ts.AssertAccountLedgerUntilWithoutNewState(3)
-	}
-
-	// Export and import for pre-latest slot.
-	{
-		writer := stream.NewByteBuffer()
-
-		err := ts.Instance.Export(writer, iotago.SlotIndex(2))
-		require.NoError(t, err)
-
-		ts.Instance = ts.initAccountLedger()
-		err = ts.Instance.Import(writer.Reader())
-		require.NoError(t, err)
-		ts.Instance.SetLatestCommittedSlot(2)
-
-		ts.AssertAccountLedgerUntilWithoutNewState(2)
+		writer := []*stream.ByteBuffer{
+			stream.NewByteBuffer(),
+			stream.NewByteBuffer(),
+			stream.NewByteBuffer(),
+			stream.NewByteBuffer(),
+		}
+
+		latestSlot := iotago.SlotIndex(3)
+
+		// Export snapshots at all slots including genesis.
+		for i := iotago.SlotIndex(0); i <= latestSlot; i++ {
+			err := ts.Instance.Export(writer[i], i)
+			require.NoError(t, err)
+		}
+
+		// Import all of the created snapshots into a new manager, assert the tree root and the states.
+		for i := iotago.SlotIndex(0); i <= latestSlot; i++ {
+			ts.Instance = ts.initAccountLedger()
+			err := ts.Instance.Import(writer[i].Reader())
+			require.NoError(t, err)
+			ts.Instance.SetLatestCommittedSlot(i)
+
+			ts.AssertAccountLedgerUntilWithoutNewState(i)
+
+			require.Equal(t, accountTreeRoots[i], ts.Instance.AccountsTreeRoot())
+		}
 	}
 }