Allow maxtotalbytes to be specified on top types

parse/directives.go

@@ -22,11 +22,12 @@ type passDirective func(gen.Method, []string, *gen.Printer) error
 // to add a directive, define a func([]string, *FileSet) error
 // and then add it to this list.
 var directives = map[string]directive{
-	"shim":       applyShim,
-	"ignore":     ignore,
-	"tuple":      astuple,
-	"sort":       sortintf,
-	"allocbound": allocbound,
+	"shim":          applyShim,
+	"ignore":        ignore,
+	"tuple":         astuple,
+	"sort":          sortintf,
+	"allocbound":    allocbound,
+	"maxtotalbytes": maxtotalbytes,
 	// _postunmarshalcheck is used to add callbacks to the end of un-marshalling that are tied to a specific Element.
 	_postunmarshalcheck: postunmarshalcheck,
 }
@@ -192,3 +193,27 @@ func allocbound(text []string, f *FileSet) error {
 	}
 	return nil
 }
+
+// maxtotalbytes is the type-level counterpart of the maxtotalbytes= field tag,
+// for named slice types that have no field to carry the tag. As with the tag,
+// Bound replaces the generated MaxSize for the slice's contents, so it must
+// include each element's msgp prefix. Directives are split on spaces, so the
+// rest of the line is rejoined to allow Bound to be an expression like
+// "(N*msgp.BytesPrefixSize) + M".
+//
+//msgp:maxtotalbytes {Type} {Bound}
+func maxtotalbytes(text []string, f *FileSet) error {
+	if len(text) < 3 {
+		return fmt.Errorf("maxtotalbytes directive should have 2 arguments; found %d", len(text)-1)
+	}
+	maxTotalBytesType := strings.TrimSpace(text[1])
+	maxTotalBytes := strings.TrimSpace(strings.Join(text[2:], " "))
+	t, ok := f.Identities[maxTotalBytesType]
+	if !ok {
+		warnf("maxtotalbytes: cannot find type %s\n", maxTotalBytesType)
+	} else {
+		t.SetMaxTotalBytes(maxTotalBytes)
+		infof("maxtotalbytes(%s): setting to %s\n", maxTotalBytesType, maxTotalBytes)
+	}
+	return nil
+}

tests/maxtotalbytes.go

@@ -0,0 +1,33 @@
+// This file exercises the maxtotalbytes directive, which gives a named slice
+// type the same tight MaxSize bound that the maxtotalbytes= field tag gives a
+// struct field.
+package tests
+
+//go:generate msgp
+
+const (
+	argCount   = 8
+	argMaxSize = 100
+	// argsTotal is much smaller than argCount*argMaxSize, so a MaxSize that
+	// uses it is visibly tighter than the product of the allocbounds.
+	argsTotal = 200
+)
+
+// BoundedArgs has only allocbounds, so its MaxSize is the product of them.
+//
+//msgp:allocbound BoundedArgs argCount,argMaxSize
+type BoundedArgs [][]byte
+
+// TotalBoundedArgs adds a maxtotalbytes directive whose bound is an
+// expression containing spaces, to check that the whole expression survives.
+//
+//msgp:allocbound TotalBoundedArgs argCount,argMaxSize
+//msgp:maxtotalbytes TotalBoundedArgs (argCount*msgp.BytesPrefixSize) + argsTotal
+type TotalBoundedArgs [][]byte
+
+// ArgsHolder holds TotalBoundedArgs, and should use its MaxSize.
+type ArgsHolder struct {
+	_struct struct{} `codec:",omitempty"`
+
+	Args TotalBoundedArgs `codec:"a"`
+}