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