feat(orm): add ormfield.Codec (cosmos#10601)

* feat(orm): add ormvalue.Codec

* WIP

* WIP

* working tests

* update dep

* support more types, add docs

* comments

* address review comments

* updates

* add comment

api/buf.yaml

@@ -1,4 +1,5 @@
 version: v1
+name: buf.build/cosmos/cosmos-sdk
 breaking:
   use:
     - FILE

api/cosmos/orm/v1alpha1/orm.proto

@@ -43,8 +43,9 @@ message PrimaryKeyDescriptor {
   //   string fields support sorted iteration.
   //   - bytes are encoded as raw bytes in terminal segments and length-prefixed
   //   with a single byte in non-terminal segments. Because of this byte arrays
-  //   longer than 255 bytes cannot be used in keys and bytes fields should not
-  //   be assumed to be lexically sorted.
+  //   longer than 255 bytes are unsupported and bytes fields should not
+  //   be assumed to be lexically sorted. If you have a byte array longer than
+  //   255 bytes that you'd like to index, you should consider hashing it first.
   //   - int32, sint32, int64, sint64 are encoding as fixed width bytes with
   //   an encoding that enables sorted iteration.
   //   - google.protobuf.Timestamp and google.protobuf.Duration are encoded

buf.work.yaml

@@ -8,3 +8,4 @@ directories:
   - api
   - proto
   - third_party/proto
+  - orm/internal

orm/encoding/ormfield/bool.go

@@ -0,0 +1,54 @@
+package ormfield
+
+import (
+	io "io"
+
+	"google.golang.org/protobuf/reflect/protoreflect"
+)
+
+// BoolCodec encodes a bool value as a single byte 0 or 1.
+type BoolCodec struct{}
+
+func (b BoolCodec) Decode(r Reader) (protoreflect.Value, error) {
+	x, err := r.ReadByte()
+	return protoreflect.ValueOfBool(x != 0), err
+}
+
+var (
+	zeroBz = []byte{0}
+	oneBz  = []byte{1}
+)
+
+func (b BoolCodec) Encode(value protoreflect.Value, w io.Writer) error {
+	var err error
+	if value.Bool() {
+		_, err = w.Write(oneBz)
+	} else {
+		_, err = w.Write(zeroBz)
+	}
+	return err
+}
+
+func (b BoolCodec) Compare(v1, v2 protoreflect.Value) int {
+	b1 := v1.Bool()
+	b2 := v2.Bool()
+	if b1 == b2 {
+		return 0
+	} else if b1 {
+		return -1
+	} else {
+		return 1
+	}
+}
+
+func (b BoolCodec) IsOrdered() bool {
+	return false
+}
+
+func (b BoolCodec) FixedBufferSize() int {
+	return 1
+}
+
+func (b BoolCodec) ComputeBufferSize(protoreflect.Value) (int, error) {
+	return b.FixedBufferSize(), nil
+}

orm/encoding/ormfield/bytes.go

@@ -0,0 +1,99 @@
+package ormfield
+
+import (
+	"bytes"
+	"io"
+
+	"google.golang.org/protobuf/reflect/protoreflect"
+
+	"github.com/cosmos/cosmos-sdk/orm/types/ormerrors"
+)
+
+// BytesCodec encodes bytes as raw bytes. It errors if the byte array is longer
+// than 255 bytes.
+type BytesCodec struct{}
+
+func (b BytesCodec) FixedBufferSize() int {
+	return -1
+}
+
+func (b BytesCodec) ComputeBufferSize(value protoreflect.Value) (int, error) {
+	return bytesSize(value)
+}
+
+func bytesSize(value protoreflect.Value) (int, error) {
+	bz := value.Bytes()
+	n := len(bz)
+	if n > 255 {
+		return -1, ormerrors.BytesFieldTooLong
+	}
+	return n, nil
+}
+
+func (b BytesCodec) IsOrdered() bool {
+	return false
+}
+
+func (b BytesCodec) Decode(r Reader) (protoreflect.Value, error) {
+	bz, err := io.ReadAll(r)
+	return protoreflect.ValueOfBytes(bz), err
+}
+
+func (b BytesCodec) Encode(value protoreflect.Value, w io.Writer) error {
+	_, err := w.Write(value.Bytes())
+	return err
+}
+
+func (b BytesCodec) Compare(v1, v2 protoreflect.Value) int {
+	return bytes.Compare(v1.Bytes(), v2.Bytes())
+}
+
+// NonTerminalBytesCodec encodes bytes as raw bytes length prefixed by a single
+// byte. It errors if the byte array is longer than 255 bytes.
+type NonTerminalBytesCodec struct{}
+
+func (b NonTerminalBytesCodec) FixedBufferSize() int {
+	return -1
+}
+
+func (b NonTerminalBytesCodec) ComputeBufferSize(value protoreflect.Value) (int, error) {
+	n, err := bytesSize(value)
+	return n + 1, err
+}
+
+func (b NonTerminalBytesCodec) IsOrdered() bool {
+	return false
+}
+
+func (b NonTerminalBytesCodec) Compare(v1, v2 protoreflect.Value) int {
+	return bytes.Compare(v1.Bytes(), v2.Bytes())
+}
+
+func (b NonTerminalBytesCodec) Decode(r Reader) (protoreflect.Value, error) {
+	n, err := r.ReadByte()
+	if err != nil {
+		return protoreflect.Value{}, err
+	}
+
+	if n == 0 {
+		return protoreflect.ValueOfBytes([]byte{}), nil
+	}
+
+	bz := make([]byte, n)
+	_, err = r.Read(bz)
+	return protoreflect.ValueOfBytes(bz), err
+}
+
+func (b NonTerminalBytesCodec) Encode(value protoreflect.Value, w io.Writer) error {
+	bz := value.Bytes()
+	n := len(bz)
+	if n > 255 {
+		return ormerrors.BytesFieldTooLong
+	}
+	_, err := w.Write([]byte{byte(n)})
+	if err != nil {
+		return err
+	}
+	_, err = w.Write(bz)
+	return err
+}

orm/encoding/ormfield/codec.go

@@ -0,0 +1,106 @@
+package ormfield
+
+import (
+	"io"
+
+	"github.com/cosmos/cosmos-sdk/orm/types/ormerrors"
+
+	"google.golang.org/protobuf/types/known/durationpb"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	"google.golang.org/protobuf/reflect/protoreflect"
+)
+
+// Codec defines an interface for decoding and encoding values in ORM index keys.
+type Codec interface {
+
+	// Decode decodes a value in a key.
+	Decode(r Reader) (protoreflect.Value, error)
+
+	// Encode encodes a value in a key.
+	Encode(value protoreflect.Value, w io.Writer) error
+
+	// Compare compares two values of this type and should primarily be used
+	// for testing.
+	Compare(v1, v2 protoreflect.Value) int
+
+	// IsOrdered returns true if callers can always assume that this ordering
+	// is suitable for sorted iteration.
+	IsOrdered() bool
+
+	// FixedBufferSize returns a positive value if encoders should assume a
+	// fixed size buffer for encoding. Encoders will use at most this much size
+	// to encode the value.
+	FixedBufferSize() int
+
+	// ComputeBufferSize estimates the buffer size needed to encode the field.
+	// Encoders will use at most this much size to encode the value.
+	ComputeBufferSize(value protoreflect.Value) (int, error)
+}
+
+type Reader interface {
+	io.Reader
+	io.ByteReader
+}
+
+var (
+	timestampMsgType  = (&timestamppb.Timestamp{}).ProtoReflect().Type()
+	timestampFullName = timestampMsgType.Descriptor().FullName()
+	durationMsgType   = (&durationpb.Duration{}).ProtoReflect().Type()
+	durationFullName  = durationMsgType.Descriptor().FullName()
+)
+
+// GetCodec returns the Codec for the provided field if one is defined.
+// nonTerminal should be set to true if this value is being encoded as a
+// non-terminal segment of a multi-part key.
+func GetCodec(field protoreflect.FieldDescriptor, nonTerminal bool) (Codec, error) {
+	if field == nil {
+		return nil, ormerrors.UnsupportedKeyField.Wrap("nil field")
+	}
+	if field.IsList() {
+		return nil, ormerrors.UnsupportedKeyField.Wrapf("repeated field %s", field.FullName())
+	}
+
+	if field.ContainingOneof() != nil {
+		return nil, ormerrors.UnsupportedKeyField.Wrapf("oneof field %s", field.FullName())
+	}
+
+	switch field.Kind() {
+	case protoreflect.BytesKind:
+		if nonTerminal {
+			return NonTerminalBytesCodec{}, nil
+		} else {
+			return BytesCodec{}, nil
+		}
+	case protoreflect.StringKind:
+		if nonTerminal {
+			return NonTerminalStringCodec{}, nil
+		} else {
+			return StringCodec{}, nil
+		}
+	case protoreflect.Uint32Kind, protoreflect.Fixed32Kind:
+		return Uint32Codec{}, nil
+	case protoreflect.Uint64Kind, protoreflect.Fixed64Kind:
+		return Uint64Codec{}, nil
+	case protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Sfixed32Kind:
+		return Int32Codec{}, nil
+	case protoreflect.Int64Kind, protoreflect.Sint64Kind, protoreflect.Sfixed64Kind:
+		return Int64Codec{}, nil
+	case protoreflect.BoolKind:
+		return BoolCodec{}, nil
+	case protoreflect.EnumKind:
+		return EnumCodec{}, nil
+	case protoreflect.MessageKind:
+		msgName := field.Message().FullName()
+		switch msgName {
+		case timestampFullName:
+			return TimestampCodec{}, nil
+		case durationFullName:
+			return DurationCodec{}, nil
+		default:
+			return nil, ormerrors.UnsupportedKeyField.Wrapf("%s of type %s", field.FullName(), msgName)
+		}
+	default:
+		return nil, ormerrors.UnsupportedKeyField.Wrapf("%s of kind %s", field.FullName(), field.Kind())
+	}
+}

orm/encoding/ormfield/codec_test.go

@@ -0,0 +1,87 @@
+package ormfield_test
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/cosmos/cosmos-sdk/orm/encoding/ormfield"
+
+	"google.golang.org/protobuf/reflect/protoreflect"
+	"gotest.tools/v3/assert"
+	"pgregory.net/rapid"
+
+	"github.com/cosmos/cosmos-sdk/orm/types/ormerrors"
+
+	"github.com/cosmos/cosmos-sdk/orm/internal/testutil"
+)
+
+func TestCodec(t *testing.T) {
+	for _, ks := range testutil.TestFieldSpecs {
+		testCodec(t, ks)
+	}
+}
+
+func testCodec(t *testing.T, spec testutil.TestFieldSpec) {
+	t.Run(fmt.Sprintf("%s %v", spec.FieldName, false), func(t *testing.T) {
+		testCodecNT(t, spec.FieldName, spec.Gen, false)
+	})
+	t.Run(fmt.Sprintf("%s %v", spec.FieldName, true), func(t *testing.T) {
+		testCodecNT(t, spec.FieldName, spec.Gen, true)
+	})
+}
+
+func testCodecNT(t *testing.T, fname protoreflect.Name, generator *rapid.Generator, nonTerminal bool) {
+	cdc, err := testutil.MakeTestCodec(fname, nonTerminal)
+	assert.NilError(t, err)
+	rapid.Check(t, func(t *rapid.T) {
+		x := protoreflect.ValueOf(generator.Draw(t, string(fname)))
+		bz1 := checkEncodeDecodeSize(t, x, cdc)
+		if cdc.IsOrdered() {
+			y := protoreflect.ValueOf(generator.Draw(t, fmt.Sprintf("%s 2", fname)))
+			bz2 := checkEncodeDecodeSize(t, y, cdc)
+			assert.Equal(t, cdc.Compare(x, y), bytes.Compare(bz1, bz2))
+		}
+	})
+}
+
+func checkEncodeDecodeSize(t *rapid.T, x protoreflect.Value, cdc ormfield.Codec) []byte {
+	buf := &bytes.Buffer{}
+	err := cdc.Encode(x, buf)
+	assert.NilError(t, err)
+	bz := buf.Bytes()
+	size, err := cdc.ComputeBufferSize(x)
+	assert.NilError(t, err)
+	assert.Assert(t, size >= len(bz))
+	fixedSize := cdc.FixedBufferSize()
+	if fixedSize > 0 {
+		assert.Equal(t, fixedSize, size)
+	}
+	y, err := cdc.Decode(bytes.NewReader(bz))
+	assert.NilError(t, err)
+	assert.Equal(t, 0, cdc.Compare(x, y))
+	return bz
+}
+
+func TestUnsupportedFields(t *testing.T) {
+	_, err := ormfield.GetCodec(nil, false)
+	assert.ErrorContains(t, err, ormerrors.UnsupportedKeyField.Error())
+	_, err = ormfield.GetCodec(testutil.GetTestField("repeated"), false)
+	assert.ErrorContains(t, err, ormerrors.UnsupportedKeyField.Error())
+	_, err = ormfield.GetCodec(testutil.GetTestField("map"), false)
+	assert.ErrorContains(t, err, ormerrors.UnsupportedKeyField.Error())
+	_, err = ormfield.GetCodec(testutil.GetTestField("msg"), false)
+	assert.ErrorContains(t, err, ormerrors.UnsupportedKeyField.Error())
+	_, err = ormfield.GetCodec(testutil.GetTestField("oneof"), false)
+	assert.ErrorContains(t, err, ormerrors.UnsupportedKeyField.Error())
+}
+
+func TestNTBytesTooLong(t *testing.T) {
+	cdc, err := ormfield.GetCodec(testutil.GetTestField("bz"), true)
+	assert.NilError(t, err)
+	buf := &bytes.Buffer{}
+	bz := protoreflect.ValueOfBytes(make([]byte, 256))
+	assert.ErrorContains(t, cdc.Encode(bz, buf), ormerrors.BytesFieldTooLong.Error())
+	_, err = cdc.ComputeBufferSize(bz)
+	assert.ErrorContains(t, err, ormerrors.BytesFieldTooLong.Error())
+}