feat: collections API (#894)

* add: collections

* add pair_test

* chore: CHANGELOG.md

* add: finish range API and range testing

* remove: unused file

* chore: lint

* change: simplify object API allow keys to be values

* change: address reviews

* change: address reviews 2

* change: typeName logic

* chore: doc

CHANGELOG.md

@@ -41,6 +41,7 @@ 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!
 ### CI
 
 * [#785](https://github.com/NibiruChain/nibiru/pull/785) - ci: create simulations job

collections/collections.go

@@ -0,0 +1,81 @@
+package collections
+
+import (
+	"bytes"
+
+	"github.com/gogo/protobuf/proto"
+
+	"github.com/cosmos/cosmos-sdk/codec"
+	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+)
+
+// Object defines an object which can marshal and unmarshal itself to and from bytes.
+type Object interface {
+	// Marshal marshals the object into bytes.
+	Marshal() (b []byte, err error)
+	// Unmarshal populates the object from bytes.
+	Unmarshal(b []byte) error
+}
+
+// storeCodec implements only the subset of functionalities
+// required for the ser/de at state layer.
+// It respects cosmos-sdk guarantees around interface unpacking.
+type storeCodec struct {
+	ir codectypes.InterfaceRegistry
+}
+
+func newStoreCodec(cdc codec.BinaryCodec) storeCodec {
+	return storeCodec{ir: cdc.(*codec.ProtoCodec).InterfaceRegistry()}
+}
+
+func (c storeCodec) marshal(o Object) []byte {
+	bytes, err := o.Marshal()
+	if err != nil {
+		panic(err)
+	}
+	return bytes
+}
+
+func (c storeCodec) unmarshal(bytes []byte, o Object) {
+	err := o.Unmarshal(bytes)
+	if err != nil {
+		panic(err)
+	}
+	err = codectypes.UnpackInterfaces(o, c.ir)
+	if err != nil {
+		panic(err)
+	}
+}
+
+// setObject is used when no object functionality is needed.
+type setObject struct{}
+
+func (n setObject) String() string {
+	panic("must never be called")
+}
+
+func (n setObject) Marshal() ([]byte, error) {
+	return []byte{}, nil
+}
+
+func (n setObject) Unmarshal(b []byte) error {
+	if !bytes.Equal(b, []byte{}) {
+		panic("bad usage")
+	}
+	return nil
+}
+
+var _ Object = (*setObject)(nil)
+
+// TODO(mercilex): improve typeName api
+func typeName(o Object) string {
+	switch o.(type) {
+	case *setObject, setObject:
+		return "no-op-object"
+	}
+	pm, ok := o.(proto.Message)
+	if !ok {
+		return "unknown"
+	}
+	return proto.MessageName(pm)
+}

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:      newStoreCodec(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      storeCodec
+	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.unmarshal(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.marshal(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 numbers of bytes read from the buffer.
+	// 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(buf []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,62 @@
+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 1, Uint8Key(b[0])
+}
+
+func (u Uint8Key) String() string { return fmt.Sprintf("%d", u) }
+
+func (u Uint8Key) Marshal() ([]byte, error) {
+	return []byte{uint8(u)}, nil
+}
+
+func (u *Uint8Key) Unmarshal(b []byte) error {
+	if len(b) != 1 {
+		return fmt.Errorf("invalid bytes type for Uint8Key")
+	}
+	*u = Uint8Key(b[0])
+	return nil
+}
+
+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 8, Uint64(binary.BigEndian.Uint64(b))
+}
+
+func (u Uint64Key) String() string {
+	return fmt.Sprintf("%d", u)
+}
+
+func (u Uint64Key) Marshal() ([]byte, error) {
+	return sdk.Uint64ToBigEndian(uint64(u)), nil
+}
+
+func (u *Uint64Key) Unmarshal(b []byte) error {
+	if len(b) != 8 {
+		return fmt.Errorf("invalid bytes type for Uint64Key")
+	}
+	*u = Uint64(binary.BigEndian.Uint64(b))
+	return nil
+}

collections/keys/pair.go

@@ -0,0 +1,88 @@
+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:])
+	return i1 + i2, 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)
+}