From 618376dbc03d703109c08d6302ade4c509c3cbb4 Mon Sep 17 00:00:00 2001
From: Josh Bleecher Snyder <josh@tailscale.com>
Date: Thu, 16 Sep 2021 15:48:40 -0700
Subject: [PATCH] util/codegen: add AssertStructUnchanged

Refactored out from cmd/cloner.

Signed-off-by: Josh Bleecher Snyder <josh@tailscale.com>
---
 cmd/cloner/cloner.go    | 19 +------------------
 util/codegen/codegen.go | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/cmd/cloner/cloner.go b/cmd/cloner/cloner.go
index 41d036eec..322118170 100644
--- a/cmd/cloner/cloner.go
+++ b/cmd/cloner/cloner.go
@@ -164,18 +164,6 @@ func gen(buf *bytes.Buffer, imports map[string]struct{}, name string, typ *types
 
 	switch t := typ.Underlying().(type) {
 	case *types.Struct:
-		// We generate two bits of code simultaneously while we walk the struct.
-		// One is the Clone method itself, which we write directly to buf.
-		// The other is a variable assignment that will fail if the struct
-		// changes without the Clone method getting regenerated.
-		// We write that to regenBuf, and then append it to buf at the end.
-		regenBuf := new(bytes.Buffer)
-		writeRegen := func(format string, args ...interface{}) {
-			fmt.Fprintf(regenBuf, format+"\n", args...)
-		}
-		writeRegen("// A compilation failure here means this code must be regenerated, with the command at the top of this file.")
-		writeRegen("var _%sNeedsRegeneration = %s(struct {", name, name)
-
 		name := typ.Obj().Name()
 		fmt.Fprintf(buf, "// Clone makes a deep copy of %s.\n", name)
 		fmt.Fprintf(buf, "// The result aliases no memory with the original.\n")
@@ -191,9 +179,6 @@ func gen(buf *bytes.Buffer, imports map[string]struct{}, name string, typ *types
 		for i := 0; i < t.NumFields(); i++ {
 			fname := t.Field(i).Name()
 			ft := t.Field(i).Type()
-
-			writeRegen("\t%s %s", fname, importedName(ft))
-
 			if !containsPointers(ft) {
 				continue
 			}
@@ -258,9 +243,7 @@ func gen(buf *bytes.Buffer, imports map[string]struct{}, name string, typ *types
 		writef("return dst")
 		fmt.Fprintf(buf, "}\n\n")
 
-		writeRegen("}{})\n")
-
-		buf.Write(regenBuf.Bytes())
+		buf.Write(codegen.AssertStructUnchanged(t, name, "", thisPkg, imports))
 	}
 }
 
diff --git a/util/codegen/codegen.go b/util/codegen/codegen.go
index b7fbbc073..95302440b 100644
--- a/util/codegen/codegen.go
+++ b/util/codegen/codegen.go
@@ -6,8 +6,10 @@
 package codegen
 
 import (
+	"bytes"
 	"fmt"
 	"go/format"
+	"go/types"
 	"os"
 )
 
@@ -38,3 +40,42 @@ func WriteFormatted(code []byte, path string) error {
 	}
 	return nil
 }
+
+// AssertStructUnchanged generates code that asserts at compile time that type t is unchanged.
+// tname is the named type corresponding to t.
+// ctx is a single-word context for this assertion, such as "Clone".
+// thisPkg is the package containing t.
+// 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, tname, ctx string, thisPkg *types.Package, imports map[string]struct{}) []byte {
+	buf := new(bytes.Buffer)
+	w := func(format string, args ...interface{}) {
+		fmt.Fprintf(buf, format+"\n", args...)
+	}
+	w("// A compilation failure here means this code must be regenerated, with the command at the top of this file.")
+	w("var _%s%sNeedsRegeneration = %s(struct {", tname, ctx, tname)
+
+	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{}{}
+		}
+		w("\t%s %s", fname, qname)
+	}
+
+	w("}{})\n")
+	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
+}