Fix prefix dropper in Go codegen (#56041)

go.mod

@@ -254,6 +254,7 @@ require (
 	github.com/google/go-github/v45 v45.2.0
 	github.com/grafana/dskit v0.0.0-20211011144203-3a88ec0b675f
 	github.com/jmoiron/sqlx v1.3.5
+	github.com/matryer/is v1.4.0
 	github.com/urfave/cli v1.22.5
 	go.etcd.io/etcd/api/v3 v3.5.4
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.31.0

pkg/codegen/astmanip_test.go

@@ -0,0 +1,263 @@
+package codegen
+
+import (
+	"bytes"
+	"go/format"
+	"go/parser"
+	"go/token"
+	"testing"
+
+	"github.com/matryer/is"
+	"golang.org/x/tools/go/ast/astutil"
+)
+
+func TestPrefixDropper(t *testing.T) {
+	tt := map[string]struct {
+		in, out string
+		skip    bool
+	}{
+		"basic": {
+			in: `package foo
+
+type Foo struct {
+	Id int64
+	Ref FooThing
+}
+
+type FooThing struct {
+	Id int64
+}`,
+			out: `package foo
+
+type Model struct {
+	Id  int64
+	Ref Thing
+}
+
+type Thing struct {
+	Id int64
+}
+`,
+		},
+		"pointer": {
+			in: `package foo
+
+type Foo struct {
+	Id int64
+	Ref *FooThing
+}
+
+type FooThing struct {
+	Id int64
+}`,
+			out: `package foo
+
+type Model struct {
+	Id  int64
+	Ref *Thing
+}
+
+type Thing struct {
+	Id int64
+}
+`,
+		},
+		"sliceref": {
+			in: `package foo
+
+type Foo struct {
+	Id int64
+	Ref []FooThing
+	PRef []*FooThing
+	SPRef *[]FooThing
+}
+
+type FooThing struct {
+	Id int64
+}`,
+			out: `package foo
+
+type Model struct {
+	Id    int64
+	Ref   []Thing
+	PRef  []*Thing
+	SPRef *[]Thing
+}
+
+type Thing struct {
+	Id int64
+}
+`,
+		},
+		"mapref": {
+			in: `package foo
+
+type Foo struct {
+	Id int64
+	KeyRef map[FooThing]string
+	ValRef map[string]FooThing
+	BothRef map[FooThing]FooThing
+}
+
+type FooThing struct {
+	Id int64
+}`,
+			out: `package foo
+
+type Model struct {
+	Id      int64
+	KeyRef  map[Thing]string
+	ValRef  map[string]Thing
+	BothRef map[Thing]Thing
+}
+
+type Thing struct {
+	Id int64
+}
+`,
+		},
+		"pmapref": {
+			in: `package foo
+
+type Foo struct {
+	Id int64
+	KeyRef map[*FooThing]string
+	ValRef map[string]*FooThing
+	BothRef map[*FooThing]*FooThing
+	PKeyRef *map[*FooThing]string
+}
+
+type FooThing struct {
+	Id int64
+}`,
+			out: `package foo
+
+type Model struct {
+	Id      int64
+	KeyRef  map[*Thing]string
+	ValRef  map[string]*Thing
+	BothRef map[*Thing]*Thing
+	PKeyRef *map[*Thing]string
+}
+
+type Thing struct {
+	Id int64
+}
+`,
+		},
+		"ignore-fieldname": {
+			in: `package foo
+
+type Foo struct {
+	Id int64
+	FooRef []string
+}`,
+			out: `package foo
+
+type Model struct {
+	Id     int64
+	FooRef []string
+}
+`,
+		},
+		"const": {
+			in: `package foo
+
+const one FooThing = "boop"
+
+const (
+	two   FooThing = "boop"
+	three FooThing = "boop"
+)
+
+type FooThing string
+`,
+			out: `package foo
+
+const one Thing = "boop"
+
+const (
+	two   Thing = "boop"
+	three Thing = "boop"
+)
+
+type Thing string
+`,
+		},
+		"var": {
+			in: `package foo
+
+var one FooThing = "boop"
+
+var (
+	two   FooThing = "boop"
+	three FooThing = "boop"
+)
+
+type FooThing string
+`,
+			out: `package foo
+
+var one Thing = "boop"
+
+var (
+	two   Thing = "boop"
+	three Thing = "boop"
+)
+
+type Thing string
+`,
+		},
+		"varp": {
+			in: `package foo
+
+var one *FooThing = "boop"
+
+var (
+	two   []FooThing = []FooThing{"boop"}
+	three map[FooThing]string = map[FooThing]string{ "beep": "boop" }
+)
+
+type FooThing string
+`,
+			out: `package foo
+
+var one *Thing = "boop"
+
+var (
+	two   []Thing = []Thing{"boop"}
+	three map[Thing]string = map[Thing]string{ "beep": "boop" }
+)
+
+type Thing string
+`,
+			// Skip this one for now - there's currently no codegen that constructs instances
+			// of objects, only types, so we shouldn't encounter this case.
+			skip: true,
+		},
+	}
+
+	for name, it := range tt {
+		item := it
+		t.Run(name, func(t *testing.T) {
+			if item.skip {
+				t.Skip()
+			}
+			is := is.New(t)
+			fset := token.NewFileSet()
+			inf, err := parser.ParseFile(fset, "input.go", item.in, parser.ParseComments)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			drop := makePrefixDropper("Foo", "Model")
+			astutil.Apply(inf, drop, nil)
+			buf := new(bytes.Buffer)
+			err = format.Node(buf, fset, inf)
+			if err != nil {
+				t.Fatal(err)
+			}
+			is.Equal(item.out, buf.String())
+		})
+	}
+}

pkg/codegen/coremodel.go

@@ -21,6 +21,7 @@ import (
 	"github.com/grafana/grafana/pkg/cuectx"
 	"github.com/grafana/thema"
 	"github.com/grafana/thema/encoding/openapi"
+	"golang.org/x/tools/go/ast/astutil"
 )
 
 // CoremodelDeclaration contains the results of statically analyzing a Grafana
@@ -279,30 +280,77 @@ type prefixDropper struct {
 	rxpsuff *regexp.Regexp
 }
 
-func makePrefixDropper(str, base string) prefixDropper {
-	return prefixDropper{
+func makePrefixDropper(str, base string) astutil.ApplyFunc {
+	return (&prefixDropper{
 		str:     str,
 		base:    base,
 		rxpsuff: regexp.MustCompile(fmt.Sprintf(`%s([a-zA-Z_]*)`, str)),
 		rxp:     regexp.MustCompile(fmt.Sprintf(`%s([\s.,;-])`, str)),
-	}
+	}).applyfunc
 }
 
-func (d prefixDropper) Visit(n ast.Node) ast.Visitor {
+func depoint(e ast.Expr) ast.Expr {
+	if star, is := e.(*ast.StarExpr); is {
+		return star.X
+	}
+	return e
+}
+
+func (d prefixDropper) applyfunc(c *astutil.Cursor) bool {
+	n := c.Node()
+
+	// fmt.Printf("%T %s\n", c.Node(), ast.Print(nil, c.Node()))
 	switch x := n.(type) {
-	case *ast.Ident:
-		if x.Name != d.str {
-			x.Name = strings.TrimPrefix(x.Name, d.str)
-		} else {
-			x.Name = d.base
+	case *ast.ValueSpec:
+		// fmt.Printf("%T %s\n", c.Node(), ast.Print(nil, c.Node()))
+		d.handleExpr(x.Type)
+		for _, id := range x.Names {
+			d.do(id)
 		}
+	case *ast.TypeSpec:
+		// Always do typespecs
+		d.do(x.Name)
+	case *ast.Field:
+		// Don't rename struct fields. We just want to rename type declarations, and
+		// field value specifications that reference those types.
+		d.handleExpr(x.Type)
+		// return false
+
 	case *ast.CommentGroup:
 		for _, c := range x.List {
 			c.Text = d.rxp.ReplaceAllString(c.Text, d.base+"$1")
 			c.Text = d.rxpsuff.ReplaceAllString(c.Text, "$1")
 		}
 	}
-	return d
+	return true
+}
+
+func (d prefixDropper) handleExpr(e ast.Expr) {
+	// Deref a StarExpr, if there is one
+	expr := depoint(e)
+	switch x := expr.(type) {
+	case *ast.Ident:
+		d.do(x)
+	case *ast.ArrayType:
+		if id, is := depoint(x.Elt).(*ast.Ident); is {
+			d.do(id)
+		}
+	case *ast.MapType:
+		if id, is := depoint(x.Key).(*ast.Ident); is {
+			d.do(id)
+		}
+		if id, is := depoint(x.Value).(*ast.Ident); is {
+			d.do(id)
+		}
+	}
+}
+
+func (d prefixDropper) do(n *ast.Ident) {
+	if n.Name != d.str {
+		n.Name = strings.TrimPrefix(n.Name, d.str)
+	} else {
+		n.Name = d.base
+	}
 }
 
 // GenerateCoremodelRegistry produces Go files that define a registry with

pkg/codegen/util_go.go

@@ -3,7 +3,6 @@ package codegen
 import (
 	"bytes"
 	"fmt"
-	"go/ast"
 	"go/format"
 	"go/parser"
 	"go/token"
@@ -11,12 +10,13 @@ import (
 	"path/filepath"
 	"strings"
 
+	"golang.org/x/tools/go/ast/astutil"
 	"golang.org/x/tools/imports"
 )
 
 type genGoFile struct {
 	path   string
-	walker ast.Visitor
+	walker astutil.ApplyFunc
 	in     []byte
 }
 
@@ -30,7 +30,7 @@ func postprocessGoFile(cfg genGoFile) ([]byte, error) {
 	}
 
 	if cfg.walker != nil {
-		ast.Walk(cfg.walker, gf)
+		astutil.Apply(gf, cfg.walker, nil)
 
 		err = format.Node(buf, fset, gf)
 		if err != nil {