Fix clear prefix

lib/runtime/storage/storagediff.go

@@ -12,7 +12,6 @@ import (
 	"golang.org/x/exp/slices"
 
 	"github.com/ChainSafe/gossamer/pkg/trie"
-	"github.com/ChainSafe/gossamer/pkg/trie/codec"
 )
 
 // storageDiff is a structure that stores the differences between consecutive
@@ -143,8 +142,6 @@ func (cs *storageDiff) clearPrefixInChild(keyToChild string, prefix []byte, chil
 // 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) {
-	prefix = codec.KeyLEToNibbles(prefix)
-	prefix = bytes.TrimSuffix(prefix, []byte{0})
 	newKeys := maps.Keys(cs.upserts)
 	allKeys := append(newKeys, trieKeys...)
 	deleted = 0
@@ -154,11 +151,12 @@ func (cs *storageDiff) clearPrefix(prefix []byte, trieKeys []string, limit int)
 			break
 		}
 		keyBytes := []byte(k)
-		bytes.HasPrefix(keyBytes, prefix)
-		cs.delete(k)
-		deleted++
-		if !slices.Contains(newKeys, k) {
-			limit--
+		if bytes.HasPrefix(keyBytes, prefix) {
+			cs.delete(k)
+			deleted++
+			if !slices.Contains(newKeys, k) {
+				limit--
+			}
 		}
 	}
 

lib/runtime/storage/storagediff_test.go

@@ -9,7 +9,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestStorageDiff_Get(t *testing.T) {
+func TestStorageDiff_MainTrie(t *testing.T) {
 	t.Run("Upsert", func(t *testing.T) {
 		changes := newStorageDiff()
 
@@ -22,7 +22,7 @@ func TestStorageDiff_Get(t *testing.T) {
 		require.Equal(t, value, val)
 	})
 
-	t.Run("Upsert then delete", func(t *testing.T) {
+	t.Run("Delete", func(t *testing.T) {
 		changes := newStorageDiff()
 
 		key := "key"
@@ -34,4 +34,120 @@ func TestStorageDiff_Get(t *testing.T) {
 		require.True(t, deleted)
 		require.Nil(t, val)
 	})
+
+	t.Run("clearPrefix", func(t *testing.T) {
+		testEntries := map[string][]byte{
+			"pre":        []byte("pre"),
+			"predict":    []byte("predict"),
+			"prediction": []byte("prediction"),
+		}
+
+		commonPrefix := []byte("pre")
+
+		cases := map[string]struct {
+			prefix    []byte
+			limit     int
+			trieKeys  []string
+			deleted   uint32
+			allDelted bool
+		}{
+			"empty_trie_limit_1": {
+				prefix:    commonPrefix,
+				limit:     1,
+				deleted:   3, // Since keys during block exec does not count
+				allDelted: true,
+			},
+			"empty_trie_limit_2": {
+				prefix:    commonPrefix,
+				limit:     2,
+				deleted:   3, // Since keys during block exec does not count
+				allDelted: true,
+			},
+			"empty_trie_same_limit_than_stored_keys": {
+				prefix:    commonPrefix,
+				limit:     3,
+				deleted:   3,
+				allDelted: true,
+			},
+			"empty_trie_no_limit": {
+				prefix:    commonPrefix,
+				limit:     -1,
+				deleted:   3,
+				allDelted: true,
+			},
+			"with_previous_state_not_sharing_prefix_limit_1": {
+				prefix:    commonPrefix,
+				limit:     1,
+				trieKeys:  []string{"bio"},
+				deleted:   3, // Since keys during block exec does not count
+				allDelted: false,
+			},
+			"with_previous_state_not_sharing_prefix_limit_2": {
+				prefix:    commonPrefix,
+				limit:     2,
+				trieKeys:  []string{"bio"},
+				deleted:   3, // Since keys during block exec does not count
+				allDelted: false,
+			},
+			"with_previous_state_not_sharing_prefix_limit_3": {
+				prefix:    commonPrefix,
+				limit:     3,
+				trieKeys:  []string{"bio"},
+				deleted:   3,
+				allDelted: false,
+			},
+			"with_previous_state_not_sharing_prefix_with_no_limit": {
+				prefix:    commonPrefix,
+				limit:     -1,
+				trieKeys:  []string{"bio"},
+				deleted:   3,
+				allDelted: false,
+			},
+			"with_previous_state_sharing_prefix_limit_1": {
+				prefix:    []byte("p"),
+				limit:     1,
+				trieKeys:  []string{"p"},
+				deleted:   1, // the "p" key only
+				allDelted: false,
+			},
+			"with_previous_state_sharing_prefix_limit_2": {
+				prefix:    []byte("p"),
+				limit:     2,
+				trieKeys:  []string{"p"},
+				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"},
+				deleted:   4,
+				allDelted: true,
+			},
+			"with_previous_state_sharing_prefix_with_no_limit": {
+				prefix:    []byte("p"),
+				limit:     -1,
+				trieKeys:  []string{"p"},
+				deleted:   4,
+				allDelted: true,
+			},
+		}
+
+		for tname, tt := range cases {
+			tt := tt
+			t.Run(tname, func(t *testing.T) {
+				t.Parallel()
+
+				changes := newStorageDiff()
+
+				for k, v := range testEntries {
+					changes.upsert(k, v)
+				}
+
+				deleted, allDeleted := changes.clearPrefix(tt.prefix, tt.trieKeys, tt.limit)
+				require.Equal(t, tt.deleted, deleted)
+				require.Equal(t, tt.allDelted, allDeleted)
+			})
+		}
+	})
 }