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>

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.
-func NamedTypes(pkg *packages.Package) map[string]*types.Named {
+// namedTypes returns all named types in pkg, keyed by their type name.
+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 imppath != "" && imports != nil {
-			imports[imppath] = struct{}{}
+		if IsInvalid(ft) {
+			continue
 		}
+		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) {
-	qual := func(pkg *types.Package) string {
-		if thisPkg == pkg {
-			return ""
-		}
-		importPkg = pkg.Path()
-		return pkg.Name()
-	}
-	return types.TypeString(t, qual), importPkg
+// IsInvalid reports whether the provided type is invalid. It is used to allow
+// codegeneration to run even when the target files have build errors or are
+// missing views.
+func IsInvalid(t types.Type) bool {
+	return t.String() == "invalid type"
 }
 
 // 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() == "ж"
+}