Fix return loops instead of deleted keys

lib/runtime/interfaces.go

@@ -18,7 +18,7 @@ type Trie interface {
 	NextKey([]byte) []byte
 	ClearPrefix(prefix []byte) (err error)
 	ClearPrefixLimit(prefix []byte, limit uint32) (
-		deleted uint32, allDeleted bool, err error)
+		loops uint32, deleted uint32, allDeleted bool, err error)
 }
 
 // ChildTrie storage interface.S
@@ -31,7 +31,7 @@ type ChildTrie interface {
 		deleted uint32, allDeleted bool, err error)
 	ClearChildStorage(keyToChild, key []byte) error
 	ClearPrefixInChild(keyToChild, prefix []byte) error
-	ClearPrefixInChildWithLimit(keyToChild, prefix []byte, limit uint32) (uint32, bool, error)
+	ClearPrefixInChildWithLimit(keyToChild, prefix []byte, limit uint32) (uint32, uint32, bool, error)
 	GetChildNextKey(keyToChild, key []byte) ([]byte, error)
 }
 

lib/runtime/storage/trie.go

@@ -231,22 +231,24 @@ func (t *TrieState) ClearPrefix(prefix []byte) error {
 
 // ClearPrefixLimit deletes key-value pairs from the trie where the key starts with the given prefix till limit reached
 func (t *TrieState) ClearPrefixLimit(prefix []byte, limit uint32) (
-	deleted uint32, allDeleted bool, err error) {
+	loops uint32, deleted uint32, allDeleted bool, err error) {
 	t.mtx.Lock()
 	defer t.mtx.Unlock()
 
-	if currentTx := t.getCurrentTransaction(); currentTx != nil {
-		keysOnState := make([]string, 0)
+	keysOnState := make([]string, 0)
+	for key := range t.state.PrefixedKeys(prefix) {
+		keysOnState = append(keysOnState, string(key))
+	}
+	loops = uint32(len(keysOnState))
 
-		for key := range t.state.PrefixedKeys(prefix) {
-			keysOnState = append(keysOnState, string(key))
-		}
+	if currentTx := t.getCurrentTransaction(); currentTx != nil {
 
 		deleted, allDeleted = currentTx.clearPrefix(prefix, keysOnState, int(limit))
-		return deleted, allDeleted, nil
+		return loops, deleted, allDeleted, nil
 	}
 
-	return t.state.ClearPrefixLimit(prefix, limit)
+	deleted, allDeleted, err = t.state.ClearPrefixLimit(prefix, limit)
+	return loops, deleted, allDeleted, err
 }
 
 // TrieEntries returns every key-value pair in the trie
@@ -452,7 +454,7 @@ func (t *TrieState) ClearPrefixInChild(keyToChild, prefix []byte) error {
 	return nil
 }
 
-func (t *TrieState) ClearPrefixInChildWithLimit(keyToChild, prefix []byte, limit uint32) (uint32, bool, error) {
+func (t *TrieState) ClearPrefixInChildWithLimit(keyToChild, prefix []byte, limit uint32) (uint32, uint32, bool, error) {
 	t.mtx.Lock()
 	defer t.mtx.Unlock()
 
@@ -461,9 +463,9 @@ func (t *TrieState) ClearPrefixInChildWithLimit(keyToChild, prefix []byte, limit
 		if err != nil {
 			if errors.Is(err, trie.ErrChildTrieDoesNotExist) {
 				deleted, allDeleted := currentTx.clearPrefixInChild(string(keyToChild), prefix, []string{}, -1)
-				return deleted, allDeleted, nil
+				return 0, deleted, allDeleted, nil
 			}
-			return 0, false, err
+			return 0, 0, false, err
 		}
 
 		var onStateKeys []string
@@ -472,15 +474,16 @@ func (t *TrieState) ClearPrefixInChildWithLimit(keyToChild, prefix []byte, limit
 		}
 
 		deleted, allDeleted := currentTx.clearPrefixInChild(string(keyToChild), prefix, onStateKeys, int(limit))
-		return deleted, allDeleted, nil
+		return uint32(len(onStateKeys)), deleted, allDeleted, nil
 	}
 
 	child, err := t.state.GetChild(keyToChild)
 	if err != nil || child == nil {
-		return 0, false, err
+		return 0, 0, false, err
 	}
 
-	return child.ClearPrefixLimit(prefix, limit)
+	deleted, allDeleted, err := child.ClearPrefixLimit(prefix, limit)
+	return 0, deleted, allDeleted, err
 }
 
 // GetChildNextKey returns the next lexicographical larger key from child storage. If it does not exist, it returns nil.

lib/runtime/storage/trie_test.go

@@ -171,15 +171,17 @@ func TestTrieState_WithAndWithoutTransactions(t *testing.T) {
 				}
 			},
 			checks: func(t *testing.T, ts *TrieState, isTransactionRunning bool) {
-				deleted, allDeleted, err := ts.ClearPrefixLimit([]byte("noo"), uint32(1))
+				loops, deleted, allDeleted, err := ts.ClearPrefixLimit([]byte("noo"), uint32(1))
 				require.NoError(t, err)
 
 				if isTransactionRunning {
 					// New keys are not considered towards the limit
 					require.Equal(t, uint32(2), deleted)
+					require.Equal(t, uint32(0), loops)
 					require.False(t, allDeleted)
 				} else {
 					require.Equal(t, uint32(1), deleted)
+					require.Equal(t, uint32(2), loops)
 					require.False(t, allDeleted)
 				}
 			},
@@ -215,15 +217,17 @@ func TestTrieState_WithAndWithoutTransactions(t *testing.T) {
 
 			},
 			checks: func(t *testing.T, ts *TrieState, isTransactionRunning bool) {
-				deleted, allDeleted, err := ts.ClearPrefixInChildWithLimit(keyToChild, []byte("noo"), uint32(1))
+				loops, deleted, allDeleted, err := ts.ClearPrefixInChildWithLimit(keyToChild, []byte("noo"), uint32(1))
 
 				require.NoError(t, err)
 				require.False(t, allDeleted)
 
 				if isTransactionRunning {
 					require.Equal(t, uint32(2), deleted)
+					require.Equal(t, uint32(0), loops)
 				} else {
 					require.Equal(t, uint32(1), deleted)
+					require.Equal(t, uint32(0), loops)
 				}
 			},
 		},

lib/runtime/wazero/imports.go

@@ -38,10 +38,7 @@ var (
 	noneEncoded            []byte = []byte{0x00}
 	allZeroesBytes                = [32]byte{}
 
-	childStorageKeyPrefix        = []byte(":child_storage:")
-	imOnlinePalletProblematicKey = []byte{0x2b, 0x06, 0xaf, 0x97, 0x19, 0xac, 0x64, 0xd7, 0x55, 0x62, 0x3c, 0xda,
-		0x8d, 0xdd, 0x9b, 0x94, 0x4e, 0x7b, 0x90, 0x12, 0x09, 0x6b, 0x41, 0xc4, 0xeb, 0x3a, 0xaf,
-		0x94, 0x7f, 0x6e, 0xa4, 0x29}
+	childStorageKeyPrefix = []byte(":child_storage:")
 )
 
 const (
@@ -1152,13 +1149,13 @@ func ext_default_child_storage_clear_prefix_version_2(ctx context.Context, m api
 
 	limitUint := binary.LittleEndian.Uint32(limit)
 
-	deleted, allDeleted, err := storage.ClearPrefixInChildWithLimit(
+	loops, _, allDeleted, err := storage.ClearPrefixInChildWithLimit(
 		keyToChild, prefix, limitUint)
 	if err != nil {
 		logger.Errorf("failed to clear prefix in child with limit: %s", err)
 	}
 
-	killStorageResult := NewKillStorageResult(deleted, allDeleted)
+	killStorageResult := NewKillStorageResult(loops, allDeleted)
 
 	encodedKillStorageResult, err := scale.Marshal(killStorageResult)
 	if err != nil {
@@ -2138,13 +2135,13 @@ func ext_storage_clear_prefix_version_2(ctx context.Context, m api.Module, prefi
 		return valueSpan
 	}
 
-	numRemoved, all, err := storage.ClearPrefixLimit(prefix, *limitPtr)
+	loops, _, all, err := storage.ClearPrefixLimit(prefix, *limitPtr)
 	if err != nil {
 		logger.Errorf("failed to clear prefix limit: %s", err)
 		panic(err)
 	}
 
-	encBytes, err := toKillStorageResultEnum(all, numRemoved)
+	encBytes, err := toKillStorageResultEnum(all, loops)
 	if err != nil {
 		logger.Errorf("failed to allocate memory: %s", err)
 		panic(err)
@@ -2317,12 +2314,6 @@ func ext_storage_set_version_1(ctx context.Context, m api.Module, keySpan, value
 	storage := rtCtx.Storage
 
 	key := read(m, keySpan)
-
-	// TODO: temporal fix if key is pallet: "ImOnline" ++ ":__storage_version__:"
-	if bytes.Equal(key, imOnlinePalletProblematicKey) {
-		return
-	}
-
 	value := read(m, valueSpan)
 
 	cp := make([]byte, len(value))

lib/runtime/wazero/imports_test.go

@@ -1973,7 +1973,7 @@ func Test_ext_storage_clear_prefix_version_2(t *testing.T) {
 	var numDeleted uint32
 	// numDeleted represents no. of actual keys deleted
 	scale.Unmarshal(decVal[1:], &numDeleted)
-	require.Equal(t, uint32(2), numDeleted)
+	require.Equal(t, uint32(4), numDeleted)
 
 	var expectedAllDeleted byte
 	// expectedAllDeleted value 0 represents all keys deleted, 1 represents keys are pending with prefix in trie