Add tracekv readOperations and getAllKeysUsedInTrace

go.mod

@@ -61,6 +61,8 @@ require (
 	sigs.k8s.io/yaml v1.3.0
 )
 
+require github.com/chrispappas/golang-generics-set v1.0.1
+
 require (
 	cloud.google.com/go v0.100.2 // indirect
 	cloud.google.com/go/compute v1.6.1 // indirect

go.sum

@@ -208,6 +208,8 @@ github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XL
 github.com/cheekybits/genny v1.0.0 h1:uGGa4nei+j20rOSeDeP5Of12XVm7TGUd4dJA9RDitfE=
 github.com/cheekybits/genny v1.0.0/go.mod h1:+tQajlRqAUrPI7DOSpB0XAqZYtQakVtB7wXkRAgjxjQ=
 github.com/cheggaaa/pb v1.0.27/go.mod h1:pQciLPpbU0oxA0h+VJYYLxO+XeDQb5pZijXscXHm81s=
+github.com/chrispappas/golang-generics-set v1.0.1 h1:91l8cInAWTxCPwZ8UNg7qkkPsdFdkYS9hytsd8UJsIU=
+github.com/chrispappas/golang-generics-set v1.0.1/go.mod h1:cp8j73+rlDyFF9PrjUkrRvi8L4jSRIsRK6Q1nPPIoqo=
 github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=
 github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI=
 github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=

store/tracekv/store.go

@@ -1,10 +1,13 @@
 package tracekv
 
 import (
+	"bytes"
 	"encoding/base64"
 	"encoding/json"
+	"fmt"
 	"io"
 
+	"github.com/chrispappas/golang-generics-set/set"
 	"github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/cosmos/cosmos-sdk/types/errors"
 )
@@ -17,6 +20,10 @@ const (
 	iterValueOp operation = "iterValue"
 )
 
+var (
+	ErrBufferEmpty = fmt.Errorf("provided buffer is empty")
+)
+
 type (
 	// Store implements the KVStore interface with tracing enabled.
 	// Operations are traced on each core KVStore call and written to the
@@ -90,6 +97,28 @@ func (tkv *Store) ReverseIterator(start, end []byte) types.Iterator {
 	return tkv.iterator(start, end, false)
 }
 
+// GetAllKeysUsedInTrace reads through all traced operations and returns
+// a set of all the keys inside the trace operations
+func (tkv *Store) GetAllKeysUsedInTrace(buf bytes.Buffer) set.Set[string] {
+
+	keys := make(set.Set[string], 0)
+	for {
+		traceOp, err := readOperation(&buf)
+		// Reached end of buffer
+		if err == ErrBufferEmpty {
+			return keys
+		}
+		if err != nil {
+			panic(err)
+		}
+		key, err := base64.StdEncoding.DecodeString(traceOp.Key)
+		if err != nil {
+			panic(errors.Wrap(err, "failed to decode key read from buf"))
+		}
+		keys.Add(string(key))
+	}
+}
+
 // iterator facilitates iteration over a KVStore. It delegates the necessary
 // calls to it's parent KVStore.
 func (tkv *Store) iterator(start, end []byte, ascending bool) types.Iterator {
@@ -202,3 +231,22 @@ func writeOperation(w io.Writer, op operation, tc types.TraceContext, key, value
 
 	io.WriteString(w, "\n")
 }
+
+// readOperation reads a KVStore operation from the underlying buffer as
+// JSON-encoded data where the key/value pair is base64 encoded.
+func readOperation(r *bytes.Buffer) (*traceOperation, error) {
+	raw, err := r.ReadString('\n')
+	if raw == "" {
+		return nil, ErrBufferEmpty
+	}
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to read trace operation")
+	}
+	traceOp := traceOperation{}
+	err = json.Unmarshal([]byte(raw), &traceOp)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to deserialize trace operation")
+	}
+
+	return &traceOp, nil
+}

store/tracekv/store_test.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"testing"
 
+	"github.com/chrispappas/golang-generics-set/set"
 	"github.com/stretchr/testify/require"
 
 	dbm "github.com/tendermint/tm-db"
@@ -113,6 +114,25 @@ func TestTraceKVStoreSet(t *testing.T) {
 	require.Panics(t, func() { store.Set(nil, []byte("value")) }, "setting a nil key should panic")
 }
 
+func TestGetAllKeysUsedInTrace(t *testing.T) {
+	expectedKeys := set.FromSlice([]string{
+		string(kvPairs[0].Key),
+		string(kvPairs[1].Key),
+		string(kvPairs[2].Key),
+	})
+
+	var buf bytes.Buffer
+	store := newEmptyTraceKVStore(&buf)
+	buf.Reset()
+
+	for _, kvPair := range kvPairs {
+		store.Set(kvPair.Key, kvPair.Value)
+	}
+
+	keys := store.GetAllKeysUsedInTrace(buf)
+	require.Equal(t, expectedKeys, keys)
+}
+
 func TestTraceKVStoreDelete(t *testing.T) {
 	testCases := []struct {
 		key         []byte