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