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  
    56  func lineAt(fset *token.FileSet, pos token.Pos) int {
    57  	return fset.PositionFor(pos, false).Line
    58  }
    59  
    60  func 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  
    68  func importName(s Spec) string {
    69  	n := s.(*ImportSpec).Name
    70  	if n == nil {
    71  		return ""
    72  	}
    73  	return n.Name
    74  }
    75  
    76  func 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.
    85  func 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  
    92  type posSpan struct {
    93  	Start token.Pos
    94  	End   token.Pos
    95  }
    96  
    97  type cgPos struct {
    98  	left bool // true if comment is to the left of the spec, false otherwise.
    99  	cg   *CommentGroup
   100  }
   101  
   102  func 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
View as plain text