Merge pull request #38 from Stebalien/feat/speed

speedup gx-go rw/uw

main.go

@@ -1017,23 +1017,32 @@ func addRewriteForDep(dep *gx.Dependency, pkg *Package, m map[string]string, und
 }
 
 func buildRewriteMapping(pkg *Package, pkgdir string, m map[string]string, undo bool) error {
-	for _, dep := range pkg.Dependencies {
-		cpkg, err := loadDep(dep, pkgdir)
-		if err != nil {
-			VLog("error loading dep %q of %q: %s", dep.Name, pkg.Name, err)
-			return fmt.Errorf("package %q not found. (dependency of %s)", dep.Name, pkg.Name)
-		}
+	seen := make(map[string]struct{})
+	var process func(pkg *Package) error
+	process = func(pkg *Package) error {
+		for _, dep := range pkg.Dependencies {
+			if _, ok := seen[dep.Hash]; ok {
+				continue
+			}
+			seen[dep.Hash] = struct{}{}
 
-		addRewriteForDep(dep, cpkg, m, undo)
+			cpkg, err := loadDep(dep, pkgdir)
+			if err != nil {
+				VLog("error loading dep %q of %q: %s", dep.Name, pkg.Name, err)
+				return fmt.Errorf("package %q not found. (dependency of %s)", dep.Name, pkg.Name)
+			}
 
-		// recurse!
-		err = buildRewriteMapping(cpkg, pkgdir, m, undo)
-		if err != nil {
-			return err
+			addRewriteForDep(dep, cpkg, m, undo)
+
+			// recurse!
+			err = process(cpkg)
+			if err != nil {
+				return err
+			}
 		}
+		return nil
 	}
-
-	return nil
+	return process(pkg)
 }
 
 func buildMap(pkg *Package, m map[string]string) error {

rewrite/rewrite.go

@@ -8,32 +8,49 @@ import (
 	"go/parser"
 	"go/printer"
 	"go/token"
+	"io"
 	"os"
 	"path/filepath"
 	"regexp"
+	"runtime"
 	"strconv"
 	"strings"
 	"sync"
 
 	fs "github.com/kr/fs"
 )
 
-var bufpool *sync.Pool
-
-func init() {
-	bufpool = &sync.Pool{
-		New: func() interface{} {
-			return new(bytes.Buffer)
-		},
-	}
+var bufpool = &sync.Pool{
+	New: func() interface{} {
+		return new(bytes.Buffer)
+	},
 }
 
+var cfg = &printer.Config{Mode: printer.UseSpaces | printer.TabIndent, Tabwidth: 8}
+
 func RewriteImports(ipath string, rw func(string) string, filter func(string) bool) error {
 	path, err := filepath.EvalSymlinks(ipath)
 	if err != nil {
 		return err
 	}
 
+	var rwLock sync.Mutex
+
+	var wg sync.WaitGroup
+	torewrite := make(chan string)
+	for i := 0; i < runtime.NumCPU(); i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for path := range torewrite {
+				err := rewriteImportsInFile(path, rw, &rwLock)
+				if err != nil {
+					fmt.Println("rewrite error: ", err)
+				}
+			}
+		}()
+	}
+
 	w := fs.Walk(path)
 	for w.Step() {
 		rel := w.Path()[len(path):]
@@ -54,28 +71,33 @@ func RewriteImports(ipath string, rw func(string) string, filter func(string) bo
 		if !filter(rel) {
 			continue
 		}
-
-		err := rewriteImportsInFile(w.Path(), rw)
-		if err != nil {
-			fmt.Println("rewrite error: ", err)
-		}
+		torewrite <- w.Path()
 	}
+	close(torewrite)
+	wg.Wait()
 	return nil
 }
 
 // inspired by godeps rewrite, rewrites import paths with gx vendored names
-func rewriteImportsInFile(fi string, rw func(string) string) error {
-	cfg := &printer.Config{Mode: printer.UseSpaces | printer.TabIndent, Tabwidth: 8}
+func rewriteImportsInFile(fi string, rw func(string) string, rwLock *sync.Mutex) error {
+	// 1. Rewrite the imports (if we have any)
 	fset := token.NewFileSet()
-	file, err := parser.ParseFile(fset, fi, nil, parser.ParseComments)
+	file, err := parser.ParseFile(fset, fi, nil, parser.ParseComments|parser.ImportsOnly)
 	if err != nil {
 		return err
 	}
+	if len(file.Imports) == 0 {
+		return nil
+	}
+
+	oldImportsEnd := fset.Position(file.Imports[len(file.Imports)-1].End()).Offset
 
+	rwLock.Lock()
 	var changed bool
 	for _, imp := range file.Imports {
 		p, err := strconv.Unquote(imp.Path.Value)
 		if err != nil {
+			rwLock.Unlock()
 			return err
 		}
 
@@ -86,42 +108,89 @@ func rewriteImportsInFile(fi string, rw func(string) string) error {
 			imp.Path.Value = strconv.Quote(np)
 		}
 	}
+	rwLock.Unlock()
 
 	if !changed {
 		return nil
 	}
 
 	buf := bufpool.Get().(*bytes.Buffer)
+	defer func() {
+		bufpool.Put(buf)
+	}()
+
+	// Write them back to a temporary buffer
+
+	buf.Reset()
 	if err = cfg.Fprint(buf, fset, file); err != nil {
 		return err
 	}
 
+	// 2. Read the imports back in to sort them.
+
 	fset = token.NewFileSet()
-	file, err = parser.ParseFile(fset, fi, buf, parser.ParseComments)
+	file, err = parser.ParseFile(fset, fi, buf, parser.ParseComments|parser.ImportsOnly)
 	if err != nil {
 		return err
 	}
 
+	ast.SortImports(fset, file)
+
+	// Write them back to a temporary buffer
+
 	buf.Reset()
-	bufpool.Put(buf)
+	if err = cfg.Fprint(buf, fset, file); err != nil {
+		return err
+	}
 
-	ast.SortImports(fset, file)
+	// 3. Read them back in to find the new end of the imports.
 
-	wpath := fi + ".temp"
-	w, err := os.Create(wpath)
+	fset = token.NewFileSet()
+	file, err = parser.ParseFile(fset, fi, buf, parser.ParseComments|parser.ImportsOnly)
+	if err != nil {
+		return err
+	}
+
+	newImportsEnd := fset.Position(file.Imports[len(file.Imports)-1].End()).Offset
+
+	// Write them back to the buffer and truncate.
+	buf.Reset()
+	if err = cfg.Fprint(buf, fset, file); err != nil {
+		return err
+	}
+	buf.Truncate(newImportsEnd)
+
+	// Finally, build the file.
+
+	tmppath := fi + ".temp"
+	tmp, err := os.Create(tmppath)
+	if err != nil {
+		return err
+	}
+
+	src, err := os.Open(fi)
 	if err != nil {
 		return err
 	}
+	defer src.Close()
 
-	if err = cfg.Fprint(w, fset, file); err != nil {
+	_, err = src.Seek(int64(oldImportsEnd), io.SeekStart)
+	if err != nil {
+		return err
+	}
+
+	buf.WriteTo(tmp)
+
+	_, err = io.Copy(tmp, src)
+	if err != nil {
 		return err
 	}
 
-	if err = w.Close(); err != nil {
+	if err = tmp.Close(); err != nil {
 		return err
 	}
 
-	return os.Rename(wpath, fi)
+	return os.Rename(tmppath, fi)
 }
 
 func fixCanonicalImports(buf []byte) (bool, error) {