Revert tx that migrates data bugfix

state/accountsDB_test.go

@@ -3065,6 +3065,85 @@ func TestAccountsDB_SaveKeyValAfterAccountIsReverted(t *testing.T) {
 	require.NotNil(t, acc)
 }
 
+func TestAccountsDB_RevertTxWhichMigratesDataRemovesMigratedData(t *testing.T) {
+	t.Parallel()
+
+	marshaller := &marshallerMock.MarshalizerMock{}
+	hasher := &hashingMocks.HasherMock{}
+	enableEpochsHandler := &enableEpochsHandlerMock.EnableEpochsHandlerStub{
+		IsAutoBalanceDataTriesEnabledField: false,
+	}
+	tsm, _ := trie.NewTrieStorageManager(storage.GetStorageManagerArgs())
+	tr, _ := trie.NewTrie(tsm, marshaller, hasher, enableEpochsHandler, uint(5))
+	spm := &stateMock.StoragePruningManagerStub{}
+	argsAccountsDB := createMockAccountsDBArgs()
+	argsAccountsDB.Trie = tr
+	argsAccountsDB.Hasher = hasher
+	argsAccountsDB.Marshaller = marshaller
+	argsAccCreator := factory.ArgsAccountCreator{
+		Hasher:              hasher,
+		Marshaller:          marshaller,
+		EnableEpochsHandler: enableEpochsHandler,
+	}
+	argsAccountsDB.AccountFactory, _ = factory.NewAccountCreator(argsAccCreator)
+	argsAccountsDB.StoragePruningManager = spm
+	adb, _ := state.NewAccountsDB(argsAccountsDB)
+
+	address := make([]byte, 32)
+	acc, err := adb.LoadAccount(address)
+	require.Nil(t, err)
+
+	// save account with data trie that is not migrated
+	userAcc := acc.(state.UserAccountHandler)
+	key := []byte("key")
+	err = userAcc.SaveKeyValue(key, []byte("value"))
+	require.Nil(t, err)
+	err = userAcc.SaveKeyValue([]byte("key1"), []byte("value"))
+	require.Nil(t, err)
+
+	err = adb.SaveAccount(userAcc)
+	userAccRootHash := userAcc.GetRootHash()
+	require.Nil(t, err)
+	_, err = adb.Commit()
+	require.Nil(t, err)
+
+	enableEpochsHandler.IsAutoBalanceDataTriesEnabledField = true
+
+	// a JournalEntry is needed so the revert can happen at snapshot 1. Creating a new account creates a new journal entry.
+	newAcc, _ := adb.LoadAccount(generateRandomByteArray(32))
+	_ = adb.SaveAccount(newAcc)
+	assert.Equal(t, 1, adb.JournalLen())
+
+	// change the account data trie. This will trigger the migration
+	acc, err = adb.LoadAccount(address)
+	require.Nil(t, err)
+	userAcc = acc.(state.UserAccountHandler)
+	value1 := []byte("value1")
+	err = userAcc.SaveKeyValue(key, value1)
+	require.Nil(t, err)
+	err = adb.SaveAccount(userAcc)
+	require.Nil(t, err)
+
+	// revert the migration
+	err = adb.RevertToSnapshot(1)
+	require.Nil(t, err)
+
+	// check that the data trie was completely reverted. The rootHash of the user account should be present
+	// in both old and new hashes. This means that after the revert, the rootHash is the same as before
+	markForEvictionCalled := false
+	spm.MarkForEvictionCalled = func(_ []byte, _ []byte, oldHashes common.ModifiedHashes, newHashes common.ModifiedHashes) error {
+		_, ok := oldHashes[string(userAccRootHash)]
+		require.True(t, ok)
+		_, ok = newHashes[string(userAccRootHash)]
+		require.True(t, ok)
+		markForEvictionCalled = true
+
+		return nil
+	}
+	_, _ = adb.Commit()
+	require.True(t, markForEvictionCalled)
+}
+
 func testAccountMethodsConcurrency(
 	t *testing.T,
 	adb state.AccountsAdapter,

state/trackableDataTrie/trackableDataTrie.go

@@ -234,12 +234,20 @@ func (tdt *trackableDataTrie) updateTrie(dtr state.DataTrie) ([]core.TrieData, e
 			return nil, err
 		}
 
-		err = tdt.modifyTrie([]byte(key), dataEntry, oldVal, dtr)
+		newKey, err := tdt.modifyTrie([]byte(key), dataEntry, oldVal, dtr)
 		if err != nil {
 			return nil, err
 		}
 
 		index++
+
+		isFirstMigration := oldVal.Version == core.NotSpecified && dataEntry.newVersion == core.AutoBalanceEnabled
+		if isFirstMigration && len(newKey) != 0 {
+			oldValues = append(oldValues, core.TrieData{
+				Key:   newKey,
+				Value: nil,
+			})
+		}
 	}
 
 	tdt.dirtyData = make(map[string]dirtyData)
@@ -327,25 +335,26 @@ func (tdt *trackableDataTrie) deleteOldEntryIfMigrated(key []byte, newData dirty
 
 	isMigration := oldEntry.Version == core.NotSpecified && newData.newVersion == core.AutoBalanceEnabled
 	if isMigration && len(newData.value) != 0 {
+		log.Trace("delete old entry if migrated", "key", key)
 		return tdt.tr.Delete(key)
 	}
 
 	return nil
 }
 
-func (tdt *trackableDataTrie) modifyTrie(key []byte, dataEntry dirtyData, oldVal core.TrieData, dtr state.DataTrie) error {
+func (tdt *trackableDataTrie) modifyTrie(key []byte, dataEntry dirtyData, oldVal core.TrieData, dtr state.DataTrie) ([]byte, error) {
 	if len(dataEntry.value) == 0 {
-		return tdt.deleteFromTrie(oldVal, key, dtr)
+		return nil, tdt.deleteFromTrie(oldVal, key, dtr)
 	}
 
 	version := dataEntry.newVersion
 	newKey := tdt.getKeyForVersion(key, version)
 	value, err := tdt.getValueForVersion(key, dataEntry.value, version)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	return dtr.UpdateWithVersion(newKey, value, version)
+	return newKey, dtr.UpdateWithVersion(newKey, value, version)
 }
 
 func (tdt *trackableDataTrie) deleteFromTrie(oldVal core.TrieData, key []byte, dtr state.DataTrie) error {

state/trackableDataTrie/trackableDataTrie_test.go

@@ -418,9 +418,11 @@ func TestTrackableDataTrie_SaveDirtyData(t *testing.T) {
 		_ = tdt.SaveKeyValue(expectedKey, expectedVal)
 		oldValues, err := tdt.SaveDirtyData(trie)
 		assert.Nil(t, err)
-		assert.Equal(t, 1, len(oldValues))
+		assert.Equal(t, 2, len(oldValues))
 		assert.Equal(t, expectedKey, oldValues[0].Key)
 		assert.Equal(t, value, oldValues[0].Value)
+		assert.Equal(t, hasher.Compute(string(expectedKey)), oldValues[1].Key)
+		assert.Equal(t, []byte(nil), oldValues[1].Value)
 		assert.True(t, deleteCalled)
 		assert.True(t, updateCalled)
 	})

trie/branchNode.go

@@ -639,6 +639,8 @@ func (bn *branchNode) setNewChild(childPos byte, newNode node) error {
 		bn.setVersionForChild(0, childPos)
 		bn.EncodedChildren[childPos] = nil
 
+		bn.revertChildrenVersionSlice()
+
 		return nil
 	}
 
@@ -647,10 +649,23 @@ func (bn *branchNode) setNewChild(childPos byte, newNode node) error {
 		return err
 	}
 	bn.setVersionForChild(childVersion, childPos)
+	if childVersion == core.NotSpecified {
+		bn.revertChildrenVersionSlice()
+	}
 
 	return nil
 }
 
+func (bn *branchNode) revertChildrenVersionSlice() {
+	for i := range bn.ChildrenVersion {
+		if bn.ChildrenVersion[i] == 1 {
+			return
+		}
+	}
+
+	bn.ChildrenVersion = []byte(nil)
+}
+
 func (bn *branchNode) reduceNode(pos int) (node, bool, error) {
 	newEn, err := newExtensionNode([]byte{byte(pos)}, bn, bn.marsh, bn.hasher)
 	if err != nil {