gofmt, gofix: sort imports

Add ast.SortImports(fset, file) to go/ast, for use by both programs.

Fixes #346.

R=golang-dev, r
CC=golang-dev
https://golang.org/cl/5330069

src/cmd/gofix/fix.go

@@ -569,9 +569,9 @@ func renameTop(f *ast.File, old, new string) bool {
 }
 
 // addImport adds the import path to the file f, if absent.
-func addImport(f *ast.File, ipath string) {
+func addImport(f *ast.File, ipath string) (added bool) {
 	if imports(f, ipath) {
-		return
+		return false
 	}
 
 	// Determine name of import.
@@ -637,10 +637,11 @@ func addImport(f *ast.File, ipath string) {
 	impdecl.Specs[insertAt] = newImport
 
 	f.Imports = append(f.Imports, newImport)
+	return true
 }
 
 // deleteImport deletes the import path from the file f, if present.
-func deleteImport(f *ast.File, path string) {
+func deleteImport(f *ast.File, path string) (deleted bool) {
 	oldImport := importSpec(f, path)
 
 	// Find the import node that imports path, if any.
@@ -657,6 +658,7 @@ func deleteImport(f *ast.File, path string) {
 
 			// We found an import spec that imports path.
 			// Delete it.
+			deleted = true
 			copy(gen.Specs[j:], gen.Specs[j+1:])
 			gen.Specs = gen.Specs[:len(gen.Specs)-1]
 
@@ -687,6 +689,19 @@ func deleteImport(f *ast.File, path string) {
 			break
 		}
 	}
+
+	return
+}
+
+// rewriteImport rewrites any import of path oldPath to path newPath.
+func rewriteImport(f *ast.File, oldPath, newPath string) (rewrote bool) {
+	for _, imp := range f.Imports {
+		if importPath(imp) == oldPath {
+			rewrote = true
+			imp.Path.Value = strconv.Quote(newPath)
+		}
+	}
+	return
 }
 
 func usesImport(f *ast.File, path string) (used bool) {

src/cmd/gofix/import_test.go

@@ -244,6 +244,26 @@ import (
 	"io"
 	"os"
 )
+`,
+	},
+	{
+		Name: "import.13",
+		Fn:   rewriteImportFn("utf8", "encoding/utf8"),
+		In: `package main
+
+import (
+	"io"
+	"os"
+	"utf8" // thanks ken
+)
+`,
+		Out: `package main
+
+import (
+	"encoding/utf8" // thanks ken
+	"io"
+	"os"
+)
 `,
 	},
 }
@@ -267,3 +287,13 @@ func deleteImportFn(path string) func(*ast.File) bool {
 		return false
 	}
 }
+
+func rewriteImportFn(old, new string) func(*ast.File) bool {
+	return func(f *ast.File) bool {
+		if imports(f, old) {
+			rewriteImport(f, old, new)
+			return true
+		}
+		return false
+	}
+}

src/cmd/gofix/main.go

@@ -9,6 +9,7 @@ import (
 	"exec"
 	"flag"
 	"fmt"
+	"go/ast"
 	"go/parser"
 	"go/printer"
 	"go/scanner"
@@ -102,11 +103,21 @@ var printConfig = &printer.Config{
 	tabWidth,
 }
 
+func gofmtFile(f *ast.File) ([]byte, error) {
+	var buf bytes.Buffer
+
+	ast.SortImports(fset, f)
+	_, err := printConfig.Fprint(&buf, fset, f)
+	if err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
 func processFile(filename string, useStdin bool) error {
 	var f *os.File
 	var err error
 	var fixlog bytes.Buffer
-	var buf bytes.Buffer
 
 	if useStdin {
 		f = os.Stdin
@@ -142,12 +153,10 @@ func processFile(filename string, useStdin bool) error {
 			// AST changed.
 			// Print and parse, to update any missing scoping
 			// or position information for subsequent fixers.
-			buf.Reset()
-			_, err = printConfig.Fprint(&buf, fset, newFile)
+			newSrc, err := gofmtFile(newFile)
 			if err != nil {
 				return err
 			}
-			newSrc := buf.Bytes()
 			newFile, err = parser.ParseFile(fset, filename, newSrc, parserMode)
 			if err != nil {
 				return err
@@ -165,12 +174,10 @@ func processFile(filename string, useStdin bool) error {
 	// output of the printer run on a standard AST generated by the parser,
 	// but the source we generated inside the loop above is the
 	// output of the printer run on a mangled AST generated by a fixer.
-	buf.Reset()
-	_, err = printConfig.Fprint(&buf, fset, newFile)
+	newSrc, err := gofmtFile(newFile)
 	if err != nil {
 		return err
 	}
-	newSrc := buf.Bytes()
 
 	if *doDiff {
 		data, err := diff(src, newSrc)

src/cmd/gofix/main_test.go

@@ -5,10 +5,8 @@
 package main
 
 import (
-	"bytes"
 	"go/ast"
 	"go/parser"
-	"go/printer"
 	"strings"
 	"testing"
 )
@@ -43,14 +41,12 @@ func parseFixPrint(t *testing.T, fn func(*ast.File) bool, desc, in string) (out
 		return
 	}
 
-	var buf bytes.Buffer
-	buf.Reset()
-	_, err = (&printer.Config{printerMode, tabWidth}).Fprint(&buf, fset, file)
+	outb, err := gofmtFile(file)
 	if err != nil {
 		t.Errorf("%s: printing: %v", desc, err)
 		return
 	}
-	if s := buf.String(); in != s && fn != fnop {
+	if s := string(outb); in != s && fn != fnop {
 		t.Errorf("%s: not gofmt-formatted.\n--- %s\n%s\n--- %s | gofmt\n%s",
 			desc, desc, in, desc, s)
 		tdiff(t, in, s)
@@ -67,14 +63,13 @@ func parseFixPrint(t *testing.T, fn func(*ast.File) bool, desc, in string) (out
 		fixed = fn(file)
 	}
 
-	buf.Reset()
-	_, err = (&printer.Config{printerMode, tabWidth}).Fprint(&buf, fset, file)
+	outb, err = gofmtFile(file)
 	if err != nil {
 		t.Errorf("%s: printing: %v", desc, err)
 		return
 	}
 
-	return buf.String(), fixed, true
+	return string(outb), fixed, true
 }
 
 func TestRewrite(t *testing.T) {

src/cmd/gofix/testdata/reflect.export.go.in

@@ -22,8 +22,8 @@ package netchan
 // BUG: can't use range clause to receive when using ImportNValues to limit the count.
 
 import (
-	"log"
 	"io"
+	"log"
 	"net"
 	"os"
 	"reflect"

src/cmd/gofix/testdata/reflect.export.go.out

@@ -22,8 +22,8 @@ package netchan
 // BUG: can't use range clause to receive when using ImportNValues to limit the count.
 
 import (
-	"log"
 	"io"
+	"log"
 	"net"
 	"os"
 	"reflect"

src/cmd/gofmt/gofmt.go

@@ -114,6 +114,8 @@ func processFile(filename string, in io.Reader, out io.Writer, stdin bool) error
 		}
 	}
 
+	ast.SortImports(fset, file)
+
 	if *simplifyAST {
 		simplify(file)
 	}

src/cmd/gofmt/gofmt_test.go

@@ -78,6 +78,7 @@ var tests = []struct {
 	{"testdata/rewrite2.input", "-r=int->bool"},
 	{"testdata/stdin*.input", "-stdin"},
 	{"testdata/comments.input", ""},
+	{"testdata/import.input", ""},
 }
 
 func TestRewrite(t *testing.T) {

src/cmd/gofmt/testdata/import.golden

+package main
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"log"
+	"math"
+)
+
+import (
+	"fmt"
+
+	"math"
+
+	"log"
+
+	"errors"
+
+	"io"
+)
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"log"
+	"math"
+
+	"fmt"
+
+	"math"
+
+	"log"
+
+	"errors"
+
+	"io"
+)
+
+import (
+	// a block with comments
+	"errors"
+	"fmt" // for Printf
+	"io"  // for Reader
+	"log" // for Fatal
+	"math"
+)
+
+import (
+	"fmt" // for Printf
+
+	"math"
+
+	"log" // for Fatal
+
+	"errors"
+
+	"io" // for Reader
+)
+
+import (
+	// for Printf
+	"fmt"
+
+	"math"
+
+	// for Fatal
+	"log"
+
+	"errors"
+
+	// for Reader
+	"io"
+)
+
+import (
+	"errors"
+	"fmt" // for Printf
+	"io"  // for Reader
+	"log" // for Fatal
+	"math"
+
+	"fmt" // for Printf
+
+	"math"
+
+	"log" // for Fatal
+
+	"errors"
+
+	"io" // for Reader
+)
+
+import (
+	"fmt" // for Printf
+
+	"errors"
+	"io"  // for Reader
+	"log" // for Fatal
+	"math"
+
+	"errors"
+	"fmt" // for Printf
+	"io"  // for Reader
+	"log" // for Fatal
+	"math"
+)

src/cmd/gofmt/testdata/import.input

+package main
+
+import (
+	"fmt"
+	"math"
+	"log"
+	"errors"
+	"io"
+)
+
+import (
+	"fmt"
+
+	"math"
+
+	"log"
+
+	"errors"
+
+	"io"
+)
+
+import (
+	"fmt"
+	"math"
+	"log"
+	"errors"
+	"io"
+
+	"fmt"
+
+	"math"
+
+	"log"
+
+	"errors"
+
+	"io"
+)
+
+import (
+	// a block with comments
+	"fmt" // for Printf
+	"math"
+	"log" // for Fatal
+	"errors"
+	"io" // for Reader
+)
+
+import (
+	"fmt" // for Printf
+
+	"math"
+
+	"log" // for Fatal
+
+	"errors"
+
+	"io" // for Reader
+)
+
+import (
+	// for Printf
+	"fmt"
+
+	"math"
+
+	// for Fatal
+	"log"
+
+	"errors"
+
+	// for Reader
+	"io"
+)
+
+import (
+	"fmt" // for Printf
+	"math"
+	"log" // for Fatal
+	"errors"
+	"io" // for Reader
+
+	"fmt" // for Printf
+
+	"math"
+
+	"log" // for Fatal
+
+	"errors"
+
+	"io" // for Reader
+)
+
+import (
+	"fmt" // for Printf
+
+	"math"
+	"log" // for Fatal
+	"errors"
+	"io" // for Reader
+
+	"fmt" // for Printf
+	"math"
+	"log" // for Fatal
+	"errors"
+	"io" // for Reader
+)

src/pkg/go/ast/Makefile

@@ -8,6 +8,7 @@ TARG=go/ast
 GOFILES=\
 	ast.go\
 	filter.go\
+	import.go\
 	print.go\
 	resolve.go\
 	scope.go\

src/pkg/go/ast/import.go

+// Copyright 2011 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package ast
+
+import (
+	"go/token"
+	"sort"
+	"strconv"
+)
+
+// SortImports sorts runs of consecutive import lines in import blocks in f.
+func SortImports(fset *token.FileSet, f *File) {
+	for _, d := range f.Decls {
+		d, ok := d.(*GenDecl)
+		if !ok || d.Tok != token.IMPORT {
+			// Not an import declaration, so we're done.
+			// Imports are always first.
+			break
+		}
+
+		if d.Lparen == token.NoPos {
+			// Not a block: sorted by default.
+			continue
+		}
+
+		// Identify and sort runs of specs on successive lines.
+		i := 0
+		for j, s := range d.Specs {
+			if j > i && fset.Position(s.Pos()).Line > 1+fset.Position(d.Specs[j-1].End()).Line {
+				// j begins a new run.  End this one.
+				sortSpecs(fset, f, d.Specs[i:j])
+				i = j
+			}
+		}
+		sortSpecs(fset, f, d.Specs[i:])
+	}
+}
+
+func importPath(s Spec) string {
+	t, err := strconv.Unquote(s.(*ImportSpec).Path.Value)
+	if err == nil {
+		return t
+	}
+	return ""
+}
+
+type posSpan struct {
+	Start token.Pos
+	End   token.Pos
+}
+
+func sortSpecs(fset *token.FileSet, f *File, specs []Spec) {
+	// Avoid work if already sorted (also catches < 2 entries).
+	sorted := true
+	for i, s := range specs {
+		if i > 0 && importPath(specs[i-1]) > importPath(s) {
+			sorted = false
+			break
+		}
+	}
+	if sorted {
+		return
+	}
+
+	// Record positions for specs.
+	pos := make([]posSpan, len(specs))
+	for i, s := range specs {
+		pos[i] = posSpan{s.Pos(), s.End()}
+	}
+
+	// Identify comments in this range.
+	// Any comment from pos[0].Start to the final line counts.
+	lastLine := fset.Position(pos[len(pos)-1].End).Line
+	cstart := len(f.Comments)
+	cend := len(f.Comments)
+	for i, g := range f.Comments {
+		if g.Pos() < pos[0].Start {
+			continue
+		}
+		if i < cstart {
+			cstart = i
+		}
+		if fset.Position(g.End()).Line > lastLine {
+			cend = i
+			break
+		}
+	}
+	comments := f.Comments[cstart:cend]
+
+	// Assign each comment to the import spec preceding it.
+	importComment := map[*ImportSpec][]*CommentGroup{}
+	specIndex := 0
+	for _, g := range comments {
+		for specIndex+1 < len(specs) && pos[specIndex+1].Start <= g.Pos() {
+			specIndex++
+		}
+		s := specs[specIndex].(*ImportSpec)
+		importComment[s] = append(importComment[s], g)
+	}
+
+	// Sort the import specs by import path.
+	// Reassign the import paths to have the same position sequence.
+	// Reassign each comment to abut the end of its spec.
+	// Sort the comments by new position.
+	sort.Sort(byImportPath(specs))
+	for i, s := range specs {
+		s := s.(*ImportSpec)
+		s.Path.ValuePos = pos[i].Start
+		s.EndPos = pos[i].End
+		for _, g := range importComment[s] {
+			for _, c := range g.List {
+				c.Slash = pos[i].End
+			}
+		}
+	}
+	sort.Sort(byCommentPos(comments))
+}
+
+type byImportPath []Spec // slice of *ImportSpec
+
+func (x byImportPath) Len() int           { return len(x) }
+func (x byImportPath) Swap(i, j int)      { x[i], x[j] = x[j], x[i] }
+func (x byImportPath) Less(i, j int) bool { return importPath(x[i]) < importPath(x[j]) }
+
+type byCommentPos []*CommentGroup
+
+func (x byCommentPos) Len() int           { return len(x) }
+func (x byCommentPos) Swap(i, j int)      { x[i], x[j] = x[j], x[i] }
+func (x byCommentPos) Less(i, j int) bool { return x[i].Pos() < x[j].Pos() }