1 // Copyright 2017 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
15 "golang.org/x/tools/go/ast/astutil"
18 var rewriteTests = [...]struct {
21 pre, post astutil.ApplyFunc
23 {name: "nop", orig: "package p\n", want: "package p\n"},
34 post: func(c *astutil.Cursor) bool {
35 if _, ok := c.Node().(*ast.ValueSpec); ok {
36 c.Replace(valspec("t", "T"))
43 {name: "set doc strings",
60 post: func(c *astutil.Cursor) bool {
61 if _, ok := c.Parent().(*ast.GenDecl); ok && c.Name() == "Doc" && c.Node() == nil {
62 c.Replace(&ast.CommentGroup{List: []*ast.Comment{{Text: "// a foo is a foo"}}})
68 {name: "insert names",
75 const a, b, c = 1, 2, 3
77 pre: func(c *astutil.Cursor) bool {
78 if _, ok := c.Parent().(*ast.ValueSpec); ok {
81 c.InsertAfter(ast.NewIdent("c"))
82 c.InsertAfter(ast.NewIdent("b"))
84 c.InsertAfter(&ast.BasicLit{Kind: token.INT, Value: "3"})
85 c.InsertAfter(&ast.BasicLit{Kind: token.INT, Value: "2"})
112 pre: func(c *astutil.Cursor) bool {
113 if _, ok := c.Node().(*ast.GenDecl); ok {
114 c.InsertBefore(vardecl("before1", "int"))
115 c.InsertAfter(vardecl("after1", "int"))
116 c.InsertAfter(vardecl("after2", "int"))
117 c.InsertBefore(vardecl("before2", "int"))
135 pre: func(c *astutil.Cursor) bool {
137 if d, ok := n.(*ast.GenDecl); ok && d.Specs[0].(*ast.ValueSpec).Names[0].Name == "x" {
144 {name: "insertafter-delete",
158 pre: func(c *astutil.Cursor) bool {
160 if d, ok := n.(*ast.GenDecl); ok && d.Specs[0].(*ast.ValueSpec).Names[0].Name == "x" {
161 c.InsertAfter(vardecl("x1", "int"))
168 {name: "delete-insertafter",
181 pre: func(c *astutil.Cursor) bool {
183 if d, ok := n.(*ast.GenDecl); ok && d.Specs[0].(*ast.ValueSpec).Names[0].Name == "x" {
185 // The cursor is now effectively atop the 'var y int' node.
186 c.InsertAfter(vardecl("x1", "int"))
193 func valspec(name, typ string) *ast.ValueSpec {
194 return &ast.ValueSpec{Names: []*ast.Ident{ast.NewIdent(name)},
195 Type: ast.NewIdent(typ),
199 func vardecl(name, typ string) *ast.GenDecl {
202 Specs: []ast.Spec{valspec(name, typ)},
206 func TestRewrite(t *testing.T) {
207 t.Run("*", func(t *testing.T) {
208 for _, test := range rewriteTests {
210 t.Run(test.name, func(t *testing.T) {
212 fset := token.NewFileSet()
213 f, err := parser.ParseFile(fset, test.name, test.orig, parser.ParseComments)
217 n := astutil.Apply(f, test.pre, test.post)
219 if err := format.Node(&buf, fset, n); err != nil {
223 if got != test.want {
224 t.Errorf("got:\n\n%s\nwant:\n\n%s\n", got, test.want)
233 func BenchmarkRewrite(b *testing.B) {
234 for _, test := range rewriteTests {
235 b.Run(test.name, func(b *testing.B) {
236 for i := 0; i < b.N; i++ {
238 fset := token.NewFileSet()
239 f, err := parser.ParseFile(fset, test.name, test.orig, parser.ParseComments)
244 sink = astutil.Apply(f, test.pre, test.post)