feat(collections): Add Triple key (#17024)

Co-authored-by: unknown unknown <unknown@unknown>

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

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

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("<nil>")
+	}
+
+	b.WriteString(", ")
+	if key.k2 != nil {
+		b.WriteByte('"')
+		b.WriteString(t.keyCodec2.Stringify(*key.k2))
+		b.WriteByte('"')
+	} else {
+		b.WriteString("<nil>")
+	}
+
+	b.WriteString(", ")
+	if key.k3 != nil {
+		b.WriteByte('"')
+		b.WriteString(t.keyCodec3.Stringify(*key.k3))
+		b.WriteByte('"')
+	} else {
+		b.WriteString("<nil>")
+	}
+
+	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),
+	}
+}

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)
+}