import.go

     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  
     5  package ast
     6  
     7  import (
     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.
    16  func 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  	// Make File.Imports order consistent.
    56  	f.Imports = f.Imports[:0]
    57  	for _, decl := range f.Decls {
    58  		if decl, ok := decl.(*GenDecl); ok && decl.Tok == token.IMPORT {
    59  			for _, spec := range decl.Specs {
    60  				f.Imports = append(f.Imports, spec.(*ImportSpec))
    61  			}
    62  		}
    63  	}
    64  }
    65  
    66  func lineAt(fset *token.FileSet, pos token.Pos) int {
    67  	return fset.PositionFor(pos, false).Line
    68  }
    69  
    70  func importPath(s Spec) string {
    71  	t, err := strconv.Unquote(s.(*ImportSpec).Path.Value)
    72  	if err == nil {
    73  		return t
    74  	}
    75  	return ""
    76  }
    77  
    78  func importName(s Spec) string {
    79  	n := s.(*ImportSpec).Name
    80  	if n == nil {
    81  		return ""
    82  	}
    83  	return n.Name
    84  }
    85  
    86  func importComment(s Spec) string {
    87  	c := s.(*ImportSpec).Comment
    88  	if c == nil {
    89  		return ""
    90  	}
    91  	return c.Text()
    92  }
    93  
    94  // collapse indicates whether prev may be removed, leaving only next.
    95  func collapse(prev, next Spec) bool {
    96  	if importPath(next) != importPath(prev) || importName(next) != importName(prev) {
    97  		return false
    98  	}
    99  	return prev.(*ImportSpec).Comment == nil
   100  }
   101  
   102  type posSpan struct {
   103  	Start token.Pos
   104  	End   token.Pos
   105  }
   106  
   107  type cgPos struct {
   108  	left bool // true if comment is to the left of the spec, false otherwise.
   109  	cg   *CommentGroup
   110  }
   111  
   112  func sortSpecs(fset *token.FileSet, f *File, specs []Spec) []Spec {
   113  	// Can't short-circuit here even if specs are already sorted,
   114  	// since they might yet need deduplication.
   115  	// A lone import, however, may be safely ignored.
   116  	if len(specs) <= 1 {
   117  		return specs
   118  	}
   119  
   120  	// Record positions for specs.
   121  	pos := make([]posSpan, len(specs))
   122  	for i, s := range specs {
   123  		pos[i] = posSpan{s.Pos(), s.End()}
   124  	}
   125  
   126  	// Identify comments in this range.
   127  	begSpecs := pos[0].Start
   128  	endSpecs := pos[len(pos)-1].End
   129  	beg := fset.File(begSpecs).LineStart(lineAt(fset, begSpecs))
   130  	endLine := lineAt(fset, endSpecs)
   131  	endFile := fset.File(endSpecs)
   132  	var end token.Pos
   133  	if endLine == endFile.LineCount() {
   134  		end = endSpecs
   135  	} else {
   136  		end = endFile.LineStart(endLine + 1) // beginning of next line
   137  	}
   138  	first := len(f.Comments)
   139  	last := -1
   140  	for i, g := range f.Comments {
   141  		if g.End() >= end {
   142  			break
   143  		}
   144  		// g.End() < end
   145  		if beg <= g.Pos() {
   146  			// comment is within the range [beg, end[ of import declarations
   147  			if i < first {
   148  				first = i
   149  			}
   150  			if i > last {
   151  				last = i
   152  			}
   153  		}
   154  	}
   155  
   156  	var comments []*CommentGroup
   157  	if last >= 0 {
   158  		comments = f.Comments[first : last+1]
   159  	}
   160  
   161  	// Assign each comment to the import spec on the same line.
   162  	importComments := map[*ImportSpec][]cgPos{}
   163  	specIndex := 0
   164  	for _, g := range comments {
   165  		for specIndex+1 < len(specs) && pos[specIndex+1].Start <= g.Pos() {
   166  			specIndex++
   167  		}
   168  		var left bool
   169  		// A block comment can appear before the first import spec.
   170  		if specIndex == 0 && pos[specIndex].Start > g.Pos() {
   171  			left = true
   172  		} else if specIndex+1 < len(specs) && // Or it can appear on the left of an import spec.
   173  			lineAt(fset, pos[specIndex].Start)+1 == lineAt(fset, g.Pos()) {
   174  			specIndex++
   175  			left = true
   176  		}
   177  		s := specs[specIndex].(*ImportSpec)
   178  		importComments[s] = append(importComments[s], cgPos{left: left, cg: g})
   179  	}
   180  
   181  	// Sort the import specs by import path.
   182  	// Remove duplicates, when possible without data loss.
   183  	// Reassign the import paths to have the same position sequence.
   184  	// Reassign each comment to the spec on the same line.
   185  	// Sort the comments by new position.
   186  	slices.SortFunc(specs, func(a, b Spec) int {
   187  		ipath := importPath(a)
   188  		jpath := importPath(b)
   189  		r := cmp.Compare(ipath, jpath)
   190  		if r != 0 {
   191  			return r
   192  		}
   193  		iname := importName(a)
   194  		jname := importName(b)
   195  		r = cmp.Compare(iname, jname)
   196  		if r != 0 {
   197  			return r
   198  		}
   199  		return cmp.Compare(importComment(a), importComment(b))
   200  	})
   201  
   202  	// Dedup. Thanks to our sorting, we can just consider
   203  	// adjacent pairs of imports.
   204  	deduped := specs[:0]
   205  	for i, s := range specs {
   206  		if i == len(specs)-1 || !collapse(s, specs[i+1]) {
   207  			deduped = append(deduped, s)
   208  		} else {
   209  			p := s.Pos()
   210  			fset.File(p).MergeLine(lineAt(fset, p))
   211  		}
   212  	}
   213  	specs = deduped
   214  
   215  	// Fix up comment positions
   216  	for i, s := range specs {
   217  		s := s.(*ImportSpec)
   218  		if s.Name != nil {
   219  			s.Name.NamePos = pos[i].Start
   220  		}
   221  		s.Path.ValuePos = pos[i].Start
   222  		s.EndPos = pos[i].End
   223  		for _, g := range importComments[s] {
   224  			for _, c := range g.cg.List {
   225  				if g.left {
   226  					c.Slash = pos[i].Start - 1
   227  				} else {
   228  					// An import spec can have both block comment and a line comment
   229  					// to its right. In that case, both of them will have the same pos.
   230  					// But while formatting the AST, the line comment gets moved to
   231  					// after the block comment.
   232  					c.Slash = pos[i].End
   233  				}
   234  			}
   235  		}
   236  	}
   237  
   238  	slices.SortFunc(comments, func(a, b *CommentGroup) int {
   239  		return cmp.Compare(a.Pos(), b.Pos())
   240  	})
   241  
   242  	return specs
   243  }
   244
View as plain text