feat(collections): genesis support (#14331)

Co-authored-by: testinginprod <testinginprod@somewhere.idk>
Co-authored-by: testinginprod <98415576+testinginprod@users.noreply.github.com>
Co-authored-by: Marko <marbar3778@yahoo.com>
Co-authored-by: Julien Robert <julien@rbrt.fr>
Co-authored-by: Aleksandr Bezobchuk <alexanderbez@users.noreply.github.com>
Co-authored-by: Rafael Tenfen <rafaeltenfen.rt@gmail.com>
Co-authored-by: Facundo Medica <14063057+facundomedica@users.noreply.github.com>
Co-authored-by: atheeshp <59333759+atheeshp@users.noreply.github.com>
Co-authored-by: Likhita Polavarapu <78951027+likhita-809@users.noreply.github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Matt Kocubinski <mkocubinski@gmail.com>
Co-authored-by: Jacob Gadikian <faddat@users.noreply.github.com>
Co-authored-by: samricotta <37125168+samricotta@users.noreply.github.com>
Co-authored-by: Noel Ukwa <noeluchechukwu@gmail.com>
Co-authored-by: JayT106 <JayT106@users.noreply.github.com>
Co-authored-by: Julián Toledano <JulianToledano@users.noreply.github.com>
Co-authored-by: Amaury <1293565+amaurym@users.noreply.github.com>
Co-authored-by: Denver <aeharvlee@gmail.com>
Co-authored-by: Hyunwoo Lee <denver@Hyunwoos-MacBook-Pro-2.local>
Co-authored-by: Federico Kunze Küllmer <31522760+fedekunze@users.noreply.github.com>
Co-authored-by: MalteHerrmann <42640438+MalteHerrmann@users.noreply.github.com>
Co-authored-by: Vladislav Varadinov <vladislav.varadinov@gmail.com>
Co-authored-by: cipherZ <dev@cipherz.com>

collections/collections.go

@@ -22,6 +22,8 @@ type collection interface {
 
 	// getPrefix is the unique prefix of the collection within a schema.
 	getPrefix() []byte
+
+	genesisHandler
 }
 
 // Prefix defines a segregation namespace
@@ -85,6 +87,10 @@ type KeyCodec[T any] interface {
 	// to return the maximum varint bytes buffer length, at the risk of
 	// over-estimating in order to pick the most performant path.
 	Size(key T) int
+	// EncodeJSON encodes the value as JSON.
+	EncodeJSON(value T) ([]byte, error)
+	// DecodeJSON decodes the provided JSON bytes into an instance of T.
+	DecodeJSON(b []byte) (T, error)
 	// Stringify returns a string representation of T.
 	Stringify(key T) string
 	// KeyType returns a string identifier for the type of the key.
@@ -117,6 +123,10 @@ type ValueCodec[T any] interface {
 	Encode(value T) ([]byte, error)
 	// Decode returns the type T given its binary representation.
 	Decode(b []byte) (T, error)
+	// EncodeJSON encodes the value as JSON.
+	EncodeJSON(value T) ([]byte, error)
+	// DecodeJSON decodes the provided JSON bytes into an instance of T.
+	DecodeJSON(b []byte) (T, error)
 	// Stringify returns a string representation of T.
 	Stringify(value T) string
 	// ValueType returns the identifier for the type.

collections/collections_test.go

@@ -92,6 +92,13 @@ func checkKeyCodec[T any](t *testing.T, keyCodec KeyCodec[T], key T) {
 	require.NoError(t, err)
 	require.Equal(t, len(buffer), read, "encoded non terminal key and pair key read bytes must have same size")
 	require.Equal(t, pairKey, decodedPairKey, "encoding and decoding produces different keys with non terminal encoding")
+
+	// check JSON
+	keyJSON, err := keyCodec.EncodeJSON(key)
+	require.NoError(t, err)
+	decoded, err := keyCodec.DecodeJSON(keyJSON)
+	require.NoError(t, err)
+	require.Equal(t, key, decoded, "json encoding and decoding did not produce the same results")
 }
 
 // checkValueCodec asserts the correct behaviour of a ValueCodec over the type T.

collections/genesis.go

@@ -0,0 +1,153 @@
+package collections
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+)
+
+type genesisHandler interface {
+	validateGenesis(r io.Reader) error
+	importGenesis(ctx context.Context, r io.Reader) error
+	exportGenesis(ctx context.Context, w io.Writer) error
+	defaultGenesis(w io.Writer) error
+}
+
+type jsonMapEntry struct {
+	Key   json.RawMessage `json:"key"`
+	Value json.RawMessage `json:"value,omitempty"`
+}
+
+func (m Map[K, V]) validateGenesis(reader io.Reader) error {
+	return m.doDecodeJson(reader, func(key K, value V) error {
+		return nil
+	})
+}
+
+func (m Map[K, V]) importGenesis(ctx context.Context, reader io.Reader) error {
+	return m.doDecodeJson(reader, func(key K, value V) error {
+		return m.Set(ctx, key, value)
+	})
+}
+
+func (m Map[K, V]) exportGenesis(ctx context.Context, writer io.Writer) error {
+	_, err := writer.Write([]byte("["))
+	if err != nil {
+		return err
+	}
+
+	it, err := m.Iterate(ctx, nil)
+	if err != nil {
+		return err
+	}
+	defer it.Close()
+
+	first := true
+	for ; it.Valid(); it.Next() {
+		// add a comma before encoding the object
+		// for all objects besides the first one.
+		if !first {
+			_, err = writer.Write([]byte(","))
+			if err != nil {
+				return err
+			}
+		}
+		first = false
+
+		key, err := it.Key()
+		if err != nil {
+			return err
+		}
+
+		keyBz, err := m.kc.EncodeJSON(key)
+		if err != nil {
+			return err
+		}
+
+		value, err := it.Value()
+		if err != nil {
+			return err
+		}
+
+		valueBz, err := m.vc.EncodeJSON(value)
+		if err != nil {
+			return err
+		}
+
+		entry := jsonMapEntry{
+			Key:   keyBz,
+			Value: valueBz,
+		}
+
+		bz, err := json.Marshal(entry)
+		if err != nil {
+			return err
+		}
+
+		_, err = writer.Write(bz)
+		if err != nil {
+			return err
+		}
+	}
+
+	_, err = writer.Write([]byte("]"))
+	return err
+}
+
+func (m Map[K, V]) doDecodeJson(reader io.Reader, onEntry func(key K, value V) error) error {
+	decoder := json.NewDecoder(reader)
+	token, err := decoder.Token()
+	if err != nil {
+		return err
+	}
+
+	if token != json.Delim('[') {
+		return fmt.Errorf("expected [ got %s", token)
+	}
+
+	for decoder.More() {
+		var rawJson json.RawMessage
+		err := decoder.Decode(&rawJson)
+		if err != nil {
+			return err
+		}
+
+		var mapEntry jsonMapEntry
+		err = json.Unmarshal(rawJson, &mapEntry)
+		if err != nil {
+			return err
+		}
+
+		key, err := m.kc.DecodeJSON(mapEntry.Key)
+		if err != nil {
+			return err
+		}
+
+		value, err := m.vc.DecodeJSON(mapEntry.Value)
+		if err != nil {
+			return err
+		}
+
+		err = onEntry(key, value)
+		if err != nil {
+			return err
+		}
+	}
+
+	token, err = decoder.Token()
+	if err != nil {
+		return err
+	}
+
+	if token != json.Delim(']') {
+		return fmt.Errorf("expected ] got %s", token)
+	}
+
+	return nil
+}
+
+func (m Map[K, V]) defaultGenesis(writer io.Writer) error {
+	_, err := writer.Write([]byte(`[]`))
+	return err
+}

collections/genesis_test.go

@@ -0,0 +1,160 @@
+package collections
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"testing"
+
+	"cosmossdk.io/core/appmodule"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDefaultGenesis(t *testing.T) {
+	f := initFixture(t)
+	var writers []*bufCloser
+	require.NoError(t, f.schema.DefaultGenesis(func(field string) (io.WriteCloser, error) {
+		w := newBufCloser(t, "")
+		writers = append(writers, w)
+		return w, nil
+	}))
+	require.Len(t, writers, 4)
+	require.Equal(t, `[]`, writers[0].Buffer.String())
+	require.Equal(t, `[]`, writers[1].Buffer.String())
+	require.Equal(t, `[]`, writers[2].Buffer.String())
+	require.Equal(t, `[]`, writers[3].Buffer.String())
+}
+
+func TestValidateGenesis(t *testing.T) {
+	f := initFixture(t)
+	require.NoError(t, f.schema.ValidateGenesis(createTestGenesisSource(t)))
+}
+
+func TestImportGenesis(t *testing.T) {
+	f := initFixture(t)
+	require.NoError(t, f.schema.InitGenesis(f.ctx, createTestGenesisSource(t)))
+	// assert map correct genesis
+	mapIt, err := f.m.Iterate(f.ctx, nil)
+	require.NoError(t, err)
+	defer mapIt.Close()
+
+	kvs, err := mapIt.KeyValues()
+	require.NoError(t, err)
+	require.Equal(t, KeyValue[string, uint64]{Key: "abc", Value: 1}, kvs[0])
+	require.Equal(t, KeyValue[string, uint64]{Key: "def", Value: 2}, kvs[1])
+
+	// assert item correct genesis
+	x, err := f.i.Get(f.ctx)
+	require.NoError(t, err)
+	require.Equal(t, "superCoolItem", x)
+
+	// assert keyset correct genesis
+	ksIt, err := f.ks.Iterate(f.ctx, nil)
+	require.NoError(t, err)
+	defer ksIt.Close()
+
+	keys, err := ksIt.Keys()
+	require.NoError(t, err)
+	require.Equal(t, []string{"0", "1", "2"}, keys)
+
+	// assert sequence correct genesis
+	seq, err := f.s.Peek(f.ctx)
+	require.NoError(t, err)
+	require.Equal(t, uint64(1000), seq)
+}
+
+func TestExportGenesis(t *testing.T) {
+	f := initFixture(t)
+	require.NoError(t, f.schema.InitGenesis(f.ctx, createTestGenesisSource(t)))
+
+	var writers []*bufCloser
+	require.NoError(t, f.schema.ExportGenesis(f.ctx, func(field string) (io.WriteCloser, error) {
+		w := newBufCloser(t, "")
+		writers = append(writers, w)
+		return w, nil
+	}))
+	require.Len(t, writers, 4)
+	require.Equal(t, expectedItemGenesis, writers[0].Buffer.String())
+	require.Equal(t, expectedKeySetGenesis, writers[1].Buffer.String())
+	require.Equal(t, expectedMapGenesis, writers[2].Buffer.String())
+	require.Equal(t, expectedSequenceGenesis, writers[3].Buffer.String())
+}
+
+type testFixture struct {
+	schema Schema
+	ctx    context.Context
+	m      Map[string, uint64]
+	i      Item[string]
+	s      Sequence
+	ks     KeySet[string]
+}
+
+func initFixture(t *testing.T) *testFixture {
+	sk, ctx := deps()
+	schemaBuilder := NewSchemaBuilder(sk)
+	m := NewMap(schemaBuilder, NewPrefix(1), "map", StringKey, Uint64Value)
+	i := NewItem(schemaBuilder, NewPrefix(2), "item", StringValue)
+	s := NewSequence(schemaBuilder, NewPrefix(3), "sequence")
+	ks := NewKeySet(schemaBuilder, NewPrefix(4), "key_set", StringKey)
+	schema, err := schemaBuilder.Build()
+	require.NoError(t, err)
+	return &testFixture{
+		schema: schema,
+		ctx:    ctx,
+		m:      m,
+		i:      i,
+		s:      s,
+		ks:     ks,
+	}
+}
+
+func createTestGenesisSource(t *testing.T) appmodule.GenesisSource {
+	expectedOrder := []string{"item", "key_set", "map", "sequence"}
+	currentIndex := 0
+	return func(field string) (io.ReadCloser, error) {
+		require.Equal(t, expectedOrder[currentIndex], field, "unordered genesis")
+		currentIndex++
+
+		switch field {
+		case "map":
+			return newBufCloser(t, expectedMapGenesis), nil
+		case "item":
+			return newBufCloser(t, expectedItemGenesis), nil
+		case "key_set":
+			return newBufCloser(t, expectedKeySetGenesis), nil
+		case "sequence":
+			return newBufCloser(t, expectedSequenceGenesis), nil
+		default:
+			return nil, nil
+		}
+	}
+}
+
+const (
+	expectedMapGenesis      = `[{"key":"abc","value":"1"},{"key":"def","value":"2"}]`
+	expectedItemGenesis     = `[{"key":"item","value":"superCoolItem"}]`
+	expectedKeySetGenesis   = `[{"key":"0"},{"key":"1"},{"key":"2"}]`
+	expectedSequenceGenesis = `[{"key":"item","value":"1000"}]`
+)
+
+type bufCloser struct {
+	*bytes.Buffer
+	closed bool
+}
+
+func (b *bufCloser) Close() error {
+	b.closed = true
+	return nil
+}
+
+func newBufCloser(t *testing.T, str string) *bufCloser {
+	b := &bufCloser{
+		Buffer: bytes.NewBufferString(str),
+		closed: false,
+	}
+	// this ensures Close was called by the implementation
+	t.Cleanup(func() {
+		require.True(t, b.closed)
+	})
+	return b
+}

collections/go.mod

@@ -10,6 +10,8 @@ require (
 )
 
 require (
+	cosmossdk.io/api v0.2.6 // indirect
+	cosmossdk.io/depinject v1.0.0-alpha.3 // indirect
 	github.com/DataDog/zstd v1.4.5 // indirect
 	github.com/HdrHistogram/hdrhistogram-go v1.1.2 // indirect
 	github.com/cespare/xxhash/v2 v2.1.2 // indirect
@@ -18,9 +20,11 @@ require (
 	github.com/cockroachdb/pebble v0.0.0-20220817183557-09c6e030a677 // indirect
 	github.com/cockroachdb/redact v1.0.8 // indirect
 	github.com/cockroachdb/sentry-go v0.6.1-cockroachdb.2 // indirect
+	github.com/cosmos/cosmos-proto v1.0.0-beta.1 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/google/btree v1.1.2 // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
@@ -38,5 +42,9 @@ require (
 	golang.org/x/exp v0.0.0-20221019170559-20944726eadf // indirect
 	golang.org/x/net v0.3.0 // indirect
 	golang.org/x/sys v0.3.0 // indirect
+	golang.org/x/text v0.5.0 // indirect
+	google.golang.org/genproto v0.0.0-20221118155620-16455021b5e6 // indirect
+	google.golang.org/grpc v1.51.0 // indirect
+	google.golang.org/protobuf v1.28.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )