refactor(perp): use collections for state management

CHANGELOG.md

@@ -46,7 +46,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#785](https://github.com/NibiruChain/nibiru/pull/785) - ci: create simulations job
 
 ### State Machine Breaking
-
+* [#952](https://github.com/NibiruChain/nibiru/pull/952) - x/perp move state logic to collections
 * [#872](https://github.com/NibiruChain/nibiru/pull/872) - x/perp remove module balances from genesis
 * [#878](https://github.com/NibiruChain/nibiru/pull/878) - rename `PremiumFraction` to `FundingRate`
 * [#900](https://github.com/NibiruChain/nibiru/pull/900) - refactor x/vpool snapshot state management

collections/index_multi.go

@@ -0,0 +1,139 @@
+package collections
+
+import (
+	"github.com/cosmos/cosmos-sdk/codec"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/NibiruChain/nibiru/collections/keys"
+)
+
+// MultiIndex represents an index in which there is no uniqueness constraint.
+// Which means that multiple primary keys with the same key can exist.
+// It is implemented using a KeySet with keys.Pair[IK, PK], where
+// IK is the index key and PK is the primary key of the object.
+// Indexing keys are simple references, meaning that
+// the Indexing key is formed as concat(index_key, primary_key)
+// Example, given an object Obj{City: milan, ID: 0}, where City is the index and ID is the primary key
+// The following is the generated KeyPair
+// keys.Pair[K1: milan, K2: 0]
+// Simulating that there are multiple objects that were indexed, the following is the Raw KV mapping
+// Key                   | Value
+// ('milan', 0)          | []byte{}
+// ('milan', 5)          | []byte{}
+// ('new york', 1)       | []byte{}
+// ('new york', 2)       | []byte{}
+// So if we want to get all the objects which had City as 'milan'
+// we would prefix over 'milan' in the raw KV to get all the primary keys => 0, 5.
+type MultiIndex[IK keys.Key, PK keys.Key, V any] struct {
+	// indexFn is used to get the secondary key (aka index key)
+	// from the object we're indexing.
+	indexFn func(V) IK
+	// secondaryKeys is a multipart key composed by the
+	// index key (IK) and the primary key (PK)
+	secondaryKeys KeySet[keys.Pair[IK, PK]]
+}
+
+// Insert inserts fetches the index key IK from the object v.
+// And then maps the index key to the primary key.
+func (i *MultiIndex[IK, PK, V]) Insert(ctx sdk.Context, pk PK, v V) {
+	// get secondary key
+	sk := i.indexFn(v)
+	// insert it
+	i.secondaryKeys.Insert(ctx, keys.Join(sk, pk))
+}
+
+// Delete removes the object from the KeySet, removing the references
+// of PK from the index.
+func (i *MultiIndex[IK, PK, V]) Delete(ctx sdk.Context, pk PK, v V) {
+	sk := i.indexFn(v)
+	i.secondaryKeys.Delete(ctx, keys.Join(sk, pk))
+}
+
+// Initialize initializes the index, objectNamespace defines the broader object (V) namespace.
+// IndexNamespace identifies the index namespace in the object namespace.
+func (i *MultiIndex[IK, PK, V]) Initialize(cdc codec.BinaryCodec, storeKey sdk.StoreKey, objectNamespace uint8, indexNamespace uint8) {
+	i.secondaryKeys = NewKeySet[keys.Pair[IK, PK]](cdc, storeKey, indexNamespace)
+	i.secondaryKeys.prefix = []byte{objectNamespace, indexNamespace}
+}
+
+// Iterate iterates over indexing keys using the provided range.
+func (i *MultiIndex[IK, PK, V]) Iterate(ctx sdk.Context, rng keys.Range[keys.Pair[IK, PK]]) IndexIterator[IK, PK] {
+	return IndexIterator[IK, PK]{
+		ks: i.secondaryKeys.Iterate(ctx, rng),
+	}
+}
+
+// Match matches and returns an iterator of all primary keys of objects containing
+// the index key provided to the function.
+func (i *MultiIndex[IK, PK, V]) Match(ctx sdk.Context, ik IK) IndexIterator[IK, PK] {
+	return i.Iterate(ctx, keys.NewRange[keys.Pair[IK, PK]]().Prefix(keys.PairPrefix[IK, PK](ik)))
+}
+
+// ReverseMatch matches and returns a reverse iterator of all primary keys of objects
+// containing the index key provided to the function.
+func (i *MultiIndex[IK, PK, V]) ReverseMatch(ctx sdk.Context, ik IK) IndexIterator[IK, PK] {
+	return i.Iterate(ctx, keys.NewRange[keys.Pair[IK, PK]]().Prefix(keys.PairPrefix[IK, PK](ik)).Descending())
+}
+
+// NewMultiIndex instantiates a new MultiIndex instance.
+// Where IK is the indexing key.
+// PK is the primary key.
+// V is the object being indexed itself.
+// An index function is proved which given the object, returns the indexing key.
+// EXAMPLE:
+// Person {
+// ID: keys.Uint64 (PrimaryKey)
+// City: keys.String (IndexKey)
+// Cap: uint64
+// }
+// indexFn: func(object Person) keys.String { return person.City }
+func NewMultiIndex[IK keys.Key, PK keys.Key, V any](indexFn func(V) IK) *MultiIndex[IK, PK, V] {
+	return &MultiIndex[IK, PK, V]{
+		indexFn: indexFn,
+	}
+}
+
+// IndexIterator wraps a KeySetIterator but provides more
+// index iterator functionalities, such as getting the primary key only.
+type IndexIterator[IK keys.Key, PK keys.Key] struct {
+	ks KeySetIterator[keys.Pair[IK, PK]]
+}
+
+// Keys fully consumes the iterator and returns only
+// the primary keys which matched the query.
+func (i IndexIterator[IK, PK]) Keys() []PK {
+	keys := i.ks.Keys()
+	primaryKeys := make([]PK, len(keys))
+	for i, key := range keys {
+		primaryKeys[i] = key.K2()
+	}
+	return primaryKeys
+}
+
+// FullKeys fully consumes the iterator and returns
+// the keys.Pair containing both index key and primary key.
+func (i IndexIterator[IK, PK]) FullKeys() []keys.Pair[IK, PK] {
+	return i.ks.Keys()
+}
+
+// Key returns the iterator current primary key.
+func (i IndexIterator[IK, PK]) Key() PK {
+	return i.FullKey().K2()
+}
+
+// FullKey returns the iterator current index key + primary key.
+func (i IndexIterator[IK, PK]) FullKey() keys.Pair[IK, PK] {
+	return i.ks.Key()
+}
+
+func (i IndexIterator[IK, PK]) Next() {
+	i.ks.Next()
+}
+
+func (i IndexIterator[IK, PK]) Close() {
+	i.ks.Close()
+}
+
+func (i IndexIterator[IK, PK]) Valid() bool {
+	return i.ks.Valid()
+}

collections/indexed_map.go

@@ -0,0 +1,88 @@
+package collections
+
+import (
+	"github.com/cosmos/cosmos-sdk/codec"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/NibiruChain/nibiru/collections/keys"
+)
+
+// Indexes defines a structure which contains indexes definitions.
+type Indexes[K keys.Key, V any] interface {
+	// IndexList must be implemented by the indexes struct and must return
+	// the indexes to query.
+	// NOTE: changing the order at which elements in IndexList are provided
+	// is breaking, and will cause state corruption.
+	IndexList() []Index[K, V]
+}
+
+// Index defines the index API needed by IndexedMap
+// to index objects of type V, with primary key PK.
+type Index[PK keys.Key, V any] interface {
+	// Insert inserts elements in the index.
+	Insert(ctx sdk.Context, k PK, v V)
+	// Delete deletes objects from the index.
+	Delete(ctx sdk.Context, k PK, v V)
+	// Initialize is called by IndexedMap to initialize the Index.
+	// id is provided by the indexed map based
+	Initialize(cdc codec.BinaryCodec, storeKey sdk.StoreKey, objectNamespace uint8, indexNamespace uint8)
+}
+
+func NewIndexedMap[K keys.Key, V any, PV interface {
+	*V
+	Object
+}, I Indexes[K, V]](cdc codec.BinaryCodec, sk sdk.StoreKey, prefix uint8, indexes I) IndexedMap[K, V, PV, I] {
+	m := NewMap[K, V, PV](cdc, sk, 0)
+	m.prefix = []byte{prefix, 0}
+	for i, index := range indexes.IndexList() {
+		index.Initialize(cdc, sk, prefix, uint8(i)+1)
+	}
+
+	return IndexedMap[K, V, PV, I]{
+		Indexes: indexes,
+		m:       m,
+	}
+}
+
+type IndexedMap[K keys.Key, V any, PV interface {
+	*V
+	Object
+}, I Indexes[K, V]] struct {
+	Indexes I
+	m       Map[K, V, PV]
+}
+
+func (i IndexedMap[K, V, PV, I]) Insert(ctx sdk.Context, k K, v V) {
+	i.m.Insert(ctx, k, v)
+	for _, index := range i.Indexes.IndexList() {
+		index.Insert(ctx, k, v)
+	}
+}
+
+func (i IndexedMap[K, V, PV, I]) Delete(ctx sdk.Context, k K) error {
+	old, err := i.m.Get(ctx, k)
+	if err != nil {
+		return err
+	}
+	err = i.m.Delete(ctx, k)
+	if err != nil {
+		panic(err)
+	}
+
+	for _, index := range i.Indexes.IndexList() {
+		index.Delete(ctx, k, old)
+	}
+	return nil
+}
+
+func (i IndexedMap[K, V, PV, I]) Get(ctx sdk.Context, k K) (V, error) {
+	return i.m.Get(ctx, k)
+}
+
+func (i IndexedMap[K, V, PV, I]) GetOr(ctx sdk.Context, k K, def V) V {
+	return i.m.GetOr(ctx, k, def)
+}
+
+func (i IndexedMap[K, V, PV, I]) Iterate(ctx sdk.Context, rng keys.Range[K]) MapIterator[K, V, PV] {
+	return i.m.Iterate(ctx, rng)
+}

collections/indexed_map_test.go

@@ -0,0 +1,83 @@
+package collections
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/NibiruChain/nibiru/collections/keys"
+)
+
+type object struct {
+	ID    keys.Uint64Key
+	Owner keys.StringKey
+}
+
+func (b object) Marshal() ([]byte, error) {
+	return json.Marshal(b)
+}
+
+func (b *object) Unmarshal(x []byte) error {
+	return json.Unmarshal(x, b)
+}
+
+type Index1 struct {
+	Owner *MultiIndex[keys.StringKey, keys.Uint64Key, object]
+}
+
+func (i Index1) IndexList() []Index[keys.Uint64Key, object] {
+	return []Index[keys.Uint64Key, object]{i.Owner}
+}
+
+func TestNewIndexedMap(t *testing.T) {
+	sk, ctx, cdc := deps()
+	im := NewIndexedMap[keys.Uint64Key, object, *object, Index1](cdc, sk, 0, Index1{
+		Owner: NewMultiIndex[keys.StringKey, keys.Uint64Key, object](func(v object) keys.StringKey {
+			return keys.String(v.Owner)
+		}),
+	})
+
+	im.Insert(ctx, 0, object{
+		ID:    0,
+		Owner: keys.String("mercilex"),
+	})
+
+	im.Insert(ctx, 1, object{
+		ID:    1,
+		Owner: keys.String("mercilex"),
+	})
+
+	im.Insert(ctx, 2, object{
+		ID:    2,
+		Owner: keys.String("heisenberg"),
+	})
+
+	im.Insert(ctx, 3, object{
+		ID:    3,
+		Owner: "mercilex",
+	})
+
+	// we want to range over "mercilex" owned objects.
+	rng := keys.NewRange[keys.Pair[keys.StringKey, keys.Uint64Key]]()
+	pfx := keys.PairPrefix[keys.StringKey, keys.Uint64Key](keys.String("mercilex"))
+	rng = rng.Prefix(pfx)
+
+	ks := im.Indexes.Owner.Iterate(ctx, rng).Keys()
+	require.Equal(t, []keys.Uint64Key{0, 1, 3}, ks)
+
+	// we want to range over "mercilex" owner objects, starting from key0 exclusive and ending key2 inclusive
+	rng = keys.NewRange[keys.Pair[keys.StringKey, keys.Uint64Key]]()
+	pfx = keys.PairPrefix[keys.StringKey, keys.Uint64Key](keys.String("mercilex"))
+	rng = rng.Prefix(pfx)
+	rng = rng.Start(keys.Exclusive(keys.PairSuffix[keys.StringKey, keys.Uint64Key](keys.Uint64(uint64(0)))))
+	rng = rng.End(keys.Inclusive(keys.PairSuffix[keys.StringKey, keys.Uint64Key](keys.Uint64(uint64(2)))))
+
+	ks = im.Indexes.Owner.Iterate(ctx, rng).Keys()
+	require.Equal(t, []keys.Uint64Key{1}, ks)
+
+	// removal of the element from the indexed map reflects on the indexes too
+	require.NoError(t, im.Delete(ctx, ks[0]))
+	ks = im.Indexes.Owner.Iterate(ctx, rng).Keys()
+	require.Empty(t, ks)
+}

collections/keys/bound.go

@@ -0,0 +1,27 @@
+package keys
+
+// Bound defines a key bound.
+type Bound[K Key] struct {
+	value     K    // value is the concrete key.
+	inclusive bool // inclusive defines if the key bound should include or not the provided value.
+}
+
+// Inclusive creates a key Bound which is inclusive,
+// which means the provided key will be included
+// in the key range (if present).
+func Inclusive[K Key](k K) Bound[K] {
+	return Bound[K]{
+		value:     k,
+		inclusive: true,
+	}
+}
+
+// Exclusive creates a key Bound which is exclusive,
+// which means the provided key will be excluded from
+// the key range.
+func Exclusive[K Key](k K) Bound[K] {
+	return Bound[K]{
+		value:     k,
+		inclusive: false,
+	}
+}

collections/keys/keys.go

@@ -8,7 +8,12 @@ import (
 type Order uint8
 
 const (
+	// OrderAscending defines an order going from the
+	// smallest key to the biggest key.
 	OrderAscending Order = iota
+	// OrderDescending defines an order going from the
+	// biggest key to the smallest. In the KVStore
+	// it equals to iterating in reverse.
 	OrderDescending
 )
 

collections/keys/pair.go

@@ -39,6 +39,30 @@ type Pair[K1 Key, K2 Key] struct {
 	p2 *K2
 }
 
+// K1 returns the first part of the key,
+// if present. If the key is not present
+// the zero value is returned.
+func (t Pair[K1, K2]) K1() K1 {
+	if t.p1 != nil {
+		return *t.p1
+	} else {
+		var x K1
+		return x
+	}
+}
+
+// K2 returns the second part of the key,
+// if present, If the key is not present
+// the zero value is returned.
+func (t Pair[K1, K2]) K2() K2 {
+	if t.p2 != nil {
+		return *t.p2
+	} else {
+		var x K2
+		return x
+	}
+}
+
 func (t Pair[K1, K2]) fkb1(b []byte) (int, K1) {
 	var k1 K1
 	i, p1 := k1.FromKeyBytes(b)