Merge pull request #140 from thanos-community/unique_unmarshal

features/unmarshal: implement unmarshal_unique

Author: vmg

README.md

@@ -53,6 +53,20 @@ The following features can be generated:
 
     - `func (p *YourProto) CloneMessageVT() proto.Message`: this function behaves like the above `p.CloneVT()`, but provides a uniform signature in order to be accessible via type assertions even if the type is not known at compile time. This allows implementing a generic `func CloneVT(proto.Message)` without reflection. If the receiver `p` is `nil`, a typed `nil` pointer of the message type will be returned inside a `proto.Message` interface.
 
+### Field Options
+
+- `unique` is a field option available on strings. If it is set to `true` then all all strings are interned using [unique.Make](https://pkg.go.dev/unique#Make). Go 1.23+ is needed. `unmarshal_unsafe` takes precendence over `unique`. Example usage:
+
+```
+import "github.com/planetscale/vtprotobuf/vtproto/ext.proto";
+
+message Label {
+    string name  = 1 [(vtproto.options).unique = true];
+    string value = 2 [(vtproto.options).unique = true];
+}
+```
+
+
 ## Usage
 
 1. Install `protoc-gen-go-vtproto`:

features/unmarshal/unmarshal.go

@@ -12,9 +12,11 @@ import (
 
 	"google.golang.org/protobuf/compiler/protogen"
 	"google.golang.org/protobuf/encoding/protowire"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 
 	"github.com/planetscale/vtprotobuf/generator"
+	"github.com/planetscale/vtprotobuf/vtproto"
 )
 
 func init() {
@@ -156,6 +158,8 @@ func (p *unmarshal) declareMapField(varName string, nullable bool, field *protog
 }
 
 func (p *unmarshal) mapField(varName string, field *protogen.Field) {
+	unique := proto.GetExtension(field.Desc.Options(), vtproto.E_Options).(*vtproto.Opts).GetUnique()
+
 	switch field.Desc.Kind() {
 	case protoreflect.DoubleKind:
 		p.P(`var `, varName, `temp uint64`)
@@ -193,13 +197,20 @@ func (p *unmarshal) mapField(varName string, field *protogen.Field) {
 		p.P(`if postStringIndex`, varName, ` > l {`)
 		p.P(`return `, p.Ident("io", `ErrUnexpectedEOF`))
 		p.P(`}`)
-		if p.unsafe {
+		switch {
+		case p.unsafe:
 			p.P(`if intStringLen`, varName, ` == 0 {`)
 			p.P(varName, ` = ""`)
 			p.P(`} else {`)
 			p.P(varName, ` = `, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen`, varName, `)`)
 			p.P(`}`)
-		} else {
+		case unique:
+			p.P(`if intStringLen`, varName, ` == 0 {`)
+			p.P(varName, ` = ""`)
+			p.P(`} else {`)
+			p.P(varName, ` = `, p.Ident("unique", `Make`), `[string](`, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen`, varName, `)).Value()`)
+			p.P(`}`)
+		default:
 			p.P(varName, ` = `, "string", `(dAtA[iNdEx:postStringIndex`, varName, `])`)
 		}
 		p.P(`iNdEx = postStringIndex`, varName)
@@ -409,6 +420,8 @@ func (p *unmarshal) fieldItem(field *protogen.Field, fieldname string, message *
 			p.P(`m.`, fieldname, ` = &b`)
 		}
 	case protoreflect.StringKind:
+		unique := proto.GetExtension(field.Desc.Options(), vtproto.E_Options).(*vtproto.Opts).GetUnique()
+
 		p.P(`var stringLen uint64`)
 		p.decodeVarint("stringLen", "uint64")
 		p.P(`intStringLen := int(stringLen)`)
@@ -423,12 +436,19 @@ func (p *unmarshal) fieldItem(field *protogen.Field, fieldname string, message *
 		p.P(`return `, p.Ident("io", `ErrUnexpectedEOF`))
 		p.P(`}`)
 		str := "string(dAtA[iNdEx:postIndex])"
-		if p.unsafe {
+		switch {
+		case p.unsafe:
 			str = "stringValue"
 			p.P(`var stringValue string`)
 			p.P(`if intStringLen > 0 {`)
 			p.P(`stringValue = `, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen)`)
 			p.P(`}`)
+		case unique:
+			str = "stringValue"
+			p.P(`var stringValue string`)
+			p.P(`if intStringLen > 0 {`)
+			p.P(`stringValue = `, p.Ident("unique", `Make`), `[string](`, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen)).Value()`)
+			p.P(`}`)
 		}
 		if oneof {
 			p.P(`m.`, fieldname, ` = &`, field.GoIdent, `{`, field.GoName, ": ", str, `}`)

include/github.com/planetscale/vtprotobuf/vtproto/ext.proto

@@ -9,4 +9,14 @@ option go_package = "github.com/planetscale/vtprotobuf/vtproto";
 
 extend google.protobuf.MessageOptions {
   optional bool mempool = 64101;
+}
+
+extend google.protobuf.FieldOptions {
+  optional Opts options = 64150;
+}
+
+// These options should be used during schema definition,
+// applying them to some of the fields in protobuf
+message Opts {
+  optional bool unique = 1;
 }

testproto/unique/unique.pb.go

@@ -0,0 +1,8 @@
+syntax = "proto3";
+option go_package = "testproto/unique";
+
+import "github.com/planetscale/vtprotobuf/vtproto/ext.proto";
+
+message UniqueFieldExtension {
+  string foo = 1 [(vtproto.options).unique = true];
+}

testproto/unique/unique.proto

@@ -0,0 +1,25 @@
+package unique
+
+import (
+	"testing"
+	"unsafe"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestUnmarshalSameMemory(t *testing.T) {
+	m := &UniqueFieldExtension{
+		Foo: "bar",
+	}
+
+	b, err := m.MarshalVTStrict()
+	require.NoError(t, err)
+
+	m2 := &UniqueFieldExtension{}
+	require.NoError(t, m2.UnmarshalVT(b))
+
+	m3 := &UniqueFieldExtension{}
+	require.NoError(t, m3.UnmarshalVT(b))
+
+	require.Equal(t, unsafe.StringData(m2.Foo), unsafe.StringData(m3.Foo))
+}