Add zero-copy Unmarshal (#389)

klauspost · web-flow · commit 4dfdcb521d8c · 2025-03-10T17:19:21.000+01:00
Adds `zerocopy` tag that will allow byte slices to be unmarshaled with zero-copy, referencing the original slice.

AFAICT not possible for other types or when decoding.
diff --git a/_generated/def.go b/_generated/def.go
@@ -25,13 +25,21 @@ type Block [32]byte
 // tests edge-cases with
 // compiling size compilation.
 type X struct {
-	Values    [32]byte    // should compile to 32*msgp.ByteSize; encoded as Bin
-	ValuesPtr *[32]byte   // check (*)[:] deref
-	More      Block       // should be identical to the above
-	Others    [][32]int32 // should compile to len(x.Others)*32*msgp.Int32Size
-	Matrix    [][]int32   // should not optimize
-	ManyFixed []Fixed
-	WeirdTag  string `msg:"\x0b"`
+	Values         [32]byte    // should compile to 32*msgp.ByteSize; encoded as Bin
+	ValuesPtr      *[32]byte   // check (*)[:] deref
+	More           Block       // should be identical to the above
+	Others         [][32]int32 // should compile to len(x.Others)*32*msgp.Int32Size
+	Matrix         [][]int32   // should not optimize
+	ManyFixed      []Fixed
+	WeirdTag       string                       `msg:"\x0b"`
+	ZCBytes        []byte                       `msg:",zerocopy"`
+	ZCBytesAN      []byte                       `msg:",zerocopy,allownil"`
+	ZCBytesOE      []byte                       `msg:",zerocopy,omitempty"`
+	ZCBytesSlice   [][]byte                     `msg:",zerocopy"`
+	ZCBytesArr     [2][]byte                    `msg:",zerocopy"`
+	ZCBytesMap     map[string][]byte            `msg:",zerocopy"`
+	ZCBytesMapDeep map[string]map[string][]byte `msg:",zerocopy"`
+	CustomBytes    CustomBytes                  `msg:",zerocopy"`
 }
 
 // test fixed-size struct
@@ -73,6 +81,14 @@ type Object struct {
 	MapMap          map[string]map[string]string
 	MapStringEmpty  map[string]struct{}
 	MapStringEmpty2 map[string]EmptyStruct
+	ZCBytes         []byte                       `msg:",zerocopy"`
+	ZCBytesAN       []byte                       `msg:",zerocopy,allownil"`
+	ZCBytesOE       []byte                       `msg:",zerocopy,omitempty"`
+	ZCBytesSlice    [][]byte                     `msg:",zerocopy"`
+	ZCBytesArr      [2][]byte                    `msg:",zerocopy"`
+	ZCBytesMap      map[string][]byte            `msg:",zerocopy"`
+	ZCBytesMapDeep  map[string]map[string][]byte `msg:",zerocopy"`
+	CustomBytes     CustomBytes                  `msg:",zerocopy"`
 }
 
 //msgp:tuple TestBench
diff --git a/gen/elem.go b/gen/elem.go
@@ -576,6 +576,7 @@ type BaseElem struct {
 	ShimFromBase string    // shim from base type, or empty
 	Value        Primitive // Type of element
 	Convert      bool      // should we do an explicit conversion?
+	zerocopy     bool      // Allow zerocopy for byte slices in unmarshal.
 	mustinline   bool      // must inline; not printable
 	needsref     bool      // needs reference for shim
 	allowNil     *bool     // Override from parent.
diff --git a/gen/unmarshal.go b/gen/unmarshal.go
@@ -89,14 +89,37 @@ func (u *unmarshalGen) tuple(s *Struct) {
 			u.p.printf("\nif msgp.IsNil(bts) {\nbts = bts[1:]\n%s = nil\n} else {", fieldElem.Varname())
 		}
 		SetIsAllowNil(fieldElem, anField)
+		if s.Fields[i].HasTagPart("zerocopy") {
+			setRecursiveZC(fieldElem, true)
+		}
 		next(u, fieldElem)
+		if s.Fields[i].HasTagPart("zerocopy") {
+			setRecursiveZC(fieldElem, false)
+		}
+
 		u.ctx.Pop()
 		if anField {
 			u.p.printf("\n}")
 		}
 	}
 }
 
+// setRecursiveZC will alloc zerocopy for byte fields that are present.
+func setRecursiveZC(e Elem, enable bool) {
+	if base, ok := e.(*BaseElem); ok {
+		base.zerocopy = enable
+	}
+	if el, ok := e.(*Slice); ok {
+		setRecursiveZC(el.Els, enable)
+	}
+	if el, ok := e.(*Array); ok {
+		setRecursiveZC(el.Els, enable)
+	}
+	if el, ok := e.(*Map); ok {
+		setRecursiveZC(el.Value, enable)
+	}
+}
+
 func (u *unmarshalGen) mapstruct(s *Struct) {
 	u.needsField()
 	sz := randIdent()
@@ -136,7 +159,13 @@ func (u *unmarshalGen) mapstruct(s *Struct) {
 			u.p.printf("\nif msgp.IsNil(bts) {\nbts = bts[1:]\n%s = nil\n} else {", fieldElem.Varname())
 		}
 		SetIsAllowNil(fieldElem, anField)
+		if s.Fields[i].HasTagPart("zerocopy") {
+			setRecursiveZC(fieldElem, true)
+		}
 		next(u, fieldElem)
+		if s.Fields[i].HasTagPart("zerocopy") {
+			setRecursiveZC(fieldElem, false)
+		}
 		u.ctx.Pop()
 		if oeCount > 0 && (s.Fields[i].HasTagPart("omitempty") || s.Fields[i].HasTagPart("omitzero")) {
 			u.p.printf("\n%s", bm.setStmt(len(oeEmittedIdx)))
@@ -188,7 +217,11 @@ func (u *unmarshalGen) gBase(b *BaseElem) {
 
 	switch b.Value {
 	case Bytes:
-		u.p.printf("\n%s, bts, err = msgp.ReadBytesBytes(bts, %s)", refname, lowered)
+		if b.zerocopy {
+			u.p.printf("\n%s, bts, err = msgp.ReadBytesZC(bts)", refname)
+		} else {
+			u.p.printf("\n%s, bts, err = msgp.ReadBytesBytes(bts, %s)", refname, lowered)
+		}
 	case Ext:
 		u.p.printf("\nbts, err = msgp.ReadExtensionBytes(bts, %s)", lowered)
 	case IDENT:
