gopls/go/ast/astutil/rewrite

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.
4
5	package astutil_test
6
7	import (
8	"bytes"
9	"go/ast"
10	"go/format"
11	"go/parser"
12	"go/token"
13	"testing"
14
15	"golang.org/x/tools/go/ast/astutil"
16	"golang.org/x/tools/internal/typeparams"
17	)
18
19	type rewriteTest struct {
20	name string
21	orig, want string
22	pre, post astutil.ApplyFunc
23	}
24
25	var rewriteTests = []rewriteTest{
26	{name: "nop", orig: "package p\n", want: "package p\n"},
27
28	{name: "replace",
29	orig: `package p
30
31	var x int
32	`,
33	want: `package p
34
35	var t T
36	`,
37	post: func(c *astutil.Cursor) bool {
38	if _, ok := c.Node().(*ast.ValueSpec); ok {
39	c.Replace(valspec("t", "T"))
40	return false
41	}
42	return true
43	},
44	},
45
46	{name: "set doc strings",
47	orig: `package p
48
49	const z = 0
50
51	type T struct{}
52
53	var x int
54	`,
55	want: `package p
56	// a foo is a foo
57	const z = 0
58	// a foo is a foo
59	type T struct{}
60	// a foo is a foo
61	var x int
62	`,
63	post: func(c *astutil.Cursor) bool {
64	if _, ok := c.Parent().(*ast.GenDecl); ok && c.Name() == "Doc" && c.Node() == nil {
65	c.Replace(&ast.CommentGroup{List: []*ast.Comment{{Text: "// a foo is a foo"}}})
66	}
67	return true
68	},
69	},
70
71	{name: "insert names",
72	orig: `package p
73
74	const a = 1
75	`,
76	want: `package p
77
78	const a, b, c = 1, 2, 3
79	`,
80	pre: func(c *astutil.Cursor) bool {
81	if _, ok := c.Parent().(*ast.ValueSpec); ok {
82	switch c.Name() {
83	case "Names":
84	c.InsertAfter(ast.NewIdent("c"))
85	c.InsertAfter(ast.NewIdent("b"))
86	case "Values":
87	c.InsertAfter(&ast.BasicLit{Kind: token.INT, Value: "3"})
88	c.InsertAfter(&ast.BasicLit{Kind: token.INT, Value: "2"})
89	}
90	}
91	return true
92	},
93	},
94
95	{name: "insert",
96	orig: `package p
97
98	var (
99	x int
100	y int
101	)
102	`,
103	want: `package p
104
105	var before1 int
106	var before2 int
107
108	var (
109	x int
110	y int
111	)
112	var after2 int
113	var after1 int
114	`,
115	pre: func(c *astutil.Cursor) bool {
116	if _, ok := c.Node().(*ast.GenDecl); ok {
117	c.InsertBefore(vardecl("before1", "int"))
118	c.InsertAfter(vardecl("after1", "int"))
119	c.InsertAfter(vardecl("after2", "int"))
120	c.InsertBefore(vardecl("before2", "int"))
121	}
122	return true
123	},
124	},
125
126	{name: "delete",
127	orig: `package p
128
129	var x int
130	var y int
131	var z int
132	`,
133	want: `package p
134
135	var y int
136	var z int
137	`,
138	pre: func(c *astutil.Cursor) bool {
139	n := c.Node()
140	if d, ok := n.(ast.GenDecl); ok && d.Specs[0].(ast.ValueSpec).Names[0].Name == "x" {
141	c.Delete()
142	}
143	return true
144	},
145	},
146
147	{name: "insertafter-delete",
148	orig: `package p
149
150	var x int
151	var y int
152	var z int
153	`,
154	want: `package p
155
156	var x1 int
157
158	var y int
159	var z int
160	`,
161	pre: func(c *astutil.Cursor) bool {
162	n := c.Node()
163	if d, ok := n.(ast.GenDecl); ok && d.Specs[0].(ast.ValueSpec).Names[0].Name == "x" {
164	c.InsertAfter(vardecl("x1", "int"))
165	c.Delete()
166	}
167	return true
168	},
169	},
170
171	{name: "delete-insertafter",
172	orig: `package p
173
174	var x int
175	var y int
176	var z int
177	`,
178	want: `package p
179
180	var y int
181	var x1 int
182	var z int
183	`,
184	pre: func(c *astutil.Cursor) bool {
185	n := c.Node()
186	if d, ok := n.(ast.GenDecl); ok && d.Specs[0].(ast.ValueSpec).Names[0].Name == "x" {
187	c.Delete()
188	// The cursor is now effectively atop the 'var y int' node.
189	c.InsertAfter(vardecl("x1", "int"))
190	}
191	return true
192	},
193	},
194	}
195
196	func init() {
197	if typeparams.Enabled {
198	rewriteTests = append(rewriteTests, rewriteTest{
199	name: "replace",
200	orig: `package p
201
202	type T[P1, P2 any] int
203
204	type R T[int, string]
205
206	func F[Q1 any](q Q1) {}
207	`,
208	// TODO: note how the rewrite adds a trailing comma in "func F".
209	// Is that a bug in the test, or in astutil.Apply?
210	want: `package p
211
212	type S[R1, P2 any] int32
213
214	type R S[int32, string]
215
216	func F[X1 any](q X1,) {}
217	`,
218	post: func(c *astutil.Cursor) bool {
219	if ident, ok := c.Node().(*ast.Ident); ok {
220	switch ident.Name {
221	case "int":
222	c.Replace(ast.NewIdent("int32"))
223	case "T":
224	c.Replace(ast.NewIdent("S"))
225	case "P1":
226	c.Replace(ast.NewIdent("R1"))
227	case "Q1":
228	c.Replace(ast.NewIdent("X1"))
229	}
230	}
231	return true
232	},
233	})
234	}
235	}
236
237	func valspec(name, typ string) *ast.ValueSpec {
238	return &ast.ValueSpec{Names: []*ast.Ident{ast.NewIdent(name)},
239	Type: ast.NewIdent(typ),
240	}
241	}
242
243	func vardecl(name, typ string) *ast.GenDecl {
244	return &ast.GenDecl{
245	Tok: token.VAR,
246	Specs: []ast.Spec{valspec(name, typ)},
247	}
248	}
249
250	func TestRewrite(t *testing.T) {
251	t.Run("", func(t testing.T) {
252	for _, test := range rewriteTests {
253	test := test
254	t.Run(test.name, func(t *testing.T) {
255	t.Parallel()
256	fset := token.NewFileSet()
257	f, err := parser.ParseFile(fset, test.name, test.orig, parser.ParseComments)
258	if err != nil {
259	t.Fatal(err)
260	}
261	n := astutil.Apply(f, test.pre, test.post)
262	var buf bytes.Buffer
263	if err := format.Node(&buf, fset, n); err != nil {
264	t.Fatal(err)
265	}
266	got := buf.String()
267	if got != test.want {
268	t.Errorf("got:\n\n%s\nwant:\n\n%s\n", got, test.want)
269	}
270	})
271	}
272	})
273	}
274
275	var sink ast.Node
276
277	func BenchmarkRewrite(b *testing.B) {
278	for _, test := range rewriteTests {
279	b.Run(test.name, func(b *testing.B) {
280	for i := 0; i < b.N; i++ {
281	b.StopTimer()
282	fset := token.NewFileSet()
283	f, err := parser.ParseFile(fset, test.name, test.orig, parser.ParseComments)
284	if err != nil {
285	b.Fatal(err)
286	}
287	b.StartTimer()
288	sink = astutil.Apply(f, test.pre, test.post)
289	}
290	})
291	}
292	}
293