Tidy up code

baseapp/baseapp.go

@@ -779,22 +779,17 @@ func generateFraudProof(store *multi.Store, storeKeyToSubstoreTraceBuf map[types
 		keys := traceKv.GetAllKeysUsedInTrace(*subStoreTraceBuf)
 
 		substoreSMT := store.GetSubStoreSMT(storeKey.Name())
-
-		root := substoreSMT.Root()
-
-		var stateWitness StateWitness
-		stateWitness.root = root
-
+		stateWitness := StateWitness{
+			WitnessData: make([]WitnessData, 0, keys.Len()),
+		}
 		for key := range keys {
-			var witnessData WitnessData
 			value := substoreSMT.Get([]byte(key))
 			proof, err := substoreSMT.GetSMTProof([]byte(key))
 			if err != nil {
 				panic(err)
 			}
-			witnessData.Key = []byte(key)
-			witnessData.Value = []byte(value)
-			witnessData.proof = *proof
+			bKey, bVal := []byte(key), []byte(value)
+			witnessData := WitnessData{bKey, bVal, *proof}
 			stateWitness.WitnessData = append(stateWitness.WitnessData, witnessData)
 		}
 		fraudProof.stateWitness[storeKey.Name()] = stateWitness

baseapp/baseapp_test.go

@@ -19,7 +19,6 @@ import (
 	snapshottypes "github.com/cosmos/cosmos-sdk/snapshots/types"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	stypes "github.com/cosmos/cosmos-sdk/store/v2alpha1"
-	types "github.com/cosmos/cosmos-sdk/store/v2alpha1"
 	"github.com/cosmos/cosmos-sdk/store/v2alpha1/multi"
 	"github.com/cosmos/cosmos-sdk/testutil"
 	"github.com/cosmos/cosmos-sdk/testutil/mock"
@@ -100,8 +99,8 @@ func setupBaseApp(t *testing.T, options ...AppOption) *BaseApp {
 
 // baseapp loaded from a fraudproof
 func setupBaseAppFromFraudProof(t *testing.T, fraudProof FraudProof, options ...AppOption) *BaseApp {
-	storeKeys := make([]types.StoreKey, 0, len(fraudProof.stateWitness))
-	routerOpts := make([]func(bapp *BaseApp), 0)
+	storeKeys := make([]stypes.StoreKey, 0, len(fraudProof.stateWitness))
+	routerOpts := make([]func(bapp *BaseApp), 0, len(fraudProof.stateWitness))
 	for storeKeyName := range fraudProof.stateWitness {
 		storeKey := sdk.NewKVStoreKey(storeKeyName)
 		storeKeys = append(storeKeys, storeKey)
@@ -120,7 +119,10 @@ func setupBaseAppFromFraudProof(t *testing.T, fraudProof FraudProof, options ...
 	for _, routerOpt := range routerOpts {
 		options = append(options, AppOptionFunc(routerOpt))
 	}
+
+	// This initial height is used in `BeginBlock` in `validateHeight`
 	options = append(options, SetInitialHeight(fraudProof.blockHeight))
+
 	for storeKey := range fraudProof.stateWitness {
 		stateWitness := fraudProof.stateWitness[storeKey]
 		witnessData := stateWitness.WitnessData
@@ -2302,7 +2304,7 @@ func TestGenerateAndLoadFraudProof(t *testing.T) {
 
 	// Exports all data inside current multistore into a fraudProof (S1) //
 
-	storeKeyToSubstoreTraceBuf := make(map[types.StoreKey]*bytes.Buffer)
+	storeKeyToSubstoreTraceBuf := make(map[stypes.StoreKey]*bytes.Buffer)
 	storeKeyToSubstoreTraceBuf[capKey2] = subStoreTraceBuf
 
 	fraudProof := generateFraudProof(cms, storeKeyToSubstoreTraceBuf)

baseapp/fraudproof.go

@@ -2,17 +2,22 @@ package baseapp
 
 import "github.com/lazyledger/smt"
 
+// Represents a single-round fraudProof
 type FraudProof struct {
+	// The block height to load state of
 	blockHeight int64
 
+	// A map from module name to state witness
 	stateWitness map[string]StateWitness
 }
 
+// State witness with a list of all witness data
 type StateWitness struct {
-	root        []byte
+	// List of witness data
 	WitnessData []WitnessData
 }
 
+// Witness data containing a key/value pair and a SMT proof for said key/value pair
 type WitnessData struct {
 	Key   []byte
 	Value []byte

baseapp/options.go

@@ -89,7 +89,7 @@ func SetIndexEvents(ie []string) AppOptionFunc {
 // SetInterBlockCache provides a BaseApp option function that sets the
 // inter-block cache.
 func SetInterBlockCache(cache sdk.MultiStorePersistentCache) AppOptionFunc {
-	opt := func(cfg *multi.StoreParams, v uint64) error {
+	opt := func(cfg *multi.StoreParams, _ uint64) error {
 		cfg.PersistentCache = cache
 		return nil
 	}

go.mod

@@ -65,7 +65,10 @@ require (
 	sigs.k8s.io/yaml v1.3.0
 )
 
-require github.com/cosmos/cosmos-sdk/store/tools/ics23 v0.0.0-00010101000000-000000000000
+require (
+	github.com/chrispappas/golang-generics-set v1.0.1
+	github.com/cosmos/cosmos-sdk/store/tools/ics23 v0.0.0-00010101000000-000000000000
+)
 
 require (
 	4d63.com/gochecknoglobals v0.1.0 // indirect

go.sum

@@ -229,6 +229,8 @@ github.com/chavacava/garif v0.0.0-20220316182200-5cad0b5181d4 h1:tFXjAxje9thrTF4
 github.com/chavacava/garif v0.0.0-20220316182200-5cad0b5181d4/go.mod h1:W8EnPSQ8Nv4fUjc/v1/8tHFqhuOJXnRub0dTfuAQktU=
 github.com/checkpoint-restore/go-criu/v5 v5.3.0/go.mod h1:E/eQpaFtUKGOOSEBZgmKAcn+zUUwWxqcaKZlF54wK8E=
 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

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"io"
 
+	"github.com/chrispappas/golang-generics-set/set"
 	"github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/cosmos/cosmos-sdk/types/errors"
 )
@@ -93,20 +94,23 @@ func (tkv *Store) ReverseIterator(start, end []byte) types.Iterator {
 
 // 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) map[string]bool {
-	var traceOp traceOperation
-	var err error
-	keys := make(map[string]bool)
+func (tkv *Store) GetAllKeysUsedInTrace(buf bytes.Buffer) set.Set[string] {
+
+	keys := make(set.Set[string], 0)
 	for {
-		traceOp, err = readOperation(&buf)
+		traceOp, err := readOperation(&buf)
 		if err != nil {
-			return keys
+			errString := err.Error()
+			if errString == "provided buffer is empty: EOF" {
+				return keys
+			}
+			panic(err)
 		}
 		key, err := base64.StdEncoding.DecodeString(traceOp.Key)
 		if err != nil {
 			panic(errors.Wrap(err, "failed to decode key read from buf"))
 		}
-		keys[string(key)] = true
+		keys.Add(string(key))
 	}
 }
 
@@ -225,19 +229,19 @@ func writeOperation(w io.Writer, op operation, tc types.TraceContext, key, value
 
 // reaOperation 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) {
-	var traceOp traceOperation
+func readOperation(r *bytes.Buffer) (*traceOperation, error) {
 	raw, err := r.ReadString('\n')
 	if raw == "" {
-		return traceOp, errors.Wrap(err, "provided buffer is empty")
+		return nil, errors.Wrap(err, "provided buffer is empty")
 	}
 	if err != nil {
-		return traceOp, errors.Wrap(err, "failed to read trace operation")
+		return nil, errors.Wrap(err, "failed to read trace operation")
 	}
+	traceOp := traceOperation{}
 	err = json.Unmarshal([]byte(raw), &traceOp)
 	if err != nil {
-		return traceOp, errors.Wrap(err, "failed to deserialize trace operation")
+		return nil, errors.Wrap(err, "failed to deserialize trace operation")
 	}
 
-	return traceOp, nil
+	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"
@@ -114,11 +115,11 @@ func TestTraceKVStoreSet(t *testing.T) {
 }
 
 func TestGetAllKeysUsedInTrace(t *testing.T) {
-	expectedKeys := map[string]bool{
-		string(kvPairs[0].Key): true,
-		string(kvPairs[1].Key): true,
-		string(kvPairs[2].Key): true,
-	}
+	expectedKeys := set.FromSlice([]string{
+		string(kvPairs[0].Key),
+		string(kvPairs[1].Key),
+		string(kvPairs[2].Key),
+	})
 
 	var buf bytes.Buffer
 	store := newEmptyTraceKVStore(&buf)