features/unmarshal: implement unmarshal_unique

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,11 @@ 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
+	}
+
 	switch field.Desc.Kind() {
 	case protoreflect.DoubleKind:
 		p.P(`var `, varName, `temp uint64`)
@@ -193,13 +200,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, ` = 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)
@@ -275,11 +289,17 @@ 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`)
@@ -423,12 +443,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 = 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 FieldOptions options = 64102;
+}
+
+// These options should be used during schema definition,
+// applying them to some of the fields in protobuf
+message FieldOptions {
+  optional bool unique = 1;
 }