diff --git a/collections/CHANGELOG.md b/collections/CHANGELOG.md index 42c9e506fb39..f0778efdfc92 100644 --- a/collections/CHANGELOG.md +++ b/collections/CHANGELOG.md @@ -31,6 +31,9 @@ Ref: https://keepachangelog.com/en/1.0.0/ ## [Unreleased] +### Features + * [#17656](https://github.com/cosmos/cosmos-sdk/pull/17656) – Introduces `Vec`, a collection type that allows to represent a growable array on top of a KVStore. + ## [v0.4.0](https://github.com/cosmos/cosmos-sdk/releases/tag/collections%2Fv0.4.0) ### Features diff --git a/collections/vec.go b/collections/vec.go new file mode 100644 index 000000000000..58ea87fe62ac --- /dev/null +++ b/collections/vec.go @@ -0,0 +1,141 @@ +package collections + +import ( + "context" + "errors" + "fmt" + + "cosmossdk.io/collections/codec" +) + +var ( + // ErrEmptyVec is returned when trying to pop an element from an empty Vec. + ErrEmptyVec = errors.New("vec is empty") + // ErrOutOfBounds is returned when trying to do an operation on an index that is out of bounds. + ErrOutOfBounds = errors.New("vec index is out of bounds") +) + +const ( + VecElementsNameSuffix = "_elements" + VecLengthNameSuffix = "_length" + VecElementsPrefixSuffix = 0x0 + VecLengthPrefixSuffix = 0x1 +) + +// NewVec creates a new Vec instance. Since Vec relies on two collections, one for the length +// and the other for the elements, it will register two state objects on the schema builder. +// The first is the length which is an item, whose prefix is the provided prefix with a suffix +// which equals to VecLengthPrefixSuffix, the name is also suffixed with VecLengthNameSuffix. +// The second is the elements which is a map, whose prefix is the provided prefix with a suffix +// which equals to VecElementsPrefixSuffix, the name is also suffixed with VecElementsNameSuffix. +func NewVec[T any](sb *SchemaBuilder, prefix Prefix, name string, vc codec.ValueCodec[T]) Vec[T] { + return Vec[T]{ + length: NewItem(sb, append(prefix, VecLengthPrefixSuffix), name+VecLengthNameSuffix, Uint64Value), + elements: NewMap(sb, append(prefix, VecElementsPrefixSuffix), name+VecElementsNameSuffix, Uint64Key, vc), + } +} + +// Vec works like a slice sitting on top of a KVStore. +// It relies on two collections, one for the length which is an Item[uint64], +// the other for the elements which is a Map[uint64, T]. +type Vec[T any] struct { + length Item[uint64] + elements Map[uint64, T] +} + +// Push adds an element to the end of the Vec. +func (v Vec[T]) Push(ctx context.Context, elem T) error { + length, err := v.length.Get(ctx) + if err != nil && !errors.Is(err, ErrNotFound) { + return err + } + err = v.elements.Set(ctx, length, elem) + if err != nil { + return err + } + err = v.length.Set(ctx, length+1) + if err != nil { + return err + } + return nil +} + +// Pop removes an element from the end of the Vec and returns it. Fails +// if the Vec is empty. +func (v Vec[T]) Pop(ctx context.Context) (elem T, err error) { + length, err := v.length.Get(ctx) + if err != nil && !errors.Is(err, ErrNotFound) { + return elem, err + } + if length == 0 { + return elem, ErrEmptyVec + } + length -= 1 + elem, err = v.elements.Get(ctx, length) + if err != nil { + return elem, err + } + err = v.elements.Remove(ctx, length) + if err != nil { + return elem, err + } + err = v.length.Set(ctx, length) + if err != nil { + return elem, err + } + return elem, nil +} + +// Replace replaces an element at a given index. Fails if the index is out of bounds. +func (v Vec[T]) Replace(ctx context.Context, index uint64, elem T) error { + length, err := v.length.Get(ctx) + if err != nil && !errors.Is(err, ErrNotFound) { + return err + } + if index >= length { + return fmt.Errorf("%w: length %d", ErrOutOfBounds, length) + } + return v.elements.Set(ctx, index, elem) +} + +// Get returns an element at a given index. Returns ErrOutOfBounds +// if the index is out of bounds. +func (v Vec[T]) Get(ctx context.Context, index uint64) (elem T, err error) { + elem, err = v.elements.Get(ctx, index) + switch { + case err == nil: + return elem, nil + case errors.Is(err, ErrNotFound): + return elem, fmt.Errorf("%w: index %d", ErrOutOfBounds, index) + default: + return elem, err + } +} + +// Len returns the length of the Vec. +func (v Vec[T]) Len(ctx context.Context) (uint64, error) { + length, err := v.length.Get(ctx) + switch { + // no error, return length as the vec is populated + case err == nil: + return length, nil + // not found, return 0 as the vec is empty + case errors.Is(err, ErrNotFound): + return 0, nil + // something else happened + default: + return 0, err + } +} + +// Iterate iterates over the Vec. It returns an Iterator whose key is the index +// and the value is the element at that index. +func (v Vec[T]) Iterate(ctx context.Context, rng Ranger[uint64]) (Iterator[uint64, T], error) { + return v.elements.Iterate(ctx, rng) +} + +// Walk walks over the Vec. It calls the walkFn for each element in the Vec, +// where the key is the index and the value is the element at that index. +func (v Vec[T]) Walk(ctx context.Context, rng Ranger[uint64], walkFn func(index uint64, elem T) (stop bool, err error)) error { + return v.elements.Walk(ctx, rng, walkFn) +} diff --git a/collections/vec_test.go b/collections/vec_test.go new file mode 100644 index 000000000000..d7640aec708d --- /dev/null +++ b/collections/vec_test.go @@ -0,0 +1,63 @@ +package collections + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestVec(t *testing.T) { + sk, ctx := deps() + schemaBuilder := NewSchemaBuilder(sk) + vec := NewVec(schemaBuilder, NewPrefix(0), "vec", StringValue) + _, err := schemaBuilder.Build() + require.NoError(t, err) + + // length when empty + length, err := vec.Len(ctx) + require.NoError(t, err) + require.Equal(t, uint64(0), length) + + // pop when empty should error with an empty vec error + _, err = vec.Pop(ctx) + require.ErrorIs(t, err, ErrEmptyVec) + + // replace when out of bounds should error with an out of bounds error + err = vec.Replace(ctx, 0, "foo") + require.ErrorIs(t, err, ErrOutOfBounds) + + // get out of bounds should error with an out of bounds error + _, err = vec.Get(ctx, 0) + require.ErrorIs(t, err, ErrOutOfBounds) + + // push + err = vec.Push(ctx, "foo") + require.NoError(t, err) + + // push more + err = vec.Push(ctx, "bar") + require.NoError(t, err) + + // check length + length, err = vec.Len(ctx) + require.NoError(t, err) + require.Equal(t, uint64(2), length) + + // get + v, err := vec.Get(ctx, 0) + require.NoError(t, err) + require.Equal(t, "foo", v) + + // replace + err = vec.Replace(ctx, 0, "bar") + require.NoError(t, err) + + v, err = vec.Get(ctx, 0) + require.NoError(t, err) + require.Equal(t, "bar", v) + + // pop + v, err = vec.Pop(ctx) + require.NoError(t, err) + require.Equal(t, "bar", v) +}