From 59a3e18bb2ef3ed3d736fe2baf4135d80ab6c0f2 Mon Sep 17 00:00:00 2001 From: sam boyer Date: Thu, 29 Sep 2022 12:01:02 -0400 Subject: [PATCH] Fix prefix dropper in Go codegen (#56041) --- go.mod | 1 + pkg/codegen/astmanip_test.go | 263 +++++++++++++++++++++++++++++++++++ pkg/codegen/coremodel.go | 68 +++++++-- pkg/codegen/util_go.go | 6 +- 4 files changed, 325 insertions(+), 13 deletions(-) create mode 100644 pkg/codegen/astmanip_test.go diff --git a/go.mod b/go.mod index 1644feccc83..0551dd7a91e 100644 --- a/go.mod +++ b/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 diff --git a/pkg/codegen/astmanip_test.go b/pkg/codegen/astmanip_test.go new file mode 100644 index 00000000000..0f44a112f39 --- /dev/null +++ b/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()) + }) + } +} diff --git a/pkg/codegen/coremodel.go b/pkg/codegen/coremodel.go index 07676ef1a24..d70ae77d33e 100644 --- a/pkg/codegen/coremodel.go +++ b/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 diff --git a/pkg/codegen/util_go.go b/pkg/codegen/util_go.go index 525b764298c..9033e29b968 100644 --- a/pkg/codegen/util_go.go +++ b/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 {