unmarshal: add tests for unique

Signed-off-by: Giedrius Statkevičius <giedrius.statkevicius@vinted.com>

features/unmarshal/unmarshal.go

@@ -158,10 +158,7 @@ func (p *unmarshal) declareMapField(varName string, nullable bool, field *protog
 }
 
 func (p *unmarshal) mapField(varName string, field *protogen.Field) {
-	var unique bool
-	if _, ok := proto.GetExtension(field.Desc.Options(), vtproto.E_Unique).(bool); ok {
-		unique = true
-	}
+	unique := proto.GetExtension(field.Desc.Options(), vtproto.E_Options).(*vtproto.Opts).GetUnique()
 
 	switch field.Desc.Kind() {
 	case protoreflect.DoubleKind:
@@ -211,7 +208,7 @@ func (p *unmarshal) mapField(varName string, field *protogen.Field) {
 			p.P(`if intStringLen`, varName, ` == 0 {`)
 			p.P(varName, ` = ""`)
 			p.P(`} else {`)
-			p.P(varName, ` = unique.Make[string](`, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen`, varName, `)).Value()`)
+			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, `])`)
@@ -289,17 +286,11 @@ func (p *unmarshal) noStarOrSliceType(field *protogen.Field) string {
 }
 
 func (p *unmarshal) fieldItem(field *protogen.Field, fieldname string, message *protogen.Message, proto3 bool) {
-	var unique bool
-
 	repeated := field.Desc.Cardinality() == protoreflect.Repeated
 	typ := p.noStarOrSliceType(field)
 	oneof := field.Oneof != nil && !field.Oneof.Desc.IsSynthetic()
 	nullable := field.Oneof != nil && field.Oneof.Desc.IsSynthetic()
 
-	if _, ok := proto.GetExtension(field.Desc.Options(), vtproto.E_Unique).(bool); ok {
-		unique = true
-	}
-
 	switch field.Desc.Kind() {
 	case protoreflect.DoubleKind:
 		p.P(`var v uint64`)
@@ -429,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)`)
@@ -454,7 +447,7 @@ func (p *unmarshal) fieldItem(field *protogen.Field, fieldname string, message *
 			str = "stringValue"
 			p.P(`var stringValue string`)
 			p.P(`if intStringLen > 0 {`)
-			p.P(`stringValue = unique.Make[string](`, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen)).Value()`)
+			p.P(`stringValue = `, p.Ident("unique", `Make`), `[string](`, p.Ident("unsafe", `String`), `(&dAtA[iNdEx], intStringLen)).Value()`)
 			p.P(`}`)
 		}
 		if oneof {

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

@@ -12,11 +12,11 @@ extend google.protobuf.MessageOptions {
 }
 
 extend google.protobuf.FieldOptions {
-  optional FieldOptions options = 64102;
+  optional Opts options = 999999;
 }
 
 // These options should be used during schema definition,
 // applying them to some of the fields in protobuf
-message FieldOptions {
+message Opts {
   optional bool unique = 1;
 }

testproto/unique/unique.proto

@@ -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_test.go

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