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	"go/token"
9	"sort"
10	"strconv"
11)
12
13// SortImports sorts runs of consecutive import lines in import blocks in f.
14// It also removes duplicate imports when it is possible to do so without data loss.
15func SortImports(fset *token.FileSet, f *File) {
16	for _, d := range f.Decls {
17		d, ok := d.(*GenDecl)
18		if !ok || d.Tok != token.IMPORT {
19			// Not an import declaration, so we're done.
20			// Imports are always first.
21			break
22		}
23
24		if !d.Lparen.IsValid() {
25			// Not a block: sorted by default.
26			continue
27		}
28
29		// Identify and sort runs of specs on successive lines.
30		i := 0
31		specs := d.Specs[:0]
32		for j, s := range d.Specs {
33			if j > i && fset.Position(s.Pos()).Line > 1+fset.Position(d.Specs[j-1].End()).Line {
34				// j begins a new run. End this one.
35				specs = append(specs, sortSpecs(fset, f, d.Specs[i:j])...)
36				i = j
37			}
38		}
39		specs = append(specs, sortSpecs(fset, f, d.Specs[i:])...)
40		d.Specs = specs
41
42		// Deduping can leave a blank line before the rparen; clean that up.
43		if len(d.Specs) > 0 {
44			lastSpec := d.Specs[len(d.Specs)-1]
45			lastLine := fset.Position(lastSpec.Pos()).Line
46			rParenLine := fset.Position(d.Rparen).Line
47			for rParenLine > lastLine+1 {
48				rParenLine--
49				fset.File(d.Rparen).MergeLine(rParenLine)
50			}
51		}
52	}
53}
54
55func importPath(s Spec) string {
56	t, err := strconv.Unquote(s.(*ImportSpec).Path.Value)
57	if err == nil {
58		return t
59	}
60	return ""
61}
62
63func importName(s Spec) string {
64	n := s.(*ImportSpec).Name
65	if n == nil {
66		return ""
67	}
68	return n.Name
69}
70
71func importComment(s Spec) string {
72	c := s.(*ImportSpec).Comment
73	if c == nil {
74		return ""
75	}
76	return c.Text()
77}
78
79// collapse indicates whether prev may be removed, leaving only next.
80func collapse(prev, next Spec) bool {
81	if importPath(next) != importPath(prev) || importName(next) != importName(prev) {
82		return false
83	}
84	return prev.(*ImportSpec).Comment == nil
85}
86
87type posSpan struct {
88	Start token.Pos
89	End   token.Pos
90}
91
92func sortSpecs(fset *token.FileSet, f *File, specs []Spec) []Spec {
93	// Can't short-circuit here even if specs are already sorted,
94	// since they might yet need deduplication.
95	// A lone import, however, may be safely ignored.
96	if len(specs) <= 1 {
97		return specs
98	}
99
100	// Record positions for specs.
101	pos := make([]posSpan, len(specs))
102	for i, s := range specs {
103		pos[i] = posSpan{s.Pos(), s.End()}
104	}
105
106	// Identify comments in this range.
107	// Any comment from pos[0].Start to the final line counts.
108	lastLine := fset.Position(pos[len(pos)-1].End).Line
109	cstart := len(f.Comments)
110	cend := len(f.Comments)
111	for i, g := range f.Comments {
112		if g.Pos() < pos[0].Start {
113			continue
114		}
115		if i < cstart {
116			cstart = i
117		}
118		if fset.Position(g.End()).Line > lastLine {
119			cend = i
120			break
121		}
122	}
123	comments := f.Comments[cstart:cend]
124
125	// Assign each comment to the import spec preceding it.
126	importComments := map[*ImportSpec][]*CommentGroup{}
127	specIndex := 0
128	for _, g := range comments {
129		for specIndex+1 < len(specs) && pos[specIndex+1].Start <= g.Pos() {
130			specIndex++
131		}
132		s := specs[specIndex].(*ImportSpec)
133		importComments[s] = append(importComments[s], g)
134	}
135
136	// Sort the import specs by import path.
137	// Remove duplicates, when possible without data loss.
138	// Reassign the import paths to have the same position sequence.
139	// Reassign each comment to abut the end of its spec.
140	// Sort the comments by new position.
141	sort.Slice(specs, func(i, j int) bool {
142		ipath := importPath(specs[i])
143		jpath := importPath(specs[j])
144		if ipath != jpath {
145			return ipath < jpath
146		}
147		iname := importName(specs[i])
148		jname := importName(specs[j])
149		if iname != jname {
150			return iname < jname
151		}
152		return importComment(specs[i]) < importComment(specs[j])
153	})
154
155	// Dedup. Thanks to our sorting, we can just consider
156	// adjacent pairs of imports.
157	deduped := specs[:0]
158	for i, s := range specs {
159		if i == len(specs)-1 || !collapse(s, specs[i+1]) {
160			deduped = append(deduped, s)
161		} else {
162			p := s.Pos()
163			fset.File(p).MergeLine(fset.Position(p).Line)
164		}
165	}
166	specs = deduped
167
168	// Fix up comment positions
169	for i, s := range specs {
170		s := s.(*ImportSpec)
171		if s.Name != nil {
172			s.Name.NamePos = pos[i].Start
173		}
174		s.Path.ValuePos = pos[i].Start
175		s.EndPos = pos[i].End
176		for _, g := range importComments[s] {
177			for _, c := range g.List {
178				c.Slash = pos[i].End
179			}
180		}
181	}
182
183	sort.Slice(comments, func(i, j int) bool {
184		return comments[i].Pos() < comments[j].Pos()
185	})
186
187	return specs
188}
189