1// Copyright 2013 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 contains common utilities for working with the Go AST.
6package astutil // import "golang.org/x/tools/go/ast/astutil"
7
8import (
9	"fmt"
10	"go/ast"
11	"go/token"
12	"strconv"
13	"strings"
14)
15
16// AddImport adds the import path to the file f, if absent.
17func AddImport(fset *token.FileSet, f *ast.File, path string) (added bool) {
18	return AddNamedImport(fset, f, "", path)
19}
20
21// AddNamedImport adds the import with the given name and path to the file f, if absent.
22// If name is not empty, it is used to rename the import.
23//
24// For example, calling
25//
26//	AddNamedImport(fset, f, "pathpkg", "path")
27//
28// adds
29//
30//	import pathpkg "path"
31func AddNamedImport(fset *token.FileSet, f *ast.File, name, path string) (added bool) {
32	if imports(f, name, path) {
33		return false
34	}
35
36	newImport := &ast.ImportSpec{
37		Path: &ast.BasicLit{
38			Kind:  token.STRING,
39			Value: strconv.Quote(path),
40		},
41	}
42	if name != "" {
43		newImport.Name = &ast.Ident{Name: name}
44	}
45
46	// Find an import decl to add to.
47	// The goal is to find an existing import
48	// whose import path has the longest shared
49	// prefix with path.
50	var (
51		bestMatch  = -1         // length of longest shared prefix
52		lastImport = -1         // index in f.Decls of the file's final import decl
53		impDecl    *ast.GenDecl // import decl containing the best match
54		impIndex   = -1         // spec index in impDecl containing the best match
55
56		isThirdPartyPath = isThirdParty(path)
57	)
58	for i, decl := range f.Decls {
59		gen, ok := decl.(*ast.GenDecl)
60		if ok && gen.Tok == token.IMPORT {
61			lastImport = i
62			// Do not add to import "C", to avoid disrupting the
63			// association with its doc comment, breaking cgo.
64			if declImports(gen, "C") {
65				continue
66			}
67
68			// Match an empty import decl if that's all that is available.
69			if len(gen.Specs) == 0 && bestMatch == -1 {
70				impDecl = gen
71			}
72
73			// Compute longest shared prefix with imports in this group and find best
74			// matched import spec.
75			// 1. Always prefer import spec with longest shared prefix.
76			// 2. While match length is 0,
77			// - for stdlib package: prefer first import spec.
78			// - for third party package: prefer first third party import spec.
79			// We cannot use last import spec as best match for third party package
80			// because grouped imports are usually placed last by goimports -local
81			// flag.
82			// See issue #19190.
83			seenAnyThirdParty := false
84			for j, spec := range gen.Specs {
85				impspec := spec.(*ast.ImportSpec)
86				p := importPath(impspec)
87				n := matchLen(p, path)
88				if n > bestMatch || (bestMatch == 0 && !seenAnyThirdParty && isThirdPartyPath) {
89					bestMatch = n
90					impDecl = gen
91					impIndex = j
92				}
93				seenAnyThirdParty = seenAnyThirdParty || isThirdParty(p)
94			}
95		}
96	}
97
98	// If no import decl found, add one after the last import.
99	if impDecl == nil {
100		impDecl = &ast.GenDecl{
101			Tok: token.IMPORT,
102		}
103		if lastImport >= 0 {
104			impDecl.TokPos = f.Decls[lastImport].End()
105		} else {
106			// There are no existing imports.
107			// Our new import, preceded by a blank line,  goes after the package declaration
108			// and after the comment, if any, that starts on the same line as the
109			// package declaration.
110			impDecl.TokPos = f.Package
111
112			file := fset.File(f.Package)
113			pkgLine := file.Line(f.Package)
114			for _, c := range f.Comments {
115				if file.Line(c.Pos()) > pkgLine {
116					break
117				}
118				// +2 for a blank line
119				impDecl.TokPos = c.End() + 2
120			}
121		}
122		f.Decls = append(f.Decls, nil)
123		copy(f.Decls[lastImport+2:], f.Decls[lastImport+1:])
124		f.Decls[lastImport+1] = impDecl
125	}
126
127	// Insert new import at insertAt.
128	insertAt := 0
129	if impIndex >= 0 {
130		// insert after the found import
131		insertAt = impIndex + 1
132	}
133	impDecl.Specs = append(impDecl.Specs, nil)
134	copy(impDecl.Specs[insertAt+1:], impDecl.Specs[insertAt:])
135	impDecl.Specs[insertAt] = newImport
136	pos := impDecl.Pos()
137	if insertAt > 0 {
138		// If there is a comment after an existing import, preserve the comment
139		// position by adding the new import after the comment.
140		if spec, ok := impDecl.Specs[insertAt-1].(*ast.ImportSpec); ok && spec.Comment != nil {
141			pos = spec.Comment.End()
142		} else {
143			// Assign same position as the previous import,
144			// so that the sorter sees it as being in the same block.
145			pos = impDecl.Specs[insertAt-1].Pos()
146		}
147	}
148	if newImport.Name != nil {
149		newImport.Name.NamePos = pos
150	}
151	newImport.Path.ValuePos = pos
152	newImport.EndPos = pos
153
154	// Clean up parens. impDecl contains at least one spec.
155	if len(impDecl.Specs) == 1 {
156		// Remove unneeded parens.
157		impDecl.Lparen = token.NoPos
158	} else if !impDecl.Lparen.IsValid() {
159		// impDecl needs parens added.
160		impDecl.Lparen = impDecl.Specs[0].Pos()
161	}
162
163	f.Imports = append(f.Imports, newImport)
164
165	if len(f.Decls) <= 1 {
166		return true
167	}
168
169	// Merge all the import declarations into the first one.
170	var first *ast.GenDecl
171	for i := 0; i < len(f.Decls); i++ {
172		decl := f.Decls[i]
173		gen, ok := decl.(*ast.GenDecl)
174		if !ok || gen.Tok != token.IMPORT || declImports(gen, "C") {
175			continue
176		}
177		if first == nil {
178			first = gen
179			continue // Don't touch the first one.
180		}
181		// We now know there is more than one package in this import
182		// declaration. Ensure that it ends up parenthesized.
183		first.Lparen = first.Pos()
184		// Move the imports of the other import declaration to the first one.
185		for _, spec := range gen.Specs {
186			spec.(*ast.ImportSpec).Path.ValuePos = first.Pos()
187			first.Specs = append(first.Specs, spec)
188		}
189		f.Decls = append(f.Decls[:i], f.Decls[i+1:]...)
190		i--
191	}
192
193	return true
194}
195
196func isThirdParty(importPath string) bool {
197	// Third party package import path usually contains "." (".com", ".org", ...)
198	// This logic is taken from golang.org/x/tools/imports package.
199	return strings.Contains(importPath, ".")
200}
201
202// DeleteImport deletes the import path from the file f, if present.
203// If there are duplicate import declarations, all matching ones are deleted.
204func DeleteImport(fset *token.FileSet, f *ast.File, path string) (deleted bool) {
205	return DeleteNamedImport(fset, f, "", path)
206}
207
208// DeleteNamedImport deletes the import with the given name and path from the file f, if present.
209// If there are duplicate import declarations, all matching ones are deleted.
210func DeleteNamedImport(fset *token.FileSet, f *ast.File, name, path string) (deleted bool) {
211	var delspecs []*ast.ImportSpec
212	var delcomments []*ast.CommentGroup
213
214	// Find the import nodes that import path, if any.
215	for i := 0; i < len(f.Decls); i++ {
216		decl := f.Decls[i]
217		gen, ok := decl.(*ast.GenDecl)
218		if !ok || gen.Tok != token.IMPORT {
219			continue
220		}
221		for j := 0; j < len(gen.Specs); j++ {
222			spec := gen.Specs[j]
223			impspec := spec.(*ast.ImportSpec)
224			if importName(impspec) != name || importPath(impspec) != path {
225				continue
226			}
227
228			// We found an import spec that imports path.
229			// Delete it.
230			delspecs = append(delspecs, impspec)
231			deleted = true
232			copy(gen.Specs[j:], gen.Specs[j+1:])
233			gen.Specs = gen.Specs[:len(gen.Specs)-1]
234
235			// If this was the last import spec in this decl,
236			// delete the decl, too.
237			if len(gen.Specs) == 0 {
238				copy(f.Decls[i:], f.Decls[i+1:])
239				f.Decls = f.Decls[:len(f.Decls)-1]
240				i--
241				break
242			} else if len(gen.Specs) == 1 {
243				if impspec.Doc != nil {
244					delcomments = append(delcomments, impspec.Doc)
245				}
246				if impspec.Comment != nil {
247					delcomments = append(delcomments, impspec.Comment)
248				}
249				for _, cg := range f.Comments {
250					// Found comment on the same line as the import spec.
251					if cg.End() < impspec.Pos() && fset.Position(cg.End()).Line == fset.Position(impspec.Pos()).Line {
252						delcomments = append(delcomments, cg)
253						break
254					}
255				}
256
257				spec := gen.Specs[0].(*ast.ImportSpec)
258
259				// Move the documentation right after the import decl.
260				if spec.Doc != nil {
261					for fset.Position(gen.TokPos).Line+1 < fset.Position(spec.Doc.Pos()).Line {
262						fset.File(gen.TokPos).MergeLine(fset.Position(gen.TokPos).Line)
263					}
264				}
265				for _, cg := range f.Comments {
266					if cg.End() < spec.Pos() && fset.Position(cg.End()).Line == fset.Position(spec.Pos()).Line {
267						for fset.Position(gen.TokPos).Line+1 < fset.Position(spec.Pos()).Line {
268							fset.File(gen.TokPos).MergeLine(fset.Position(gen.TokPos).Line)
269						}
270						break
271					}
272				}
273			}
274			if j > 0 {
275				lastImpspec := gen.Specs[j-1].(*ast.ImportSpec)
276				lastLine := fset.PositionFor(lastImpspec.Path.ValuePos, false).Line
277				line := fset.PositionFor(impspec.Path.ValuePos, false).Line
278
279				// We deleted an entry but now there may be
280				// a blank line-sized hole where the import was.
281				if line-lastLine > 1 || !gen.Rparen.IsValid() {
282					// There was a blank line immediately preceding the deleted import,
283					// so there's no need to close the hole. The right parenthesis is
284					// invalid after AddImport to an import statement without parenthesis.
285					// Do nothing.
286				} else if line != fset.File(gen.Rparen).LineCount() {
287					// There was no blank line. Close the hole.
288					fset.File(gen.Rparen).MergeLine(line)
289				}
290			}
291			j--
292		}
293	}
294
295	// Delete imports from f.Imports.
296	for i := 0; i < len(f.Imports); i++ {
297		imp := f.Imports[i]
298		for j, del := range delspecs {
299			if imp == del {
300				copy(f.Imports[i:], f.Imports[i+1:])
301				f.Imports = f.Imports[:len(f.Imports)-1]
302				copy(delspecs[j:], delspecs[j+1:])
303				delspecs = delspecs[:len(delspecs)-1]
304				i--
305				break
306			}
307		}
308	}
309
310	// Delete comments from f.Comments.
311	for i := 0; i < len(f.Comments); i++ {
312		cg := f.Comments[i]
313		for j, del := range delcomments {
314			if cg == del {
315				copy(f.Comments[i:], f.Comments[i+1:])
316				f.Comments = f.Comments[:len(f.Comments)-1]
317				copy(delcomments[j:], delcomments[j+1:])
318				delcomments = delcomments[:len(delcomments)-1]
319				i--
320				break
321			}
322		}
323	}
324
325	if len(delspecs) > 0 {
326		panic(fmt.Sprintf("deleted specs from Decls but not Imports: %v", delspecs))
327	}
328
329	return
330}
331
332// RewriteImport rewrites any import of path oldPath to path newPath.
333func RewriteImport(fset *token.FileSet, f *ast.File, oldPath, newPath string) (rewrote bool) {
334	for _, imp := range f.Imports {
335		if importPath(imp) == oldPath {
336			rewrote = true
337			// record old End, because the default is to compute
338			// it using the length of imp.Path.Value.
339			imp.EndPos = imp.End()
340			imp.Path.Value = strconv.Quote(newPath)
341		}
342	}
343	return
344}
345
346// UsesImport reports whether a given import is used.
347func UsesImport(f *ast.File, path string) (used bool) {
348	spec := importSpec(f, path)
349	if spec == nil {
350		return
351	}
352
353	name := spec.Name.String()
354	switch name {
355	case "<nil>":
356		// If the package name is not explicitly specified,
357		// make an educated guess. This is not guaranteed to be correct.
358		lastSlash := strings.LastIndex(path, "/")
359		if lastSlash == -1 {
360			name = path
361		} else {
362			name = path[lastSlash+1:]
363		}
364	case "_", ".":
365		// Not sure if this import is used - err on the side of caution.
366		return true
367	}
368
369	ast.Walk(visitFn(func(n ast.Node) {
370		sel, ok := n.(*ast.SelectorExpr)
371		if ok && isTopName(sel.X, name) {
372			used = true
373		}
374	}), f)
375
376	return
377}
378
379type visitFn func(node ast.Node)
380
381func (fn visitFn) Visit(node ast.Node) ast.Visitor {
382	fn(node)
383	return fn
384}
385
386// imports reports whether f has an import with the specified name and path.
387func imports(f *ast.File, name, path string) bool {
388	for _, s := range f.Imports {
389		if importName(s) == name && importPath(s) == path {
390			return true
391		}
392	}
393	return false
394}
395
396// importSpec returns the import spec if f imports path,
397// or nil otherwise.
398func importSpec(f *ast.File, path string) *ast.ImportSpec {
399	for _, s := range f.Imports {
400		if importPath(s) == path {
401			return s
402		}
403	}
404	return nil
405}
406
407// importName returns the name of s,
408// or "" if the import is not named.
409func importName(s *ast.ImportSpec) string {
410	if s.Name == nil {
411		return ""
412	}
413	return s.Name.Name
414}
415
416// importPath returns the unquoted import path of s,
417// or "" if the path is not properly quoted.
418func importPath(s *ast.ImportSpec) string {
419	t, err := strconv.Unquote(s.Path.Value)
420	if err != nil {
421		return ""
422	}
423	return t
424}
425
426// declImports reports whether gen contains an import of path.
427func declImports(gen *ast.GenDecl, path string) bool {
428	if gen.Tok != token.IMPORT {
429		return false
430	}
431	for _, spec := range gen.Specs {
432		impspec := spec.(*ast.ImportSpec)
433		if importPath(impspec) == path {
434			return true
435		}
436	}
437	return false
438}
439
440// matchLen returns the length of the longest path segment prefix shared by x and y.
441func matchLen(x, y string) int {
442	n := 0
443	for i := 0; i < len(x) && i < len(y) && x[i] == y[i]; i++ {
444		if x[i] == '/' {
445			n++
446		}
447	}
448	return n
449}
450
451// isTopName returns true if n is a top-level unresolved identifier with the given name.
452func isTopName(n ast.Expr, name string) bool {
453	id, ok := n.(*ast.Ident)
454	return ok && id.Name == name && id.Obj == nil
455}
456
457// Imports returns the file imports grouped by paragraph.
458func Imports(fset *token.FileSet, f *ast.File) [][]*ast.ImportSpec {
459	var groups [][]*ast.ImportSpec
460
461	for _, decl := range f.Decls {
462		genDecl, ok := decl.(*ast.GenDecl)
463		if !ok || genDecl.Tok != token.IMPORT {
464			break
465		}
466
467		group := []*ast.ImportSpec{}
468
469		var lastLine int
470		for _, spec := range genDecl.Specs {
471			importSpec := spec.(*ast.ImportSpec)
472			pos := importSpec.Path.ValuePos
473			line := fset.Position(pos).Line
474			if lastLine > 0 && pos > 0 && line-lastLine > 1 {
475				groups = append(groups, group)
476				group = []*ast.ImportSpec{}
477			}
478			group = append(group, importSpec)
479			lastLine = line
480		}
481		groups = append(groups, group)
482	}
483
484	return groups
485}
486