main.go

/*
Package goreuse provides a tool for reuse Go code, that bundles a whole package inside a single
file. It allows to rename certain identifiers and keep the changed definitions.
Also it supports "go generate" by adding the necessary comments.

Usage

To just bundle another package and write the code to a file:
	$ goreuse -o some-file.go the/package/to/bundle
For adding a certain prefix to all identifiers inside the bundled code (default value is the package name):
	$ goreuse -o some-file.go -prefix prefx_ the/package/to/bundle
To rename some individual identifiers ("-rn" flag can be repeated):
	$ goreuse -rn newName=oldName -o some-file.go -prefix prefx_ the/package/to/bundle
To rename some individual identifiers and preserve the changes that are made to it's definition (note the "+" sign in "newName+=oldName"):
	$ goreuse -rn newName+=oldName -o some-file.go -prefix prefx_ the/package/to/bundle

"goreuse" runs "goreturns" on the output file which can be installed by:
	$ go get -u -v sourcegraph.com/sqs/goreturns

Sample Scenario

Assume we have implemented a template for pipeline pattern (at this step it contains a bug intentially):
	package reusable
	type IN = interface{}
	type OUT = interface{}
	func Pipeline(in <-chan IN) <-chan OUT {
		out := make(chan OUT)
		go func() {
			close(out)
			for v := range in {
				out <- Process(v)
			}
		}()
		return out
	}
	func Process(in IN) OUT {
		var o OUT
		return o
	}
And by using this goreuse command:
	$ goreuse -o pipe.go -rn num+=IN -rn str+=OUT -rn numToStr=Pipeline -rn processNum+=Process sample/reusable
We will have an output file like:
	// Code generated by github.com/dc0d/goreuse. DO NOT EDIT.
	//go:generate goreuse -o pipe.go -rn num+=IN -rn str+=OUT -rn numToStr=Pipeline -rn processNum+=Process sample/reusable
	package actual
	type num = interface{}
	type str = interface{}
	func numToStr(in <-chan num) <-chan str {
		out := make(chan str)
		go func() {
			close(out)
			for v := range in {
				out <- processNum(v)
			}
		}()
		return out
	}
	func processNum(in num) str {
		var o str
		return o
	}
So far just the structure of the two files are the same with some renaming.
Now we can run "go generate" instead.

Next step can be using strict types instead of "interface{}". Let's change "type num = interface{}" to "type num = int".
If we run "go generate" again, "goreuse" will notice we have changed the definition of "num" and will preserve it.

This way we can have "generic code".

Now back to the bug. As you see in the origin code we have "close(out)" instead of "defer close(out)". We just have to fix it inside the original code, then run "go generate" again.

Which part is the template/generic part? It's specified using the "-rn" flag, both the renaming or redefining the identifier.
There is no need for applying any specifc convention. That's why any valid Go package that we use currently, is a valid template/generic code too!
*/
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/printer"
	"go/token"
	"go/types"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"golang.org/x/tools/go/loader"
)

var (
	outputFile = flag.String("o", "", "write output to `file` (default standard output)")
	dstPath    = flag.String("dst", "", "set destination import `path` (default taken from current directory)")
	pkgName    = flag.String("pkg", "", "set destination package `name` (default taken from current directory)")
	prefix     = flag.String("prefix", "", "set bundled identifier prefix to `p` (default source package name + \"_\")")
	underscore = flag.Bool("underscore", false, "rewrite golang.org to golang_org in imports; temporary workaround for golang.org/issue/16333")

	importMap = map[string]string{}
)

func init() {
	flag.Var(flagFunc(addImportMap), "import", "rewrite import using `map`, of form old=new (can be repeated)")
}

func addImportMap(s string) {
	if strings.Count(s, "=") != 1 {
		log.Fatal("-import argument must be of the form old=new")
	}
	i := strings.Index(s, "=")
	old, new := s[:i], s[i+1:]
	if old == "" || new == "" {
		log.Fatal("-import argument must be of the form old=new; old and new must be non-empty")
	}
	importMap[old] = new
}

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: goreuse [options] <src>\n")
	flag.PrintDefaults()
}

func main() {
	log.SetPrefix("goreuse: ")
	log.SetFlags(0)

	flag.Usage = usage
	flag.Parse()
	args := flag.Args()
	if len(args) != 1 {
		usage()
		os.Exit(2)
	}

	if *dstPath != "" {
		if *pkgName == "" {
			*pkgName = path.Base(*dstPath)
		}
	} else {
		wd, _ := os.Getwd()
		pkg, err := build.ImportDir(wd, 0)
		if err != nil {
			log.Fatalf("cannot find package in current directory: %v", err)
		}
		*dstPath = pkg.ImportPath
		if *pkgName == "" {
			*pkgName = pkg.Name
		}
	}

	// changed
	before()

	code, err := bundle(args[0], *dstPath, *pkgName, *prefix)
	if err != nil {
		log.Fatal(err)
	}
	if *outputFile != "" {
		err := ioutil.WriteFile(*outputFile, code, 0666)
		if err != nil {
			log.Fatal(err)
		}

		// changed
		after()
	} else {
		_, err := os.Stdout.Write(code)
		if err != nil {
			log.Fatal(err)
		}
	}
}

// isStandardImportPath is copied from cmd/go in the standard library.
func isStandardImportPath(path string) bool {
	i := strings.Index(path, "/")
	if i < 0 {
		i = len(path)
	}
	elem := path[:i]
	return !strings.Contains(elem, ".")
}

var ctxt = &build.Default

func bundle(src, dst, dstpkg, prefix string) ([]byte, error) {
	// Load the initial package.
	conf := loader.Config{ParserMode: parser.ParseComments, Build: ctxt}
	conf.TypeCheckFuncBodies = func(p string) bool { return p == src }
	conf.Import(src)

	lprog, err := conf.Load()
	if err != nil {
		return nil, err
	}

	// Because there was a single Import call and Load succeeded,
	// InitialPackages is guaranteed to hold the sole requested package.
	info := lprog.InitialPackages()[0]
	if prefix == "" {
		pkgName := info.Files[0].Name.Name
		prefix = pkgName + "_"
	}

	objsToUpdate := make(map[types.Object]bool)
	var rename func(from types.Object)
	rename = func(from types.Object) {
		if !objsToUpdate[from] {
			objsToUpdate[from] = true

			// Renaming a type that is used as an embedded field
			// requires renaming the field too. e.g.
			// 	type T int // if we rename this to U..
			// 	var s struct {T}
			// 	print(s.T) // ...this must change too
			if _, ok := from.(*types.TypeName); ok {
				for id, obj := range info.Uses {
					if obj == from {
						if field := info.Defs[id]; field != nil {
							rename(field)
						}
					}
				}
			}
		}
	}

	// Rename each package-level object.
	scope := info.Pkg.Scope()
	for _, name := range scope.Names() {
		rename(scope.Lookup(name))
	}

	var out bytes.Buffer

	fmt.Fprintf(&out, "// Code generated by github.com/dc0d/goreuse. DO NOT EDIT.\n")
	if *outputFile != "" {
		fmt.Fprintf(&out, "//go:generate goreuse %s\n", strings.Join(os.Args[1:], " "))
	} else {
		fmt.Fprintf(&out, "//   $ goreuse %s\n", strings.Join(os.Args[1:], " "))
	}
	fmt.Fprintf(&out, "\n")

	// Concatenate package comments from all files...
	for _, f := range info.Files {
		if doc := f.Doc.Text(); strings.TrimSpace(doc) != "" {
			for _, line := range strings.Split(doc, "\n") {
				fmt.Fprintf(&out, "// %s\n", line)
			}
		}
	}
	// ...but don't let them become the actual package comment.
	fmt.Fprintln(&out)

	fmt.Fprintf(&out, "package %s\n\n", dstpkg)

	// BUG(adonovan,shurcooL): bundle may generate incorrect code
	// due to shadowing between identifiers and imported package names.
	//
	// The generated code will either fail to compile or
	// (unlikely) compile successfully but have different behavior
	// than the original package. The risk of this happening is higher
	// when the original package has renamed imports (they're typically
	// renamed in order to resolve a shadow inside that particular .go file).

	// TODO(adonovan,shurcooL):
	// - detect shadowing issues, and either return error or resolve them
	// - preserve comments from the original import declarations.

	// pkgStd and pkgExt are sets of printed import specs. This is done
	// to deduplicate instances of the same import name and path.
	var pkgStd = make(map[string]bool)
	var pkgExt = make(map[string]bool)
	for _, f := range info.Files {
		for _, imp := range f.Imports {
			path, err := strconv.Unquote(imp.Path.Value)
			if err != nil {
				log.Fatalf("invalid import path string: %v", err) // Shouldn't happen here since conf.Load succeeded.
			}
			if path == dst {
				continue
			}
			if newPath, ok := importMap[path]; ok {
				path = newPath
			}

			var name string
			if imp.Name != nil {
				name = imp.Name.Name
			}
			spec := fmt.Sprintf("%s %q", name, path)
			if isStandardImportPath(path) {
				pkgStd[spec] = true
			} else {
				if *underscore {
					spec = strings.Replace(spec, "golang.org/", "golang_org/", 1)
				}
				pkgExt[spec] = true
			}
		}
	}

	// Print a single declaration that imports all necessary packages.
	fmt.Fprintln(&out, "import (")
	for p := range pkgStd {
		fmt.Fprintf(&out, "\t%s\n", p)
	}
	if len(pkgExt) > 0 {
		fmt.Fprintln(&out)
	}
	for p := range pkgExt {
		fmt.Fprintf(&out, "\t%s\n", p)
	}
	fmt.Fprint(&out, ")\n\n")

	// Modify and print each file.
	for _, f := range info.Files {
		// Update renamed identifiers.
		for id, obj := range info.Defs {
			if objsToUpdate[obj] {
				// changed
				if nn, ok := oldnewSymbolRename[id.Name]; ok {
					mark(obj, nn.newName, id.Name)
					id.Name = nn.newName
				} else {
					mark(obj, prefix+obj.Name(), id.Name)
					id.Name = prefix + obj.Name()
				}

				// origin
				// id.Name = prefix + obj.Name()
			}
		}
		for id, obj := range info.Uses {
			if objsToUpdate[obj] {
				// changed
				if nn, ok := oldnewSymbolRename[id.Name]; ok {
					mark(obj, nn.newName, id.Name)
					id.Name = nn.newName
				} else {
					mark(obj, prefix+obj.Name(), id.Name)
					id.Name = prefix + obj.Name()
				}
				// origin
				// id.Name = prefix + obj.Name()
			}
		}

		// For each qualified identifier that refers to the
		// destination package, remove the qualifier.
		// The "" strings are removed in postprocessing.
		ast.Inspect(f, func(n ast.Node) bool {
			if sel, ok := n.(*ast.SelectorExpr); ok {
				if id, ok := sel.X.(*ast.Ident); ok {
					if obj, ok := info.Uses[id].(*types.PkgName); ok {
						if obj.Imported().Path() == dst {
							id.Name = "@@@"
						}
					}
				}
			}
			return true
		})

		last := f.Package
		if len(f.Imports) > 0 {
			imp := f.Imports[len(f.Imports)-1]
			last = imp.End()
			if imp.Comment != nil {
				if e := imp.Comment.End(); e > last {
					last = e
				}
			}
		}

		// Pretty-print package-level declarations.
		// but no package or import declarations.
		var buf bytes.Buffer
		for _, decl := range f.Decls {
			if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
				continue
			}

			beg, end := sourceRange(decl)

			printComments(&out, f.Comments, last, beg)

			buf.Reset()
			format.Node(&buf, lprog.Fset, &printer.CommentedNode{Node: decl, Comments: f.Comments})
			// Remove each "" in the output.
			// TODO(adonovan): not hygienic.
			out.Write(bytes.Replace(buf.Bytes(), []byte(""), nil, -1))

			last = printSameLineComment(&out, f.Comments, lprog.Fset, end)

			out.WriteString("\n\n")
		}

		printLastComments(&out, f.Comments, last)
	}

	// Now format the entire thing.
	result, err := format.Source(out.Bytes())
	if err != nil {
		log.Fatalf("formatting failed: %v", err)
	}

	return result, nil
}

// sourceRange returns the [beg, end) interval of source code
// belonging to decl (incl. associated comments).
func sourceRange(decl ast.Decl) (beg, end token.Pos) {
	beg = decl.Pos()
	end = decl.End()

	var doc, com *ast.CommentGroup

	switch d := decl.(type) {
	case *ast.GenDecl:
		doc = d.Doc
		if len(d.Specs) > 0 {
			switch spec := d.Specs[len(d.Specs)-1].(type) {
			case *ast.ValueSpec:
				com = spec.Comment
			case *ast.TypeSpec:
				com = spec.Comment
			}
		}
	case *ast.FuncDecl:
		doc = d.Doc
	}

	if doc != nil {
		beg = doc.Pos()
	}
	if com != nil && com.End() > end {
		end = com.End()
	}

	return beg, end
}

func printComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos, end token.Pos) {
	for _, cg := range comments {
		if pos <= cg.Pos() && cg.Pos() < end {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			fmt.Fprintln(out)
		}
	}
}

const infinity = 1 << 30

func printLastComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos token.Pos) {
	printComments(out, comments, pos, infinity)
}

func printSameLineComment(out *bytes.Buffer, comments []*ast.CommentGroup, fset *token.FileSet, pos token.Pos) token.Pos {
	tf := fset.File(pos)
	for _, cg := range comments {
		if pos <= cg.Pos() && tf.Line(cg.Pos()) == tf.Line(pos) {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			return cg.End()
		}
	}
	return pos
}

type flagFunc func(string)

func (f flagFunc) Set(s string) error {
	f(s)
	return nil
}

func (f flagFunc) String() string { return "" }

//-----------------------------------------------------------------------------

func before() {
	if len(oldnewSymbolRename) == 0 {
		return
	}
	if info, err := os.Stat(*outputFile); err != nil || info.IsDir() {
		return
	}
	fset := token.NewFileSet()
	fast, err := parser.ParseFile(
		fset,
		*outputFile,
		nil,
		parser.AllErrors)
	if err != nil {
		log.Fatal(err)
	}
	content, err := ioutil.ReadFile(*outputFile)
	if err != nil {
		log.Fatal(err)
	}
	_before.content = content
	_before.set = fset
	_before.ast = fast
}

func after() {
	if info, err := os.Stat(*outputFile); err != nil || info.IsDir() {
		return
	}
	defer func() {
		cmd := exec.Command("goreturns", "-w", *outputFile)
		if cmd == nil {
			return
		}
		if err := cmd.Start(); err != nil {
			return
		}
		cmd.Wait()
	}()
	fset := token.NewFileSet()
	fast, err := parser.ParseFile(
		fset,
		*outputFile,
		nil,
		parser.AllErrors)
	if err != nil {
		log.Fatal(err)
	}
	content, err := ioutil.ReadFile(*outputFile)
	if err != nil {
		log.Fatal(err)
	}
	_after.content = content
	_after.set = fset
	_after.ast = fast

	filteredBefore := pickSymbols(_before.set, _before.ast, true)
	filteredAfter := pickSymbols(_after.set, _after.ast, true)

	var defs sortableDefs
	for kb, vb := range filteredBefore {
		va, ok := filteredAfter[kb]
		if !ok {
			continue
		}
		if !vb.rename.preserve {
			continue
		}
		epb := vb.endposer()
		epa := va.endposer()
		if epa == nil || epb == nil {
			continue
		}
		var d definf
		sa, ea := offsets(_after.set, epa)
		d.afterPos = [2]int{sa, ea}
		sb, eb := offsets(_before.set, epb)
		d.beforePos = [2]int{sb, eb}
		d.beforeContent = _before.content[sb:eb]
		defs = append(defs, d)
	}
	if len(defs) == 0 {
		return
	}
	sort.Sort(defs)

	lastPos := 0
	finaldst := &bytes.Buffer{}
	for _, vk := range defs {
		finaldst.Write(_after.content[lastPos:vk.afterPos[0]])
		finaldst.Write(vk.beforeContent)
		lastPos = vk.afterPos[1]
	}
	if lastPos > 0 {
		finaldst.Write(_after.content[lastPos:])
	}
	if err := ioutil.WriteFile(*outputFile, finaldst.Bytes(), 777); err != nil {
		log.Fatal(err)
	}
}

type definf struct {
	beforePos     [2]int
	beforeContent []byte
	afterPos      [2]int
}

type sortableDefs []definf

func (e sortableDefs) Len() int           { return len(e) }
func (e sortableDefs) Swap(i, j int)      { e[i], e[j] = e[j], e[i] }
func (e sortableDefs) Less(i, j int) bool { return e[i].beforePos[0] < e[j].beforePos[0] }

func offsets(set *token.FileSet, v endposer) (start, end int) {
	start = int(set.Position(v.Pos()).Offset)
	end = int(set.Position(v.End()).Offset)
	return
}

func pickSymbols(
	fset *token.FileSet,
	fast *ast.File,
	pickNew ...bool) map[string]*pickedNode {
	filtered := make(map[string]*pickedNode)
	if fset == nil || fast == nil {
		return filtered
	}
	var _pickNew bool
	if len(pickNew) > 0 && pickNew[0] {
		_pickNew = true
	}
	pickName := func(name string) (rename, string, bool) {
		var (
			rn rename
			ok bool
		)
		var dic map[string]rename
		if _pickNew {
			dic = newoldSymbolRename
		} else {
			dic = oldnewSymbolRename
		}
		rn, ok = dic[name]
		return rn, name, ok
	}
	for _, vdecl := range fast.Decls {
		switch x := vdecl.(type) {
		case *ast.FuncDecl:
			rn, name, ok := pickName(x.Name.Name)
			if ok {
				filtered[name] = &pickedNode{
					rename:   rn,
					funcDecl: x,
				}
			}
		case *ast.GenDecl:
			for _, vspec := range x.Specs {
				switch xspec := vspec.(type) {
				case *ast.TypeSpec:
					rn, name, ok := pickName(xspec.Name.Name)
					if ok {
						filtered[name] = &pickedNode{
							rename:   rn,
							typeSpec: xspec,
						}
					}
				case *ast.ValueSpec:
					for kval, vval := range xspec.Names {
						rn, name, ok := pickName(vval.Name)
						if ok {
							v := pickedNode{
								rename: rn,
							}
							v.valueSpec = xspec
							v.valueIndex = kval
							filtered[name] = &v
						}
					}
				default:
				}
			}
		default:
		}
	}
	return filtered
}

type pickedNode struct {
	rename     rename
	funcDecl   *ast.FuncDecl
	typeSpec   *ast.TypeSpec
	valueSpec  *ast.ValueSpec
	valueIndex int
}

func (pn *pickedNode) endposer() endposer {
	var item endposer
	switch {
	case pn.funcDecl != nil:
		item = pn.funcDecl
	case pn.typeSpec != nil:
		item = pn.typeSpec
	case pn.valueSpec != nil:
		item = pn.valueSpec.Values[pn.valueIndex]
	}
	return item
}

type endposer interface {
	Pos() token.Pos
	End() token.Pos
}

func mark(obj types.Object, n, o string) {
	if _, ok := oldnew[obj]; ok {
		return
	}
	oldnew[obj] = [2]string{n, o}
}

var (
	oldnew  = map[types.Object][2]string{}
	_before struct {
		content []byte
		set     *token.FileSet
		ast     *ast.File
	}
	_after struct {
		content []byte
		set     *token.FileSet
		ast     *ast.File
	}
)

//-----------------------------------------------------------------------------

var (
	oldnewSymbolRename = map[string]rename{}
	newoldSymbolRename = map[string]rename{}
)

type rename struct {
	newName, oldName string
	preserve         bool
}

func (rn rename) String() string {
	if len(rn.newName) == 0 ||
		len(rn.oldName) == 0 {
		return "N/A"
	}
	var ps string
	if rn.preserve {
		ps = "+"
	}
	return rn.newName + ps + "=" + rn.oldName
}

func addSymbolRename(s string) {
	var rename rename
	rgx := regexp.MustCompile("^(?P<new>\\w+)(?P<preserve>\\+)*=(?P<old>\\w+)$")
	nl := rgx.SubexpNames()
	rr := rgx.FindAllStringSubmatch(s, -1)
	if len(rr) == 0 {
		log.Fatalf("invalid pair: %v", s)
	}
	parts := make(map[string]string)
	for k, v := range rr[0] {
		parts[nl[k]] = v
	}
	rename.newName = parts["new"]
	rename.oldName = parts["old"]
	if parts["preserve"] == "+" {
		rename.preserve = true
	}
	oldnewSymbolRename[rename.oldName] = rename
	newoldSymbolRename[rename.newName] = rename
}

func init() {
	flag.Var(flagFunc(addSymbolRename), "rn", "needs -o, rename identifier new=old or new+=old to preserve the modifications (can be repeated)")
}

//-----------------------------------------------------------------------------