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/storagediff.go

@@ -126,21 +126,22 @@ func (cs *storageDiff) deleteChildLimit(keyToChild string,
 
 // clearPrefixInChild clears keys with a specific prefix within a child trie.
 func (cs *storageDiff) clearPrefixInChild(keyToChild string, prefix []byte,
-	childKeys []string, limit int) (deleted uint32, allDeleted bool) {
+	childKeys []string, limit int) (loops, deleted uint32, allDeleted bool) {
 	childChanges := cs.childChangeSet[keyToChild]
 	if childChanges == nil {
 		childChanges = newStorageDiff()
 	}
-	deleted, allDeleted = childChanges.clearPrefix(prefix, childKeys, limit)
+	loops, deleted, allDeleted = childChanges.clearPrefix(prefix, childKeys, limit)
 	cs.childChangeSet[keyToChild] = childChanges
 
-	return deleted, allDeleted
+	return loops, deleted, allDeleted
 }
 
 // clearPrefix removes all keys matching a specified prefix, within an
 // optional limit. It returns the number of keys deleted and a boolean
 // indicating if all keys with the prefix were removed.
-func (cs *storageDiff) clearPrefix(prefix []byte, trieKeys []string, limit int) (deleted uint32, allDeleted bool) {
+func (cs *storageDiff) clearPrefix(prefix []byte, trieKeys []string, limit int) (
+	loops, deleted uint32, allDeleted bool) {
 	newKeys := maps.Keys(cs.upserts)
 	keysToClear := maps.Keys(cs.upserts)
 	for _, k := range trieKeys {
@@ -152,20 +153,20 @@ func (cs *storageDiff) clearPrefix(prefix []byte, trieKeys []string, limit int)
 	deleted = 0
 	sort.Strings(keysToClear)
 	for _, k := range keysToClear {
-		if limit == 0 {
+		if loops == uint32(limit) {
 			break
 		}
 		keyBytes := []byte(k)
 		if bytes.HasPrefix(keyBytes, prefix) {
 			cs.delete(k)
 			deleted++
 			if !slices.Contains(newKeys, k) {
-				limit--
+				loops++
 			}
 		}
 	}
 
-	return deleted, deleted == uint32(len(keysToClear))
+	return loops, deleted, deleted == uint32(len(keysToClear))
 }
 
 // getFromChild attempts to retrieve a value associated with a specific key

lib/runtime/storage/storagediff_test.go

@@ -89,6 +89,7 @@ func Test_MainTrie(t *testing.T) {
 			prefix    []byte
 			limit     int
 			trieKeys  []string
+			loops     uint32
 			deleted   uint32
 			allDelted bool
 		}{
@@ -148,27 +149,31 @@ func Test_MainTrie(t *testing.T) {
 				prefix:    []byte("p"),
 				limit:     1,
 				trieKeys:  []string{"p"},
+				loops:     1,
 				deleted:   1, // the "p" key only
 				allDelted: false,
 			},
 			"with_previous_state_sharing_prefix_limit_2": {
 				prefix:    []byte("p"),
 				limit:     2,
 				trieKeys:  []string{"p"},
+				loops:     1,
 				deleted:   4, // Since keys during block exec does not count
 				allDelted: true,
 			},
 			"with_previous_state_sharing_prefix_limit_3": {
 				prefix:    []byte("p"),
 				limit:     3,
 				trieKeys:  []string{"p"},
+				loops:     1,
 				deleted:   4,
 				allDelted: true,
 			},
 			"with_previous_state_sharing_prefix_with_no_limit": {
 				prefix:    []byte("p"),
 				limit:     -1,
 				trieKeys:  []string{"p"},
+				loops:     1,
 				deleted:   4,
 				allDelted: true,
 			},
@@ -185,8 +190,9 @@ func Test_MainTrie(t *testing.T) {
 					changes.upsert(k, v)
 				}
 
-				deleted, allDeleted := changes.clearPrefix(tt.prefix, tt.trieKeys, tt.limit)
+				loops, deleted, allDeleted := changes.clearPrefix(tt.prefix, tt.trieKeys, tt.limit)
 				require.Equal(t, tt.deleted, deleted)
+				require.Equal(t, tt.loops, loops)
 				require.Equal(t, tt.allDelted, allDeleted)
 			})
 		}
@@ -365,6 +371,7 @@ func Test_ChildTrie(t *testing.T) {
 			prefix    []byte
 			limit     int
 			trieKeys  []string
+			loops     uint32
 			deleted   uint32
 			allDelted bool
 		}{
@@ -424,27 +431,31 @@ func Test_ChildTrie(t *testing.T) {
 				prefix:    []byte("p"),
 				limit:     1,
 				trieKeys:  []string{"p"},
+				loops:     1,
 				deleted:   1, // the "p" key only
 				allDelted: false,
 			},
 			"with_previous_state_sharing_prefix_limit_2": {
 				prefix:    []byte("p"),
 				limit:     2,
 				trieKeys:  []string{"p"},
+				loops:     1,
 				deleted:   4, // Since keys during block exec does not count
 				allDelted: true,
 			},
 			"with_previous_state_sharing_prefix_limit_3": {
 				prefix:    []byte("p"),
 				limit:     3,
 				trieKeys:  []string{"p"},
+				loops:     1,
 				deleted:   4,
 				allDelted: true,
 			},
 			"with_previous_state_sharing_prefix_with_no_limit": {
 				prefix:    []byte("p"),
 				limit:     -1,
 				trieKeys:  []string{"p"},
+				loops:     1,
 				deleted:   4,
 				allDelted: true,
 			},
@@ -461,8 +472,9 @@ func Test_ChildTrie(t *testing.T) {
 					changes.upsertChild("child", k, v)
 				}
 
-				deleted, allDeleted := changes.clearPrefixInChild("child", tt.prefix, tt.trieKeys, tt.limit)
+				loops, deleted, allDeleted := changes.clearPrefixInChild("child", tt.prefix, tt.trieKeys, tt.limit)
 				require.Equal(t, tt.deleted, deleted)
+				require.Equal(t, tt.loops, loops)
 				require.Equal(t, tt.allDelted, allDeleted)
 			})
 		}

lib/runtime/storage/trie.go

@@ -231,22 +231,22 @@ 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)
-
 		for key := range t.state.PrefixedKeys(prefix) {
 			keysOnState = append(keysOnState, string(key))
 		}
 
-		deleted, allDeleted = currentTx.clearPrefix(prefix, keysOnState, int(limit))
-		return deleted, allDeleted, nil
+		loops, deleted, allDeleted = currentTx.clearPrefix(prefix, keysOnState, int(limit))
+		return loops, deleted, allDeleted, nil
 	}
 
-	return t.state.ClearPrefixLimit(prefix, limit)
+	deleted, allDeleted, err = t.state.ClearPrefixLimit(prefix, limit)
+	return 0, deleted, allDeleted, err
 }
 
 // TrieEntries returns every key-value pair in the trie
@@ -452,35 +452,36 @@ 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()
 
 	if currentTx := t.getCurrentTransaction(); currentTx != nil {
 		child, err := t.state.GetChild(keyToChild)
 		if err != nil {
 			if errors.Is(err, trie.ErrChildTrieDoesNotExist) {
-				deleted, allDeleted := currentTx.clearPrefixInChild(string(keyToChild), prefix, []string{}, -1)
-				return deleted, allDeleted, nil
+				loops, deleted, allDeleted := currentTx.clearPrefixInChild(string(keyToChild), prefix, []string{}, -1)
+				return loops, deleted, allDeleted, nil
 			}
-			return 0, false, err
+			return 0, 0, false, err
 		}
 
 		var onStateKeys []string
 		for key := range child.PrefixedKeys(prefix) {
 			onStateKeys = append(onStateKeys, string(key))
 		}
 
-		deleted, allDeleted := currentTx.clearPrefixInChild(string(keyToChild), prefix, onStateKeys, int(limit))
-		return deleted, allDeleted, nil
+		loops, deleted, allDeleted := currentTx.clearPrefixInChild(string(keyToChild), prefix, onStateKeys, int(limit))
+		return loops, 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(0), 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

@@ -1970,10 +1970,10 @@ func Test_ext_storage_clear_prefix_version_2(t *testing.T) {
 	var decVal []byte
 	scale.Unmarshal(encValue, &decVal)
 
-	var numDeleted uint32
-	// numDeleted represents no. of actual keys deleted
-	scale.Unmarshal(decVal[1:], &numDeleted)
-	require.Equal(t, uint32(2), numDeleted)
+	var loops uint32
+	// loops represents no. of loops when deleting keys
+	scale.Unmarshal(decVal[1:], &loops)
+	require.Equal(t, uint32(0), loops)
 
 	var expectedAllDeleted byte
 	// expectedAllDeleted value 0 represents all keys deleted, 1 represents keys are pending with prefix in trie
@@ -1992,8 +1992,8 @@ func Test_ext_storage_clear_prefix_version_2(t *testing.T) {
 	require.NoError(t, err)
 
 	scale.Unmarshal(encValue, &decVal)
-	scale.Unmarshal(decVal[1:], &numDeleted)
-	require.Equal(t, uint32(2), numDeleted)
+	scale.Unmarshal(decVal[1:], &loops)
+	require.Equal(t, uint32(0), loops)
 
 	expectedAllDeleted = 0
 	require.Equal(t, expectedAllDeleted, decVal[0])