cmd/cloner,util/codegen: refactor cloner internals to allow reuse

Also run go generate again for Copyright updates.

Signed-off-by: Maisem Ali <maisem@tailscale.com>

cmd/cloner/cloner.go

@@ -22,13 +22,11 @@ import (
 	"os"
 	"strings"
 
-	"golang.org/x/tools/go/packages"
 	"tailscale.com/util/codegen"
 )
 
 var (
 	flagTypes     = flag.String("type", "", "comma-separated list of types; required")
-	flagOutput    = flag.String("output", "", "output file; required")
 	flagBuildTags = flag.String("tags", "", "compiler build tags to apply")
 	flagCloneFunc = flag.Bool("clonefunc", false, "add a top-level Clone func")
 )
@@ -43,30 +41,18 @@ func main() {
 	}
 	typeNames := strings.Split(*flagTypes, ",")
 
-	cfg := &packages.Config{
+	pkg, namedTypes, err := codegen.LoadTypes(*flagBuildTags, ".")
-		Mode:  packages.NeedTypes | packages.NeedTypesInfo | packages.NeedSyntax | packages.NeedName,
-		Tests: false,
-	}
-	if *flagBuildTags != "" {
-		cfg.BuildFlags = []string{"-tags=" + *flagBuildTags}
-	}
-	pkgs, err := packages.Load(cfg, ".")
 	if err != nil {
 		log.Fatal(err)
 	}
-	if len(pkgs) != 1 {
+	it := codegen.NewImportTracker(pkg.Types)
-		log.Fatalf("wrong number of packages: %d", len(pkgs))
-	}
-	pkg := pkgs[0]
 	buf := new(bytes.Buffer)
-	imports := make(map[string]struct{})
-	namedTypes := codegen.NamedTypes(pkg)
 	for _, typeName := range typeNames {
 		typ, ok := namedTypes[typeName]
 		if !ok {
 			log.Fatalf("could not find type %s", typeName)
 		}
-		gen(buf, imports, typ, pkg.Types)
+		gen(buf, it, typ)
 	}
 
 	w := func(format string, args ...any) {
@@ -93,62 +79,13 @@ func main() {
 		w("	return false")
 		w("}")
 	}
-
+	cloneOutput := pkg.Name + "_clone.go"
-	contents := new(bytes.Buffer)
+	if err := codegen.WritePackageFile("tailscale.com/cmd/cloner", pkg, cloneOutput, it, buf); err != nil {
-	var flagArgs []string
-	if *flagTypes != "" {
-		flagArgs = append(flagArgs, "-type="+*flagTypes)
-	}
-	if *flagOutput != "" {
-		flagArgs = append(flagArgs, "-output="+*flagOutput)
-	}
-	if *flagBuildTags != "" {
-		flagArgs = append(flagArgs, "-tags="+*flagBuildTags)
-	}
-	if *flagCloneFunc {
-		flagArgs = append(flagArgs, "-clonefunc")
-	}
-	fmt.Fprintf(contents, header, strings.Join(flagArgs, " "), pkg.Name)
-	fmt.Fprintf(contents, "import (\n")
-	for s := range imports {
-		fmt.Fprintf(contents, "\t%q\n", s)
-	}
-	fmt.Fprintf(contents, ")\n\n")
-	contents.Write(buf.Bytes())
-
-	output := *flagOutput
-	if output == "" {
-		flag.Usage()
-		os.Exit(2)
-	}
-	if err := codegen.WriteFormatted(contents.Bytes(), output); err != nil {
 		log.Fatal(err)
 	}
 }
 
-const header = `// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+func gen(buf *bytes.Buffer, it *codegen.ImportTracker, typ *types.Named) {
-// Use of this source code is governed by a BSD-style
-// license that can be found in the LICENSE file.
-
-// Code generated by tailscale.com/cmd/cloner; DO NOT EDIT.
-//` + `go:generate` + ` go run tailscale.com/cmd/cloner %s
-
-package %s
-
-`
-
-func gen(buf *bytes.Buffer, imports map[string]struct{}, typ *types.Named, thisPkg *types.Package) {
-	pkgQual := func(pkg *types.Package) string {
-		if thisPkg == pkg {
-			return ""
-		}
-		imports[pkg.Path()] = struct{}{}
-		return pkg.Name()
-	}
-	importedName := func(t types.Type) string {
-		return types.TypeString(t, pkgQual)
-	}
-
 	t, ok := typ.Underlying().(*types.Struct)
 	if !ok {
 		return
@@ -169,11 +106,11 @@ func gen(buf *bytes.Buffer, imports map[string]struct{}, typ *types.Named, thisP
 	for i := 0; i < t.NumFields(); i++ {
 		fname := t.Field(i).Name()
 		ft := t.Field(i).Type()
-		if !codegen.ContainsPointers(ft) {
+		if !codegen.ContainsPointers(ft) || codegen.HasNoClone(t.Tag(i)) {
 			continue
 		}
 		if named, _ := ft.(*types.Named); named != nil {
-			if isViewType(ft) {
+			if codegen.IsViewType(ft) {
 				writef("dst.%s = src.%s", fname, fname)
 				continue
 			}
@@ -185,7 +122,7 @@ func gen(buf *bytes.Buffer, imports map[string]struct{}, typ *types.Named, thisP
 		switch ft := ft.Underlying().(type) {
 		case *types.Slice:
 			if codegen.ContainsPointers(ft.Elem()) {
-				n := importedName(ft.Elem())
+				n := it.QualifiedName(ft.Elem())
 				writef("dst.%s = make([]%s, len(src.%s))", fname, n, fname)
 				writef("for i := range dst.%s {", fname)
 				if _, isPtr := ft.Elem().(*types.Pointer); isPtr {
@@ -202,7 +139,7 @@ func gen(buf *bytes.Buffer, imports map[string]struct{}, typ *types.Named, thisP
 				writef("dst.%s = src.%s.Clone()", fname, fname)
 				continue
 			}
-			n := importedName(ft.Elem())
+			n := it.QualifiedName(ft.Elem())
 			writef("if dst.%s != nil {", fname)
 			writef("\tdst.%s = new(%s)", fname, n)
 			writef("\t*dst.%s = *src.%s", fname, fname)
@@ -212,9 +149,9 @@ func gen(buf *bytes.Buffer, imports map[string]struct{}, typ *types.Named, thisP
 			writef("}")
 		case *types.Map:
 			writef("if dst.%s != nil {", fname)
-			writef("\tdst.%s = map[%s]%s{}", fname, importedName(ft.Key()), importedName(ft.Elem()))
+			writef("\tdst.%s = map[%s]%s{}", fname, it.QualifiedName(ft.Key()), it.QualifiedName(ft.Elem()))
 			if sliceType, isSlice := ft.Elem().(*types.Slice); isSlice {
-				n := importedName(sliceType.Elem())
+				n := it.QualifiedName(sliceType.Elem())
 				writef("\tfor k := range src.%s {", fname)
 				// use zero-length slice instead of nil to ensure
 				// the key is always copied.
@@ -237,20 +174,10 @@ func gen(buf *bytes.Buffer, imports map[string]struct{}, typ *types.Named, thisP
 	writef("return dst")
 	fmt.Fprintf(buf, "}\n\n")
 
-	buf.Write(codegen.AssertStructUnchanged(t, thisPkg, name, "Clone", imports))
+	buf.Write(codegen.AssertStructUnchanged(t, name, "Clone", it))
-}
-
-func isViewType(typ types.Type) bool {
-	t, ok := typ.Underlying().(*types.Struct)
-	if !ok {
-		return false
-	}
-	if t.NumFields() != 1 {
-		return false
-	}
-	return t.Field(0).Name() == "ж"
 }
 
+// hasBasicUnderlying reports true when typ.Underlying() is a slice or a map.
 func hasBasicUnderlying(typ types.Type) bool {
 	switch typ.Underlying().(type) {
 	case *types.Slice, *types.Map:

ipn/ipn_clone.go

@@ -1,9 +1,8 @@
-// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
 // Code generated by tailscale.com/cmd/cloner; DO NOT EDIT.
-//go:generate go run tailscale.com/cmd/cloner -type=Prefs -output=prefs_clone.go
 
 package ipn
 

ipn/prefs.go

@@ -27,7 +27,7 @@ import (
 	"tailscale.com/util/dnsname"
 )
 
-//go:generate go run tailscale.com/cmd/cloner -type=Prefs -output=prefs_clone.go
+//go:generate go run tailscale.com/cmd/cloner -type=Prefs
 
 // DefaultControlURL is the URL base of the control plane
 // ("coordination server") for use when no explicit one is configured.

tailcfg/tailcfg.go

@@ -4,7 +4,7 @@
 
 package tailcfg
 
-//go:generate go run tailscale.com/cmd/cloner --type=User,Node,Hostinfo,NetInfo,Login,DNSConfig,RegisterResponse,DERPRegion,DERPMap,DERPNode --clonefunc=true --output=tailcfg_clone.go
+//go:generate go run tailscale.com/cmd/cloner --type=User,Node,Hostinfo,NetInfo,Login,DNSConfig,RegisterResponse,DERPRegion,DERPMap,DERPNode --clonefunc
 
 import (
 	"encoding/hex"

tailcfg/tailcfg_clone.go

@@ -1,9 +1,8 @@
-// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-// Code generated by tailscale.com/cmd/cloner; DO NOT EDIT.
+// Code generated by tailscale.com/util/codegen (cloner); DO NOT EDIT.
-//go:generate go run tailscale.com/cmd/cloner -type=User,Node,Hostinfo,NetInfo,Login,DNSConfig,RegisterResponse,DERPRegion,DERPMap,DERPNode -output=tailcfg_clone.go -clonefunc
 
 package tailcfg
 

types/dnstype/dnstype.go

@@ -5,7 +5,7 @@
 // Package dnstype defines types for working with DNS.
 package dnstype
 
-//go:generate go run tailscale.com/cmd/cloner --type=Resolver --clonefunc=true --output=dnstype_clone.go
+//go:generate go run tailscale.com/cmd/cloner --type=Resolver --clonefunc=true
 
 import "inet.af/netaddr"
 

types/dnstype/dnstype_clone.go

@@ -1,9 +1,8 @@
-// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
 // Code generated by tailscale.com/cmd/cloner; DO NOT EDIT.
-//go:generate go run tailscale.com/cmd/cloner -type=Resolver -output=dnstype_clone.go -clonefunc
 
 package dnstype
 

types/persist/persist.go

@@ -12,7 +12,7 @@ import (
 	"tailscale.com/types/structs"
 )
 
-//go:generate go run tailscale.com/cmd/cloner -type=Persist -output=persist_clone.go
+//go:generate go run tailscale.com/cmd/cloner -type=Persist
 
 // Persist is the JSON type stored on disk on nodes to remember their
 // settings between runs.

types/persist/persist_clone.go

@@ -1,9 +1,8 @@
-// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
 // Code generated by tailscale.com/cmd/cloner; DO NOT EDIT.
-//go:generate go run tailscale.com/cmd/cloner -type=Persist -output=persist_clone.go
 
 package persist
 

util/codegen/codegen.go

@@ -11,13 +11,116 @@ import (
 	"go/ast"
 	"go/token"
 	"go/types"
+	"io"
 	"os"
+	"reflect"
+	"strings"
+	"time"
 
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/imports"
+	"tailscale.com/util/mak"
 )
 
-// WriteFormatted writes code to path.
+// LoadTypes returns all named types in pkgName, keyed by their type name.
+func LoadTypes(buildTags string, pkgName string) (*packages.Package, map[string]*types.Named, error) {
+	cfg := &packages.Config{
+		Mode:  packages.NeedTypes | packages.NeedTypesInfo | packages.NeedSyntax | packages.NeedName,
+		Tests: false,
+	}
+	if buildTags != "" {
+		cfg.BuildFlags = []string{"-tags=" + buildTags}
+	}
+
+	pkgs, err := packages.Load(cfg, pkgName)
+	if err != nil {
+		return nil, nil, err
+	}
+	if len(pkgs) != 1 {
+		return nil, nil, fmt.Errorf("wrong number of packages: %d", len(pkgs))
+	}
+	pkg := pkgs[0]
+	return pkg, namedTypes(pkg), nil
+}
+
+// HasNoClone reports whether the provided tag has `codegen:noclone`.
+func HasNoClone(structTag string) bool {
+	val := reflect.StructTag(structTag).Get("codegen")
+	for _, v := range strings.Split(val, ",") {
+		if v == "noclone" {
+			return true
+		}
+	}
+	return false
+}
+
+const header = `// Copyright (c) %d Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Code generated by %v; DO NOT EDIT.
+
+package %s
+`
+
+func NewImportTracker(thisPkg *types.Package) *ImportTracker {
+	return &ImportTracker{
+		thisPkg: thisPkg,
+	}
+}
+
+// ImportTracker provides a mechanism to track and build import paths.
+type ImportTracker struct {
+	thisPkg  *types.Package
+	packages map[string]bool
+}
+
+func (it *ImportTracker) Import(pkg string) {
+	if pkg != "" && !it.packages[pkg] {
+		mak.Set(&it.packages, pkg, true)
+	}
+}
+
+func (it *ImportTracker) qualifier(pkg *types.Package) string {
+	if it.thisPkg == pkg {
+		return ""
+	}
+	it.Import(pkg.Path())
+	// TODO(maisem): handle conflicts?
+	return pkg.Name()
+}
+
+// QualifiedName returns the string representation of t in the package.
+func (it *ImportTracker) QualifiedName(t types.Type) string {
+	return types.TypeString(t, it.qualifier)
+}
+
+// Write prints all the tracked imports in a single import block to w.
+func (it *ImportTracker) Write(w io.Writer) {
+	fmt.Fprintf(w, "import (\n")
+	for s := range it.packages {
+		fmt.Fprintf(w, "\t%q\n", s)
+	}
+	fmt.Fprintf(w, ")\n\n")
+}
+
+func writeHeader(w io.Writer, tool, pkg string) {
+	fmt.Fprintf(w, header, time.Now().Year(), tool, pkg)
+}
+
+// WritePackageFile adds a file with the provided imports and contents to package.
+// The tool param is used to identify the tool that generated package file.
+func WritePackageFile(tool string, pkg *packages.Package, path string, it *ImportTracker, contents *bytes.Buffer) error {
+	buf := new(bytes.Buffer)
+	writeHeader(buf, tool, pkg.Name)
+	it.Write(buf)
+	if _, err := buf.Write(contents.Bytes()); err != nil {
+		return err
+	}
+	return writeFormatted(buf.Bytes(), path)
+}
+
+// writeFormatted writes code to path.
 // It runs gofmt on it before writing;
 // if gofmt fails, it writes code unchanged.
 // Errors can include I/O errors and gofmt errors.
@@ -28,7 +131,7 @@ import (
 // It is nicer to work with it in a file than a terminal.
 // It is also easier to interpret gofmt errors
 // with an editor providing file and line numbers.
-func WriteFormatted(code []byte, path string) error {
+func writeFormatted(code []byte, path string) error {
 	out, fmterr := imports.Process(path, code, &imports.Options{
 		Comments:   true,
 		TabIndent:  true,
@@ -50,8 +153,8 @@ func WriteFormatted(code []byte, path string) error {
 	return nil
 }
 
-// NamedTypes returns all named types in pkg, keyed by their type name.
+// namedTypes returns all named types in pkg, keyed by their type name.
-func NamedTypes(pkg *packages.Package) map[string]*types.Named {
+func namedTypes(pkg *packages.Package) map[string]*types.Named {
 	nt := make(map[string]*types.Named)
 	for _, file := range pkg.Syntax {
 		for _, d := range file.Decls {
@@ -64,7 +167,10 @@ func NamedTypes(pkg *packages.Package) map[string]*types.Named {
 				if !ok {
 					continue
 				}
-				typeNameObj := pkg.TypesInfo.Defs[spec.Name]
+				typeNameObj, ok := pkg.TypesInfo.Defs[spec.Name]
+				if !ok {
+					continue
+				}
 				typ, ok := typeNameObj.Type().(*types.Named)
 				if !ok {
 					continue
@@ -82,7 +188,7 @@ func NamedTypes(pkg *packages.Package) map[string]*types.Named {
 // ctx is a single-word context for this assertion, such as "Clone".
 // If non-nil, AssertStructUnchanged will add elements to imports
 // for each package path that the caller must import for the returned code to compile.
-func AssertStructUnchanged(t *types.Struct, thisPkg *types.Package, tname, ctx string, imports map[string]struct{}) []byte {
+func AssertStructUnchanged(t *types.Struct, tname, ctx string, it *ImportTracker) []byte {
 	buf := new(bytes.Buffer)
 	w := func(format string, args ...any) {
 		fmt.Fprintf(buf, format+"\n", args...)
@@ -93,10 +199,10 @@ func AssertStructUnchanged(t *types.Struct, thisPkg *types.Package, tname, ctx s
 	for i := 0; i < t.NumFields(); i++ {
 		fname := t.Field(i).Name()
 		ft := t.Field(i).Type()
-		qname, imppath := importedName(ft, thisPkg)
+		if IsInvalid(ft) {
-		if imppath != "" && imports != nil {
+			continue
-			imports[imppath] = struct{}{}
 		}
+		qname := it.QualifiedName(ft)
 		w("\t%s %s", fname, qname)
 	}
 
@@ -104,15 +210,11 @@ func AssertStructUnchanged(t *types.Struct, thisPkg *types.Package, tname, ctx s
 	return buf.Bytes()
 }
 
-func importedName(t types.Type, thisPkg *types.Package) (qualifiedName, importPkg string) {
+// IsInvalid reports whether the provided type is invalid. It is used to allow
-	qual := func(pkg *types.Package) string {
+// codegeneration to run even when the target files have build errors or are
-		if thisPkg == pkg {
+// missing views.
-			return ""
+func IsInvalid(t types.Type) bool {
-		}
+	return t.String() == "invalid type"
-		importPkg = pkg.Path()
-		return pkg.Name()
-	}
-	return types.TypeString(t, qual), importPkg
 }
 
 // ContainsPointers reports whether typ contains any pointers,
@@ -149,3 +251,15 @@ func ContainsPointers(typ types.Type) bool {
 	}
 	return false
 }
+
+// IsViewType reports whether the provided typ is a View.
+func IsViewType(typ types.Type) bool {
+	t, ok := typ.Underlying().(*types.Struct)
+	if !ok {
+		return false
+	}
+	if t.NumFields() != 1 {
+		return false
+	}
+	return t.Field(0).Name() == "ж"
+}

wgengine/filter/filter_clone.go

@@ -1,9 +1,8 @@
-// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
 // Code generated by tailscale.com/cmd/cloner; DO NOT EDIT.
-//go:generate go run tailscale.com/cmd/cloner -type=Match -output=match_clone.go
 
 package filter
 

wgengine/filter/match.go

@@ -13,7 +13,7 @@ import (
 	"tailscale.com/types/ipproto"
 )
 
-//go:generate go run tailscale.com/cmd/cloner --type=Match --output=match_clone.go
+//go:generate go run tailscale.com/cmd/cloner --type=Match
 
 // PortRange is a range of TCP and UDP ports.
 type PortRange struct {

wgengine/wgcfg/config.go

@@ -10,7 +10,7 @@ import (
 	"tailscale.com/types/key"
 )
 
-//go:generate go run tailscale.com/cmd/cloner -type=Config,Peer -output=clone.go
+//go:generate go run tailscale.com/cmd/cloner -type=Config,Peer
 
 // Config is a WireGuard configuration.
 // It only supports the set of things Tailscale uses.

wgengine/wgcfg/wgcfg_clone.go

@@ -1,9 +1,8 @@
-// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
 // Code generated by tailscale.com/cmd/cloner; DO NOT EDIT.
-//go:generate go run tailscale.com/cmd/cloner -type=Config,Peer -output=clone.go
 
 package wgcfg