feat(lib/trie): Implement limit for trie.ClearPrefix (ChainSafe#1905)

* Implement limit for trie.ClearPrefix

lib/runtime/constants.go

@@ -44,10 +44,10 @@ const (
 	POLKADOT_RUNTIME_FP  = "polkadot_runtime.compact.wasm"
 	POLKADOT_RUNTIME_URL = "https://github.com/noot/polkadot/blob/noot/v0.8.25/polkadot_runtime.wasm?raw=true"
 
-	// v0.8 test API wasm
+	// v0.9 test API wasm
 	HOST_API_TEST_RUNTIME     = "hostapi_runtime"
 	HOST_API_TEST_RUNTIME_FP  = "hostapi_runtime.compact.wasm"
-	HOST_API_TEST_RUNTIME_URL = "https://github.com/ChainSafe/polkadot-spec/blob/80fa2be272820731b5159e9dc2a3eec3cca02b4d/test/hostapi_runtime.compact.wasm?raw=true"
+	HOST_API_TEST_RUNTIME_URL = "https://github.com/ChainSafe/polkadot-spec/blob/9cc27bf7b7f21c106000103f8f6b6c51f7fb8353/test/runtimes/hostapi/hostapi_runtime.compact.wasm?raw=true"
 
 	// v0.8 substrate runtime with modified name and babe C=(1, 1)
 	DEV_RUNTIME     = "dev_runtime"

lib/runtime/interface.go

@@ -77,6 +77,7 @@ type Storage interface {
 	GetChildNextKey(keyToChild, key []byte) ([]byte, error)
 	GetChild(keyToChild []byte) (*trie.Trie, error)
 	ClearPrefix(prefix []byte) error
+	ClearPrefixLimit(prefix []byte, limit uint32) (uint32, bool)
 	BeginStorageTransaction()
 	CommitStorageTransaction()
 	RollbackStorageTransaction()

lib/runtime/storage/trie.go

@@ -137,6 +137,15 @@ func (s *TrieState) ClearPrefix(prefix []byte) error {
 	return nil
 }
 
+// ClearPrefixLimit deletes key-value pairs from the trie where the key starts with the given prefix till limit reached
+func (s *TrieState) ClearPrefixLimit(prefix []byte, limit uint32) (uint32, bool) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	num, del := s.t.ClearPrefixLimit(prefix, limit)
+	return num, del
+}
+
 // TrieEntries returns every key-value pair in the trie
 func (s *TrieState) TrieEntries() map[string][]byte {
 	s.lock.RLock()

lib/runtime/wasmer/imports.go

@@ -1785,24 +1785,44 @@ func ext_storage_clear_prefix_version_1(context unsafe.Pointer, prefixSpan C.int
 }
 
 //export ext_storage_clear_prefix_version_2
-func ext_storage_clear_prefix_version_2(context unsafe.Pointer, prefixSpan, _ C.int64_t) C.int64_t {
+func ext_storage_clear_prefix_version_2(context unsafe.Pointer, prefixSpan, lim C.int64_t) C.int64_t {
 	logger.Trace("[ext_storage_clear_prefix_version_2] executing...")
-	logger.Warn("[ext_storage_clear_prefix_version_2] somewhat unimplemented")
-	// TODO: need to use unused `limit` parameter (#1792)
 
 	instanceContext := wasm.IntoInstanceContext(context)
 	ctx := instanceContext.Data().(*runtime.Context)
 	storage := ctx.Storage
 
 	prefix := asMemorySlice(instanceContext, prefixSpan)
-	logger.Debug("[ext_storage_clear_prefix_version_1]", "prefix", fmt.Sprintf("0x%x", prefix))
+	logger.Debug("[ext_storage_clear_prefix_version_2]", "prefix", fmt.Sprintf("0x%x", prefix))
 
-	err := storage.ClearPrefix(prefix)
+	limitBytes := asMemorySlice(instanceContext, lim)
+
+	var limit []byte
+	err := scale.Unmarshal(limitBytes, &limit)
 	if err != nil {
-		logger.Error("[ext_storage_clear_prefix_version_1]", "error", err)
+		logger.Warn("[ext_storage_clear_prefix_version_2] cannot generate limit", "error", err)
+		ret, _ := toWasmMemory(instanceContext, nil)
+		return C.int64_t(ret)
 	}
 
-	return 1
+	limitUint := binary.LittleEndian.Uint32(limit)
+	numRemoved, all := storage.ClearPrefixLimit(prefix, limitUint)
+	encBytes, err := toKillStorageResultEnum(all, numRemoved)
+
+	if err != nil {
+		logger.Error("[ext_storage_clear_prefix_version_2] failed to allocate memory", err)
+		ret, _ := toWasmMemory(instanceContext, nil)
+		return C.int64_t(ret)
+	}
+
+	valueSpan, err := toWasmMemory(instanceContext, encBytes)
+	if err != nil {
+		logger.Error("[ext_storage_clear_prefix_version_2] failed to allocate", "error", err)
+		ptr, _ := toWasmMemory(instanceContext, nil)
+		return C.int64_t(ptr)
+	}
+
+	return C.int64_t(valueSpan)
 }
 
 //export ext_storage_exists_version_1
@@ -2059,6 +2079,27 @@ func toWasmMemoryOptionalUint32(context wasm.InstanceContext, data *uint32) (int
 	return toWasmMemory(context, enc)
 }
 
+// toKillStorageResult returns enum encoded value
+func toKillStorageResultEnum(allRemoved bool, numRemoved uint32) ([]byte, error) {
+	var b, sbytes []byte
+	sbytes, err := scale.Marshal(numRemoved)
+	if err != nil {
+		return nil, err
+	}
+
+	if allRemoved {
+		// No key remains in the child trie.
+		b = append(b, byte(0))
+	} else {
+		// At least one key still resides in the child trie due to the supplied limit.
+		b = append(b, byte(1))
+	}
+
+	b = append(b, sbytes...)
+
+	return b, err
+}
+
 // Wraps slice in optional.FixedSizeBytes and copies result to wasm memory. Returns resulting 64bit span descriptor
 func toWasmMemoryFixedSizeOptional(context wasm.InstanceContext, data []byte) (int64, error) {
 	var opt [64]byte

lib/runtime/wasmer/imports_test.go

@@ -23,6 +23,11 @@ import (
 	"sort"
 	"testing"
 
+	log "github.com/ChainSafe/log15"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/wasmerio/go-ext-wasm/wasmer"
+
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/common/types"
 	"github.com/ChainSafe/gossamer/lib/crypto"
@@ -34,10 +39,6 @@ import (
 	"github.com/ChainSafe/gossamer/lib/runtime/storage"
 	"github.com/ChainSafe/gossamer/lib/trie"
 	"github.com/ChainSafe/gossamer/pkg/scale"
-	log "github.com/ChainSafe/log15"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-	"github.com/wasmerio/go-ext-wasm/wasmer"
 )
 
 var testChildKey = []byte("childKey")
@@ -275,6 +276,78 @@ func Test_ext_storage_clear_prefix_version_1(t *testing.T) {
 	require.NotNil(t, val)
 }
 
+func Test_ext_storage_clear_prefix_version_2(t *testing.T) {
+	inst := NewTestInstance(t, runtime.HOST_API_TEST_RUNTIME)
+
+	testkey := []byte("noot")
+	inst.ctx.Storage.Set(testkey, []byte{1})
+
+	testkey2 := []byte("noot1")
+	inst.ctx.Storage.Set(testkey2, []byte{1})
+
+	testkey3 := []byte("noot2")
+	inst.ctx.Storage.Set(testkey3, []byte{1})
+
+	testkey4 := []byte("noot3")
+	inst.ctx.Storage.Set(testkey4, []byte{1})
+
+	testkey5 := []byte("spaghet")
+	testValue5 := []byte{2}
+	inst.ctx.Storage.Set(testkey5, testValue5)
+
+	enc, err := scale.Marshal(testkey[:3])
+	require.NoError(t, err)
+
+	testLimit := uint32(2)
+	testLimitBytes := make([]byte, 4)
+	binary.LittleEndian.PutUint32(testLimitBytes, testLimit)
+
+	optLimit, err := scale.Marshal(&testLimitBytes)
+	require.NoError(t, err)
+
+	// clearing prefix for "noo" prefix with limit 2
+	encValue, err := inst.Exec("rtm_ext_storage_clear_prefix_version_2", append(enc, optLimit...))
+	require.NoError(t, err)
+
+	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 expectedAllDeleted byte
+	// expectedAllDeleted value 0 represents all keys deleted, 1 represents keys are pending with prefix in trie
+	expectedAllDeleted = 1
+	require.Equal(t, expectedAllDeleted, decVal[0])
+
+	val := inst.ctx.Storage.Get(testkey)
+	require.NotNil(t, val)
+
+	val = inst.ctx.Storage.Get(testkey5)
+	require.NotNil(t, val)
+	require.Equal(t, testValue5, val)
+
+	// clearing prefix again for "noo" prefix with limit 2
+	encValue, err = inst.Exec("rtm_ext_storage_clear_prefix_version_2", append(enc, optLimit...))
+	require.NoError(t, err)
+
+	scale.Unmarshal(encValue, &decVal)
+	scale.Unmarshal(decVal[1:], &numDeleted)
+	require.Equal(t, uint32(2), numDeleted)
+
+	expectedAllDeleted = 0
+	require.Equal(t, expectedAllDeleted, decVal[0])
+
+	val = inst.ctx.Storage.Get(testkey)
+	require.Nil(t, val)
+
+	val = inst.ctx.Storage.Get(testkey5)
+	require.NotNil(t, val)
+	require.Equal(t, testValue5, val)
+}
+
 func Test_ext_storage_get_version_1(t *testing.T) {
 	inst := NewTestInstance(t, runtime.HOST_API_TEST_RUNTIME)
 

lib/trie/trie.go

@@ -18,6 +18,7 @@ package trie
 
 import (
 	"bytes"
+	"fmt"
 
 	"github.com/ChainSafe/gossamer/lib/common"
 )
@@ -509,6 +510,134 @@ func (t *Trie) retrieve(parent node, key []byte) *leaf {
 	return value
 }
 
+// ClearPrefixLimit deletes the keys having the prefix till limit reached
+func (t *Trie) ClearPrefixLimit(prefix []byte, limit uint32) (uint32, bool) {
+	if limit == 0 {
+		return 0, false
+	}
+
+	p := keyToNibbles(prefix)
+	if len(p) > 0 && p[len(p)-1] == 0 {
+		p = p[:len(p)-1]
+	}
+
+	l := limit
+	var allDeleted bool
+	t.root, _, allDeleted = t.clearPrefixLimit(t.root, p, &limit)
+	return l - limit, allDeleted
+}
+
+// clearPrefixLimit deletes the keys having the prefix till limit reached and returns updated trie root node,
+// true if any node in the trie got updated, and next bool returns true if there is no keys left with prefix.
+func (t *Trie) clearPrefixLimit(cn node, prefix []byte, limit *uint32) (node, bool, bool) {
+	curr := t.maybeUpdateGeneration(cn)
+
+	switch c := curr.(type) {
+	case *branch:
+		length := lenCommonPrefix(c.key, prefix)
+		if length == len(prefix) {
+			n, _ := t.deleteNodes(c, []byte{}, limit)
+			if n == nil {
+				return nil, true, true
+			}
+			return n, true, false
+		}
+
+		if len(prefix) == len(c.key)+1 && length == len(prefix)-1 {
+			i := prefix[len(c.key)]
+			c.children[i], _ = t.deleteNodes(c.children[i], []byte{}, limit)
+
+			c.setDirty(true)
+			curr = handleDeletion(c, prefix)
+
+			if c.children[i] == nil {
+				return curr, true, true
+			}
+			return c, true, false
+		}
+
+		if len(prefix) <= len(c.key) || length < len(c.key) {
+			// this node doesn't have the prefix, return
+			return c, false, true
+		}
+
+		i := prefix[len(c.key)]
+
+		var wasUpdated, allDeleted bool
+		c.children[i], wasUpdated, allDeleted = t.clearPrefixLimit(c.children[i], prefix[len(c.key)+1:], limit)
+		if wasUpdated {
+			c.setDirty(true)
+			curr = handleDeletion(c, prefix)
+		}
+
+		return curr, curr.isDirty(), allDeleted
+	case *leaf:
+		length := lenCommonPrefix(c.key, prefix)
+		if length == len(prefix) {
+			*limit--
+			return nil, true, true
+		}
+		// Prefix not found might be all deleted
+		return curr, false, true
+
+	case nil:
+		return nil, false, true
+	}
+
+	return nil, false, true
+}
+
+func (t *Trie) deleteNodes(cn node, prefix []byte, limit *uint32) (node, bool) {
+	curr := t.maybeUpdateGeneration(cn)
+
+	switch c := curr.(type) {
+	case *leaf:
+		if *limit == 0 {
+			return c, false
+		}
+		*limit--
+		return nil, true
+	case *branch:
+		if len(c.key) != 0 {
+			prefix = append(prefix, c.key...)
+		}
+
+		for i, child := range c.children {
+			if child == nil {
+				continue
+			}
+
+			var isDel bool
+			if c.children[i], isDel = t.deleteNodes(child, prefix, limit); !isDel {
+				continue
+			}
+
+			c.setDirty(true)
+			curr = handleDeletion(c, prefix)
+			isAllNil := c.numChildren() == 0
+			if isAllNil && c.value == nil {
+				curr = nil
+			}
+
+			if *limit == 0 {
+				return curr, true
+			}
+		}
+
+		if *limit == 0 {
+			return c, true
+		}
+
+		// Delete the current node as well
+		if c.value != nil {
+			*limit--
+		}
+		return nil, true
+	}
+
+	return curr, true
+}
+
 // ClearPrefix deletes all key-value pairs from the trie where the key starts with the given prefix
 func (t *Trie) ClearPrefix(prefix []byte) {
 	if len(prefix) == 0 {
@@ -611,10 +740,10 @@ func (t *Trie) delete(parent node, key []byte) (node, bool) {
 		// Key doesn't exist.
 		return p, false
 	case nil:
-		// do nothing
+		return nil, false
+	default:
+		panic(fmt.Sprintf("%T: invalid node: %v (%v)", p, p, key))
 	}
-	// This should never happen.
-	return nil, false
 }
 
 // handleDeletion is called when a value is deleted from a branch