curry.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"log"
	"path/filepath"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
)

const (
	fileHeader = "// Code generated by \"curry\" " +
		"(http://github.com/maxsz/curry) - DO NOT EDIT\n\n"
)

type file struct {
	Name      string
	Directory string
}

func (f file) Path() string {
	return filepath.Join(f.Directory, f.Name)
}

func main() {
	var (
		fileSuffix       string
		functionModifier string
	)
	flag.StringVar(&fileSuffix, "file-suffix", "_curry.go",
		"Curried file suffix.",
	)
	flag.StringVar(&functionModifier, "function-modifier", "C",
		"The string appended to the curried function name.",
	)
	flag.Parse()
	args := flag.Args()
	if len(args) == 0 {
		args = []string{"."}
	}

	files, err := parsePackageDir(args[0], fileSuffix)
	if err != nil {
		log.Fatal("Failed to parse package dir:", err)
	}
	fs := token.NewFileSet()
	for _, f := range files {
		src, err := curryFile(fs, f, functionModifier)
		if err != nil {
			log.Fatal(err)
		}
		src, err = removeUnusedImports(fs, src)
		if err != nil {
			log.Fatal(err)
		}
		curriedFile := file{
			Name:      strings.Replace(f.Name, ".go", fileSuffix, -1),
			Directory: f.Directory,
		}
		ioutil.WriteFile(curriedFile.Path(), src, 0644)
	}
}

// parsePackageDir checks the `directory` and collects all relevant go files
// that do not have the `fileSuffix`.
func parsePackageDir(directory, fileSuffix string) ([]file, error) {
	pkg, err := build.Default.ImportDir(directory, 0)
	if err != nil {
		return nil, err
	}

	var files []file
	addFiles := func(x []string) {
		for _, filename := range x {
			if !strings.HasSuffix(filename, ".go") ||
				strings.HasSuffix(filename, fileSuffix) {
				continue
			}
			files = append(files, file{Name: filename, Directory: directory})
		}
	}
	addFiles(pkg.GoFiles)
	addFiles(pkg.CgoFiles)
	addFiles(pkg.SFiles)

	return files, nil
}

// removeUnusedImports removes all unused imports from `src`.
func removeUnusedImports(fs *token.FileSet, src []byte) ([]byte, error) {
	f, err := parser.ParseFile(fs, "", src, 0)
	if err != nil {
		return nil, err
	}
	for _, imp := range f.Imports {
		path := strings.Trim(imp.Path.Value, "\"")
		if !astutil.UsesImport(f, path) {
			astutil.DeleteImport(fs, f, path)
		}
	}

	var srcOut bytes.Buffer
	err = format.Node(&srcOut, fs, f)
	if err != nil {
		return nil, err
	}

	return srcOut.Bytes(), nil
}

// curryFile analyzes the given `file` from fileset `fs` and returns a string
// containing all curried function versions. The `modifier` is appended to the
// function names, so `func bla(arg1, arg2 string)` becomes
// `func blaC(arg1 string)... `if modifier is "C".
func curryFile(fs *token.FileSet, file file, modifier string) ([]byte, error) {
	f, err := parser.ParseFile(fs, file.Path(), nil, 0)
	if err != nil {
		return nil, err
	}
	info := types.Info{
		Types: make(map[ast.Expr]types.TypeAndValue),
		Defs:  make(map[*ast.Ident]types.Object),
		Uses:  make(map[*ast.Ident]types.Object),
	}
	conf := types.Config{Importer: importer.Default()}
	conf.Check(file.Directory, fs, []*ast.File{f}, &info)

	var src bytes.Buffer
	fmt.Fprint(&src, fileHeader)
	fmt.Fprintf(&src, "package %s\n\n", f.Name.String())

	for _, imp := range f.Imports {
		fmt.Fprintf(&src, "import %s\n", imp.Path.Value)
	}

	nameMod := strings.TrimSpace(modifier)
	if len(modifier) == 0 {
		nameMod = "C"
	}
	ast.Inspect(f, func(n ast.Node) bool {
		switch n := n.(type) {
		case *ast.FuncDecl:
			curryFunction(&src, n.Name.String(), nameMod, fs,
				info.Defs[n.Name].Type().(*types.Signature),
			)
		}
		return true
	})
	return format.Source(src.Bytes())
}

// curryFunction writes all function versions (one for each argument) to `src`.
func curryFunction(src io.Writer, name, modifier string,
	fs *token.FileSet, funcSig *types.Signature,
) error {
	if funcSig.Params().Len() < 2 {
		return nil
	}
	typeName := func(v *types.Var) string {
		var buf bytes.Buffer
		types.WriteType(&buf, v.Type(), func(p *types.Package) string {
			if p.Name() != "main" {
				return p.Name()
			}
			return ""
		})
		return buf.String()
	}
	params := funcSig.Params()
	for i := 0; i < params.Len(); i++ {
		param := params.At(i)
		if i == 0 {
			if funcSig.Recv() != nil {
				recv := funcSig.Recv()
				fmt.Fprintf(src, "func (%s %s) %s%s",
					recv.Name(),
					typeName(recv),
					name,
					modifier,
				)
			} else {
				fmt.Fprintf(src, "func %s%s", name, modifier)
			}
		} else {
			fmt.Fprintf(src, "return func ")
		}
		fmt.Fprintf(src, "(%s %s) ", param.Name(), typeName(param))
		for j := i + 1; j < params.Len(); j++ {
			remaining := params.At(j)
			fmt.Fprintf(src, "func (%s) ", typeName(remaining))
		}
		fmt.Fprintf(src, "%s{\n", funcSig.Results().String())
	}
	receiver := ""
	if funcSig.Recv() != nil {
		receiver = fmt.Sprintf("%s.", funcSig.Recv().Name())
	}
	fmt.Fprintf(src, "return %s%s(", receiver, name)
	for i := 0; i < params.Len(); i++ {
		param := params.At(i)
		fmt.Fprintf(src, "%s", param.Name())
		if i < params.Len()-1 {
			fmt.Fprint(src, ", ")
		} else {
			fmt.Fprint(src, ")\n")
		}
	}
	for i := 0; i < params.Len(); i++ {
		fmt.Fprintf(src, "}\n")
	}
	return nil
}