feat: collections API

CHANGELOG.md

@@ -41,6 +41,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+* [#894](https://github.com/NibiruChain/nibiru/pull/894) - add the collections package!
+
 ### State Machine Breaking
 
 * [#872](https://github.com/NibiruChain/nibiru/pull/872) - x/perp remove module balances from genesis

collections/collections.go

@@ -0,0 +1,64 @@
+package collections
+
+import (
+	"bytes"
+
+	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/gogo/protobuf/proto"
+)
+
+type Object interface {
+	codec.ProtoMarshaler
+}
+
+// setObject is used when no object functionality is needed.
+type setObject struct{}
+
+func (n setObject) Reset() {
+	panic("must never be called")
+}
+
+func (n setObject) String() string {
+	panic("must never be called")
+}
+
+func (n setObject) ProtoMessage() {
+	panic("must never be called")
+}
+
+func (n setObject) Marshal() ([]byte, error) {
+	return []byte{}, nil
+}
+
+func (n setObject) MarshalTo(_ []byte) (_ int, _ error) {
+	panic("must never be called")
+}
+
+func (n setObject) MarshalToSizedBuffer(_ []byte) (int, error) {
+	panic("must never be called")
+}
+
+func (n setObject) Size() int {
+	panic("must never be called")
+}
+
+func (n setObject) Unmarshal(b []byte) error {
+	if !bytes.Equal(b, []byte{}) {
+		panic("bad usage")
+	}
+	return nil
+}
+
+var _ Object = (*setObject)(nil)
+
+func typeName(o Object) string {
+	switch o.(type) {
+	case *setObject, setObject:
+		return "no-op-object"
+	}
+	n := proto.MessageName(o)
+	if n == "" {
+		panic("invalid Object implementation")
+	}
+	return n
+}

collections/errors.go

@@ -0,0 +1,12 @@
+package collections
+
+import (
+	"errors"
+	"fmt"
+)
+
+var ErrNotFound = errors.New("collections: not found")
+
+func notFoundError(name string, key string) error {
+	return fmt.Errorf("%w object '%s' with key %s", ErrNotFound, name, key)
+}

collections/item.go

@@ -0,0 +1,73 @@
+package collections
+
+import (
+	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+// itemKey is a constant byte key which maps an Item object.
+var itemKey = []byte{0x0}
+
+// NewItem instantiates a new Item instance.
+func NewItem[V any, PV interface {
+	*V
+	Object
+}](cdc codec.BinaryCodec, sk sdk.StoreKey, prefix uint8) Item[V, PV] {
+	return Item[V, PV]{
+		prefix:   []byte{prefix},
+		sk:       sk,
+		cdc:      cdc,
+		typeName: typeName(PV(new(V))),
+	}
+}
+
+// Item represents a state object which will always have one instance
+// of itself saved in the namespace.
+// Examples are:
+// 	- config
+// 	- parameters
+// 	- a sequence
+type Item[V any, PV interface {
+	*V
+	Object
+}] struct {
+	_        V
+	prefix   []byte
+	sk       sdk.StoreKey
+	cdc      codec.BinaryCodec
+	typeName string
+}
+
+func (i Item[V, PV]) getStore(ctx sdk.Context) sdk.KVStore {
+	return prefix.NewStore(ctx.KVStore(i.sk), i.prefix)
+}
+
+// Get gets the item V or returns an error.
+func (i Item[V, PV]) Get(ctx sdk.Context) (V, error) {
+	s := i.getStore(ctx)
+	bytes := s.Get(itemKey)
+	if bytes == nil {
+		var v V
+		return v, notFoundError(i.typeName, "item")
+	}
+
+	var v V
+	i.cdc.MustUnmarshal(bytes, PV(&v))
+	return v, nil
+}
+
+// GetOr either returns the provided default
+// if it's not present in state, or the value found in state.
+func (i Item[V, PV]) GetOr(ctx sdk.Context, def V) V {
+	got, err := i.Get(ctx)
+	if err != nil {
+		return def
+	}
+	return got
+}
+
+// Set sets the item value to v.
+func (i Item[V, PV]) Set(ctx sdk.Context, v V) {
+	i.getStore(ctx).Set(itemKey, i.cdc.MustMarshal(PV(&v)))
+}

collections/keys/keys.go

@@ -0,0 +1,45 @@
+package keys
+
+import (
+	"fmt"
+)
+
+// Order defines the ordering of keys.
+type Order uint8
+
+const (
+	OrderAscending Order = iota
+	OrderDescending
+)
+
+// Key defines a type which can be converted to and from bytes.
+// Constraints:
+// - It's ordered, meaning, for example:
+//		StringKey("a").KeyBytes() < StringKey("b").KeyBytes().
+//      Int64Key(100).KeyBytes() > Int64Key(-100).KeyBytes()
+// - Going back and forth using KeyBytes and FromKeyBytes produces the same results.
+// - It's prefix safe, meaning that bytes.Contains(StringKey("a").KeyBytes(), StringKey("aa").KeyBytes()) = false.
+type Key interface {
+	// KeyBytes returns the key as bytes.
+	KeyBytes() []byte
+	// FromKeyBytes parses the Key from bytes.
+	// returns i which is the index of the end of the key.
+	// Constraint: Key == Self (aka the interface implementer).
+	// NOTE(mercilex): we in theory should return Key[T any] and constrain
+	// in the collections.Map, collections.IndexedMap, collections.Set
+	// that T is in fact the Key itself.
+	// We don't do it otherwise all our APIs would get messy
+	// due to golang's compiler type inference.
+	FromKeyBytes(b []byte) (i int, k Key)
+	// Stringer is implemented to allow human-readable formats, especially important in errors.
+	fmt.Stringer
+}
+
+func validString[T ~string](s T) error {
+	for i, c := range s {
+		if c == 0 {
+			return fmt.Errorf("invalid null character at index %d: %s", i, s)
+		}
+	}
+	return nil
+}

collections/keys/numeric.go

@@ -0,0 +1,38 @@
+package keys
+
+import (
+	"encoding/binary"
+	"fmt"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+type Uint8Key uint8
+
+func (u Uint8Key) KeyBytes() []byte {
+	return []byte{uint8(u)}
+}
+
+func (u Uint8Key) FromKeyBytes(b []byte) (i int, k Key) {
+	return 0, Uint8Key(b[0])
+}
+
+func (u Uint8Key) String() string { return fmt.Sprintf("%d", u) }
+
+func Uint64[T ~uint64](u T) Uint64Key {
+	return Uint64Key(u)
+}
+
+type Uint64Key uint64
+
+func (u Uint64Key) KeyBytes() []byte {
+	return sdk.Uint64ToBigEndian(uint64(u))
+}
+
+func (u Uint64Key) FromKeyBytes(b []byte) (i int, k Key) {
+	return 7, Uint64(binary.BigEndian.Uint64(b))
+}
+
+func (u Uint64Key) String() string {
+	return fmt.Sprintf("%d", u)
+}

collections/keys/pair.go

@@ -0,0 +1,89 @@
+package keys
+
+import (
+	"fmt"
+)
+
+// Join joins the two parts of a Pair key.
+func Join[K1 Key, K2 Key](k1 K1, k2 K2) Pair[K1, K2] {
+	return Pair[K1, K2]{
+		p1: &k1,
+		p2: &k2,
+	}
+}
+
+// PairPrefix is used to provide only the K1 part of the Pair.
+// Usually used in Range.Prefix where Key is Pair.
+func PairPrefix[K1 Key, K2 Key](k1 K1) Pair[K1, K2] {
+	return Pair[K1, K2]{
+		p1: &k1,
+		p2: nil,
+	}
+}
+
+// PairSuffix is used to provide only the K2 part of the Pair.
+// Usually used in Range.Start or Range.End where Key is Pair.
+func PairSuffix[K1 Key, K2 Key](k2 K2) Pair[K1, K2] {
+	return Pair[K1, K2]{
+		p1: nil,
+		p2: &k2,
+	}
+}
+
+// Pair represents a multipart key composed of
+// two Key of different or equal types.
+type Pair[K1 Key, K2 Key] struct {
+	// p1 is the first part of the Pair.
+	p1 *K1
+	// p2 is the second part of the Pair.
+	p2 *K2
+}
+
+func (t Pair[K1, K2]) fkb1(b []byte) (int, K1) {
+	var k1 K1
+	i, p1 := k1.FromKeyBytes(b)
+	return i, p1.(K1)
+}
+
+func (t Pair[K1, K2]) fkb2(b []byte) (int, K2) {
+	var k2 K2
+	i, p2 := k2.FromKeyBytes(b)
+	return i, p2.(K2)
+}
+
+func (t Pair[K1, K2]) FromKeyBytes(b []byte) (int, Key) {
+	// NOTE(mercilex): is it always safe to assume that when we get a part
+	// of the key it's going to always contain the full key and not only a part?
+	i1, k1 := t.fkb1(b)
+	i2, k2 := t.fkb2(b[i1+1:]) // add one to not pass last index
+	// add one back as the indexes reported back will start from the last index + 1
+	return i1 + i2 + 1, Pair[K1, K2]{
+		p1: &k1,
+		p2: &k2,
+	}
+}
+
+func (t Pair[K1, K2]) KeyBytes() []byte {
+	if t.p1 != nil && t.p2 != nil {
+		return append((*t.p1).KeyBytes(), (*t.p2).KeyBytes()...)
+	} else if t.p1 != nil && t.p2 == nil {
+		return (*t.p1).KeyBytes()
+	} else if t.p1 == nil && t.p2 != nil {
+		return (*t.p2).KeyBytes()
+	} else {
+		panic("empty Pair key")
+	}
+}
+
+func (t Pair[K1, K2]) String() string {
+	p1 := "<nil>"
+	p2 := "<nil>"
+	if t.p1 != nil {
+		p1 = (*t.p1).String()
+	}
+	if t.p2 != nil {
+		p2 = (*t.p2).String()
+	}
+
+	return fmt.Sprintf("('%s', '%s')", p1, p2)
+}

collections/keys/pair_test.go

@@ -0,0 +1,39 @@
+package keys
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestPair(t *testing.T) {
+	// we only care about bijectivity
+	// as Pair is strictly K1, K2 implementation reliant.
+
+	t.Run("joined", func(t *testing.T) {
+		p := Join(StringKey("hi"), Join(StringKey("hi"), StringKey("hi")))
+		bytes := p.KeyBytes()
+		idx, result := p.FromKeyBytes(bytes)
+		require.Equalf(t, p, result, "%s <-> %s", p.String(), result.String())
+		require.Equal(t, len(bytes)-1, idx)
+	})
+
+	t.Run("pair prefix", func(t *testing.T) {
+		k1 := StringKey("hi")
+		prefix := PairPrefix[StringKey, Uint64Key](k1)
+		require.Equal(t, k1.KeyBytes(), prefix.KeyBytes())
+	})
+
+	t.Run("pair suffix", func(t *testing.T) {
+		k1 := Uint64Key(10)
+		suffix := PairSuffix[StringKey, Uint64Key](k1)
+		require.Equal(t, k1.KeyBytes(), suffix.KeyBytes())
+	})
+
+	t.Run("empty", func(t *testing.T) {
+		var p Pair[StringKey, StringKey]
+		require.Panics(t, func() {
+			p.KeyBytes()
+		})
+	})
+}