From dd585bac075596cc075643e36366e4ed2381c6c9 Mon Sep 17 00:00:00 2001 From: testinginprod <98415576+testinginprod@users.noreply.github.com> Date: Mon, 17 Jul 2023 19:49:39 +0200 Subject: [PATCH] feat(collections): Add Triple key (#17024) Co-authored-by: unknown unknown --- collections/CHANGELOG.md | 6 + collections/README.md | 52 +++++++ collections/triple.go | 310 +++++++++++++++++++++++++++++++++++++ collections/triple_test.go | 52 +++++++ 4 files changed, 420 insertions(+) create mode 100644 collections/triple.go create mode 100644 collections/triple_test.go diff --git a/collections/CHANGELOG.md b/collections/CHANGELOG.md index 92a7e92c5890..9b04bc86e9cd 100644 --- a/collections/CHANGELOG.md +++ b/collections/CHANGELOG.md @@ -29,6 +29,12 @@ Ref: https://keepachangelog.com/en/1.0.0/ # Changelog +## [Unreleased] + +### Features + +* [#17024](https://github.com/cosmos/cosmos-sdk/pull/17024) - Introduces `Triple`, a composite key with three keys. + ## [v0.3.0](https://github.com/cosmos/cosmos-sdk/releases/tag/collections%2Fv0.3.0) ### Features diff --git a/collections/README.md b/collections/README.md index bd25d2e1cccf..72d5f7172a58 100644 --- a/collections/README.md +++ b/collections/README.md @@ -1118,6 +1118,58 @@ func (k Keeper) GetAccount(ctx sdk.context, addr sdk.AccAddress) (sdk.AccountI, } ``` +## Triple key + +The `collections.Triple` is a special type of key composed of three keys, it's identical to `collections.Pair`. + +Let's see an example. + +```go +package example + +import ( + "context" + + "cosmossdk.io/collections" + storetypes "cosmossdk.io/store/types" + "github.com/cosmos/cosmos-sdk/codec" +) + +type AccAddress = string +type ValAddress = string + +type Keeper struct { + // let's simulate we have redelegations which are stored as a triple key composed of + // the delegator, the source validator and the destination validator. + Redelegations collections.KeySet[collections.Triple[AccAddress, ValAddress, ValAddress]] +} + +func NewKeeper(storeKey *storetypes.KVStoreKey) Keeper { + sb := collections.NewSchemaBuilder(sdk.OpenKVStore(storeKey)) + return Keeper{ + Redelegations: collections.NewKeySet(sb, collections.NewPrefix(0), "redelegations", collections.TripleKeyCodec(collections.StringKey, collections.StringKey, collections.StringKey) + } +} + +// RedelegationsByDelegator iterates over all the redelegations of a given delegator and calls onResult providing +// each redelegation from source validator towards the destination validator. +func (k Keeper) RedelegationsByDelegator(ctx context.Context, delegator AccAddress, onResult func(src, dst ValAddress) (stop bool, err error)) error { + rng := collections.NewPrefixedTripleRange[AccAddress, ValAddress, ValAddress](delegator) + return k.Redelegations.Walk(ctx, rng, func(key collections.Triple[AccAddress, ValAddress, ValAddress]) (stop bool, err error) { + return onResult(key.K2(), key.K3()) + }) +} + +// RedelegationsByDelegatorAndValidator iterates over all the redelegations of a given delegator and its source validator and calls onResult for each +// destination validator. +func (k Keeper) RedelegationsByDelegatorAndValidator(ctx context.Context, delegator AccAddress, validator ValAddress, onResult func(dst ValAddress) (stop bool, err error)) error { + rng := collections.NewSuperPrefixedTripleRange[AccAddress, ValAddress, ValAddress](delegator, validator) + return k.Redelegations.Walk(ctx, rng, func(key collections.Triple[AccAddress, ValAddress, ValAddress]) (stop bool, err error) { + return onResult(key.K3()) + }) +} +``` + ## Advanced Usages ### Alternative Value Codec diff --git a/collections/triple.go b/collections/triple.go new file mode 100644 index 000000000000..b465b25b877d --- /dev/null +++ b/collections/triple.go @@ -0,0 +1,310 @@ +package collections + +import ( + "encoding/json" + "fmt" + "strings" + + "cosmossdk.io/collections/codec" +) + +// Triple defines a multipart key composed of three keys. +type Triple[K1, K2, K3 any] struct { + k1 *K1 + k2 *K2 + k3 *K3 +} + +// Join3 instantiates a new Triple instance composed of the three provided keys, in order. +func Join3[K1, K2, K3 any](k1 K1, k2 K2, k3 K3) Triple[K1, K2, K3] { + return Triple[K1, K2, K3]{&k1, &k2, &k3} +} + +// K1 returns the first part of the key. If nil, the zero value is returned. +func (t Triple[K1, K2, K3]) K1() (x K1) { + if t.k1 != nil { + return *t.k1 + } + return x +} + +// K2 returns the second part of the key. If nil, the zero value is returned. +func (t Triple[K1, K2, K3]) K2() (x K2) { + if t.k2 != nil { + return *t.k2 + } + return x +} + +// K3 returns the third part of the key. If nil, the zero value is returned. +func (t Triple[K1, K2, K3]) K3() (x K3) { + if t.k3 != nil { + return *t.k3 + } + return x +} + +// TriplePrefix creates a new Triple instance composed only of the first part of the key. +func TriplePrefix[K1, K2, K3 any](k1 K1) Triple[K1, K2, K3] { + return Triple[K1, K2, K3]{k1: &k1} +} + +// TripleSuperPrefix creates a new Triple instance composed only of the first two parts of the key. +func TripleSuperPrefix[K1, K2, K3 any](k1 K1, k2 K2) Triple[K1, K2, K3] { + return Triple[K1, K2, K3]{k1: &k1, k2: &k2} +} + +// TripleKeyCodec instantiates a new KeyCodec instance that can encode the Triple, given +// the KeyCodecs of the three parts of the key, in order. +func TripleKeyCodec[K1, K2, K3 any](keyCodec1 codec.KeyCodec[K1], keyCodec2 codec.KeyCodec[K2], keyCodec3 codec.KeyCodec[K3]) codec.KeyCodec[Triple[K1, K2, K3]] { + return tripleKeyCodec[K1, K2, K3]{ + keyCodec1: keyCodec1, + keyCodec2: keyCodec2, + keyCodec3: keyCodec3, + } +} + +type tripleKeyCodec[K1, K2, K3 any] struct { + keyCodec1 codec.KeyCodec[K1] + keyCodec2 codec.KeyCodec[K2] + keyCodec3 codec.KeyCodec[K3] +} + +type jsonTripleKey [3]json.RawMessage + +func (t tripleKeyCodec[K1, K2, K3]) EncodeJSON(value Triple[K1, K2, K3]) ([]byte, error) { + json1, err := t.keyCodec1.EncodeJSON(*value.k1) + if err != nil { + return nil, err + } + + json2, err := t.keyCodec2.EncodeJSON(*value.k2) + if err != nil { + return nil, err + } + + json3, err := t.keyCodec3.EncodeJSON(*value.k3) + if err != nil { + return nil, err + } + + return json.Marshal(jsonTripleKey{json1, json2, json3}) +} + +func (t tripleKeyCodec[K1, K2, K3]) DecodeJSON(b []byte) (Triple[K1, K2, K3], error) { + var jsonKey jsonTripleKey + err := json.Unmarshal(b, &jsonKey) + if err != nil { + return Triple[K1, K2, K3]{}, err + } + + key1, err := t.keyCodec1.DecodeJSON(jsonKey[0]) + if err != nil { + return Triple[K1, K2, K3]{}, err + } + + key2, err := t.keyCodec2.DecodeJSON(jsonKey[1]) + if err != nil { + return Triple[K1, K2, K3]{}, err + } + + key3, err := t.keyCodec3.DecodeJSON(jsonKey[2]) + if err != nil { + return Triple[K1, K2, K3]{}, err + } + + return Join3(key1, key2, key3), nil +} + +func (t tripleKeyCodec[K1, K2, K3]) Stringify(key Triple[K1, K2, K3]) string { + b := new(strings.Builder) + b.WriteByte('(') + if key.k1 != nil { + b.WriteByte('"') + b.WriteString(t.keyCodec1.Stringify(*key.k1)) + b.WriteByte('"') + } else { + b.WriteString("") + } + + b.WriteString(", ") + if key.k2 != nil { + b.WriteByte('"') + b.WriteString(t.keyCodec2.Stringify(*key.k2)) + b.WriteByte('"') + } else { + b.WriteString("") + } + + b.WriteString(", ") + if key.k3 != nil { + b.WriteByte('"') + b.WriteString(t.keyCodec3.Stringify(*key.k3)) + b.WriteByte('"') + } else { + b.WriteString("") + } + + b.WriteByte(')') + return b.String() +} + +func (t tripleKeyCodec[K1, K2, K3]) KeyType() string { + return fmt.Sprintf("Triple[%s,%s,%s]", t.keyCodec1.KeyType(), t.keyCodec2.KeyType(), t.keyCodec3.KeyType()) +} + +func (t tripleKeyCodec[K1, K2, K3]) Encode(buffer []byte, key Triple[K1, K2, K3]) (int, error) { + writtenTotal := 0 + if key.k1 != nil { + written, err := t.keyCodec1.EncodeNonTerminal(buffer, *key.k1) + if err != nil { + return 0, err + } + writtenTotal += written + } + if key.k2 != nil { + written, err := t.keyCodec2.EncodeNonTerminal(buffer[writtenTotal:], *key.k2) + if err != nil { + return 0, err + } + writtenTotal += written + } + if key.k3 != nil { + written, err := t.keyCodec3.Encode(buffer[writtenTotal:], *key.k3) + if err != nil { + return 0, err + } + writtenTotal += written + } + return writtenTotal, nil +} + +func (t tripleKeyCodec[K1, K2, K3]) Decode(buffer []byte) (int, Triple[K1, K2, K3], error) { + readTotal := 0 + read, key1, err := t.keyCodec1.DecodeNonTerminal(buffer) + if err != nil { + return 0, Triple[K1, K2, K3]{}, err + } + readTotal += read + read, key2, err := t.keyCodec2.DecodeNonTerminal(buffer[readTotal:]) + if err != nil { + return 0, Triple[K1, K2, K3]{}, err + } + readTotal += read + read, key3, err := t.keyCodec3.Decode(buffer[readTotal:]) + if err != nil { + return 0, Triple[K1, K2, K3]{}, err + } + readTotal += read + return readTotal, Join3(key1, key2, key3), nil +} + +func (t tripleKeyCodec[K1, K2, K3]) Size(key Triple[K1, K2, K3]) int { + size := 0 + if key.k1 != nil { + size += t.keyCodec1.SizeNonTerminal(*key.k1) + } + if key.k2 != nil { + size += t.keyCodec2.SizeNonTerminal(*key.k2) + } + if key.k3 != nil { + size += t.keyCodec3.Size(*key.k3) + } + return size +} + +func (t tripleKeyCodec[K1, K2, K3]) EncodeNonTerminal(buffer []byte, key Triple[K1, K2, K3]) (int, error) { + writtenTotal := 0 + if key.k1 != nil { + written, err := t.keyCodec1.EncodeNonTerminal(buffer, *key.k1) + if err != nil { + return 0, err + } + writtenTotal += written + } + if key.k2 != nil { + written, err := t.keyCodec2.EncodeNonTerminal(buffer[writtenTotal:], *key.k2) + if err != nil { + return 0, err + } + writtenTotal += written + } + if key.k3 != nil { + written, err := t.keyCodec3.EncodeNonTerminal(buffer[writtenTotal:], *key.k3) + if err != nil { + return 0, err + } + writtenTotal += written + } + return writtenTotal, nil +} + +func (t tripleKeyCodec[K1, K2, K3]) DecodeNonTerminal(buffer []byte) (int, Triple[K1, K2, K3], error) { + readTotal := 0 + read, key1, err := t.keyCodec1.DecodeNonTerminal(buffer) + if err != nil { + return 0, Triple[K1, K2, K3]{}, err + } + readTotal += read + read, key2, err := t.keyCodec2.DecodeNonTerminal(buffer[readTotal:]) + if err != nil { + return 0, Triple[K1, K2, K3]{}, err + } + readTotal += read + read, key3, err := t.keyCodec3.DecodeNonTerminal(buffer[readTotal:]) + if err != nil { + return 0, Triple[K1, K2, K3]{}, err + } + readTotal += read + return readTotal, Join3(key1, key2, key3), nil +} + +func (t tripleKeyCodec[K1, K2, K3]) SizeNonTerminal(key Triple[K1, K2, K3]) int { + size := 0 + if key.k1 != nil { + size += t.keyCodec1.SizeNonTerminal(*key.k1) + } + if key.k2 != nil { + size += t.keyCodec2.SizeNonTerminal(*key.k2) + } + if key.k3 != nil { + size += t.keyCodec3.SizeNonTerminal(*key.k3) + } + return size +} + +func PrefixTripleRange[K1, K2, K3 any](k1 K1) Ranger[Triple[K1, K2, K3]] { + key := TriplePrefix[K1, K2, K3](k1) + return &Range[Triple[K1, K2, K3]]{ + start: RangeKeyExact(key), + end: RangeKeyPrefixEnd(key), + } +} + +func SuperPrefixTripleRange[K1, K2, K3 any](k1 K1, k2 K2) Ranger[Triple[K1, K2, K3]] { + key := TripleSuperPrefix[K1, K2, K3](k1, k2) + return &Range[Triple[K1, K2, K3]]{ + start: RangeKeyExact(key), + end: RangeKeyPrefixEnd(key), + } +} + +// NewPrefixedTripleRange provides a Range for all keys prefixed with the given +// first part of the Triple key. +func NewPrefixedTripleRange[K1, K2, K3 any](prefix K1) Ranger[Triple[K1, K2, K3]] { + key := TriplePrefix[K1, K2, K3](prefix) + return &Range[Triple[K1, K2, K3]]{ + start: RangeKeyExact(key), + end: RangeKeyPrefixEnd(key), + } +} + +// NewSuperPrefixedTripleRange provides a Range for all keys prefixed with the given +// first and second parts of the Triple key. +func NewSuperPrefixedTripleRange[K1, K2, K3 any](k1 K1, k2 K2) Ranger[Triple[K1, K2, K3]] { + key := TripleSuperPrefix[K1, K2, K3](k1, k2) + return &Range[Triple[K1, K2, K3]]{ + start: RangeKeyExact(key), + end: RangeKeyPrefixEnd(key), + } +} diff --git a/collections/triple_test.go b/collections/triple_test.go new file mode 100644 index 000000000000..fea1b662c15c --- /dev/null +++ b/collections/triple_test.go @@ -0,0 +1,52 @@ +package collections_test + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "cosmossdk.io/collections" + "cosmossdk.io/collections/colltest" +) + +func TestTriple(t *testing.T) { + kc := collections.TripleKeyCodec(collections.Uint64Key, collections.StringKey, collections.BytesKey) + + t.Run("conformance", func(t *testing.T) { + colltest.TestKeyCodec(t, kc, collections.Join3(uint64(1), "2", []byte("3"))) + }) +} + +func TestTripleRange(t *testing.T) { + sk, ctx := colltest.MockStore() + schema := collections.NewSchemaBuilder(sk) + // this is a key composed of 3 parts: uint64, string, []byte + kc := collections.TripleKeyCodec(collections.Uint64Key, collections.StringKey, collections.BytesKey) + + keySet := collections.NewKeySet(schema, collections.NewPrefix(0), "triple", kc) + + keys := []collections.Triple[uint64, string, []byte]{ + collections.Join3(uint64(1), "A", []byte("1")), + collections.Join3(uint64(1), "A", []byte("2")), + collections.Join3(uint64(1), "B", []byte("3")), + collections.Join3(uint64(2), "B", []byte("4")), + } + + for _, k := range keys { + require.NoError(t, keySet.Set(ctx, k)) + } + + // we prefix over (1) we expect 3 results + iter, err := keySet.Iterate(ctx, collections.NewPrefixedTripleRange[uint64, string, []byte](uint64(1))) + require.NoError(t, err) + gotKeys, err := iter.Keys() + require.NoError(t, err) + require.Equal(t, keys[:3], gotKeys) + + // we super prefix over Join(1, "A") we expect 2 results + iter, err = keySet.Iterate(ctx, collections.NewSuperPrefixedTripleRange[uint64, string, []byte](1, "A")) + require.NoError(t, err) + gotKeys, err = iter.Keys() + require.NoError(t, err) + require.Equal(t, keys[:2], gotKeys) +}