• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1// Copyright 2011 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
5package ast
6
7import (
8	"cmp"
9	"go/token"
10	"slices"
11	"strconv"
12)
13
14// SortImports sorts runs of consecutive import lines in import blocks in f.
15// It also removes duplicate imports when it is possible to do so without data loss.
16func SortImports(fset *token.FileSet, f *File) {
17	for _, d := range f.Decls {
18		d, ok := d.(*GenDecl)
19		if !ok || d.Tok != token.IMPORT {
20			// Not an import declaration, so we're done.
21			// Imports are always first.
22			break
23		}
24
25		if !d.Lparen.IsValid() {
26			// Not a block: sorted by default.
27			continue
28		}
29
30		// Identify and sort runs of specs on successive lines.
31		i := 0
32		specs := d.Specs[:0]
33		for j, s := range d.Specs {
34			if j > i && lineAt(fset, s.Pos()) > 1+lineAt(fset, d.Specs[j-1].End()) {
35				// j begins a new run. End this one.
36				specs = append(specs, sortSpecs(fset, f, d.Specs[i:j])...)
37				i = j
38			}
39		}
40		specs = append(specs, sortSpecs(fset, f, d.Specs[i:])...)
41		d.Specs = specs
42
43		// Deduping can leave a blank line before the rparen; clean that up.
44		if len(d.Specs) > 0 {
45			lastSpec := d.Specs[len(d.Specs)-1]
46			lastLine := lineAt(fset, lastSpec.Pos())
47			rParenLine := lineAt(fset, d.Rparen)
48			for rParenLine > lastLine+1 {
49				rParenLine--
50				fset.File(d.Rparen).MergeLine(rParenLine)
51			}
52		}
53	}
54}
55
56func lineAt(fset *token.FileSet, pos token.Pos) int {
57	return fset.PositionFor(pos, false).Line
58}
59
60func importPath(s Spec) string {
61	t, err := strconv.Unquote(s.(*ImportSpec).Path.Value)
62	if err == nil {
63		return t
64	}
65	return ""
66}
67
68func importName(s Spec) string {
69	n := s.(*ImportSpec).Name
70	if n == nil {
71		return ""
72	}
73	return n.Name
74}
75
76func importComment(s Spec) string {
77	c := s.(*ImportSpec).Comment
78	if c == nil {
79		return ""
80	}
81	return c.Text()
82}
83
84// collapse indicates whether prev may be removed, leaving only next.
85func collapse(prev, next Spec) bool {
86	if importPath(next) != importPath(prev) || importName(next) != importName(prev) {
87		return false
88	}
89	return prev.(*ImportSpec).Comment == nil
90}
91
92type posSpan struct {
93	Start token.Pos
94	End   token.Pos
95}
96
97type cgPos struct {
98	left bool // true if comment is to the left of the spec, false otherwise.
99	cg   *CommentGroup
100}
101
102func sortSpecs(fset *token.FileSet, f *File, specs []Spec) []Spec {
103	// Can't short-circuit here even if specs are already sorted,
104	// since they might yet need deduplication.
105	// A lone import, however, may be safely ignored.
106	if len(specs) <= 1 {
107		return specs
108	}
109
110	// Record positions for specs.
111	pos := make([]posSpan, len(specs))
112	for i, s := range specs {
113		pos[i] = posSpan{s.Pos(), s.End()}
114	}
115
116	// Identify comments in this range.
117	begSpecs := pos[0].Start
118	endSpecs := pos[len(pos)-1].End
119	beg := fset.File(begSpecs).LineStart(lineAt(fset, begSpecs))
120	endLine := lineAt(fset, endSpecs)
121	endFile := fset.File(endSpecs)
122	var end token.Pos
123	if endLine == endFile.LineCount() {
124		end = endSpecs
125	} else {
126		end = endFile.LineStart(endLine + 1) // beginning of next line
127	}
128	first := len(f.Comments)
129	last := -1
130	for i, g := range f.Comments {
131		if g.End() >= end {
132			break
133		}
134		// g.End() < end
135		if beg <= g.Pos() {
136			// comment is within the range [beg, end[ of import declarations
137			if i < first {
138				first = i
139			}
140			if i > last {
141				last = i
142			}
143		}
144	}
145
146	var comments []*CommentGroup
147	if last >= 0 {
148		comments = f.Comments[first : last+1]
149	}
150
151	// Assign each comment to the import spec on the same line.
152	importComments := map[*ImportSpec][]cgPos{}
153	specIndex := 0
154	for _, g := range comments {
155		for specIndex+1 < len(specs) && pos[specIndex+1].Start <= g.Pos() {
156			specIndex++
157		}
158		var left bool
159		// A block comment can appear before the first import spec.
160		if specIndex == 0 && pos[specIndex].Start > g.Pos() {
161			left = true
162		} else if specIndex+1 < len(specs) && // Or it can appear on the left of an import spec.
163			lineAt(fset, pos[specIndex].Start)+1 == lineAt(fset, g.Pos()) {
164			specIndex++
165			left = true
166		}
167		s := specs[specIndex].(*ImportSpec)
168		importComments[s] = append(importComments[s], cgPos{left: left, cg: g})
169	}
170
171	// Sort the import specs by import path.
172	// Remove duplicates, when possible without data loss.
173	// Reassign the import paths to have the same position sequence.
174	// Reassign each comment to the spec on the same line.
175	// Sort the comments by new position.
176	slices.SortFunc(specs, func(a, b Spec) int {
177		ipath := importPath(a)
178		jpath := importPath(b)
179		r := cmp.Compare(ipath, jpath)
180		if r != 0 {
181			return r
182		}
183		iname := importName(a)
184		jname := importName(b)
185		r = cmp.Compare(iname, jname)
186		if r != 0 {
187			return r
188		}
189		return cmp.Compare(importComment(a), importComment(b))
190	})
191
192	// Dedup. Thanks to our sorting, we can just consider
193	// adjacent pairs of imports.
194	deduped := specs[:0]
195	for i, s := range specs {
196		if i == len(specs)-1 || !collapse(s, specs[i+1]) {
197			deduped = append(deduped, s)
198		} else {
199			p := s.Pos()
200			fset.File(p).MergeLine(lineAt(fset, p))
201		}
202	}
203	specs = deduped
204
205	// Fix up comment positions
206	for i, s := range specs {
207		s := s.(*ImportSpec)
208		if s.Name != nil {
209			s.Name.NamePos = pos[i].Start
210		}
211		s.Path.ValuePos = pos[i].Start
212		s.EndPos = pos[i].End
213		for _, g := range importComments[s] {
214			for _, c := range g.cg.List {
215				if g.left {
216					c.Slash = pos[i].Start - 1
217				} else {
218					// An import spec can have both block comment and a line comment
219					// to its right. In that case, both of them will have the same pos.
220					// But while formatting the AST, the line comment gets moved to
221					// after the block comment.
222					c.Slash = pos[i].End
223				}
224			}
225		}
226	}
227
228	slices.SortFunc(comments, func(a, b *CommentGroup) int {
229		return cmp.Compare(a.Pos(), b.Pos())
230	})
231
232	return specs
233}
234