generator.go

package main // import "github.com/azr/generators/recycler"

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

	"github.com/azr/generators/utils"
)

// Usage is a replacement usage function for the flags package.
func Usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\trecycler [flags] -type T [directory]\n")
	fmt.Fprintf(os.Stderr, "\trecycler [flags] -type T files... # Must be a single package\n")
	fmt.Fprintf(os.Stderr, "\trecycler [flags] -template pool.gotpl -type T files... # To generate a basic pool\n")
	fmt.Fprintf(os.Stderr, "\trecycler [flags] -template freelist.gotpl -type T files... # To generate a basic freelist\n")
	fmt.Fprintf(os.Stderr, "\trecycler [flags] -template /path/to/template -type T files... # To generate your own freelist\n")
	fmt.Fprintf(os.Stderr, "For more information, see:\n")
	fmt.Fprintf(os.Stderr, "\thttp://godoc.org/github.com/azr/recycler\n")
	fmt.Fprintf(os.Stderr, "Flags:\n")
	flag.PrintDefaults()
}

func main() {
	var (
		typeNames  = flag.String("type", "", "comma-separated list of type names; must be set")
		output     = flag.String("output", "", "output file name; default srcdir/<type[0]>_recycler.go")
		tpl        = flag.String("template", "pool.gotpl", "go template to generate your recycler with. Defined ones are pool and freelists. Full path also works.\nAvailable template vars:\n\t*Type: type to recycle\n\t*Size: size of the freelist. Not used in pool.")
		size       = flag.Int("size", 50, "Max number of items kept. used for freelist")
		importSync = flag.Bool("sync", false, "Should the generated file import the sync pkg ?")
	)
	log.SetFlags(0)
	log.SetPrefix("recycler: ")
	flag.Usage = Usage
	flag.Parse()
	if len(*typeNames) == 0 {
		flag.Usage()
		os.Exit(2)
	}
	types := strings.Split(*typeNames, ",")

	// We accept either one directory or a list of files. Which do we have?
	args := flag.Args()
	if len(args) == 0 {
		// Default: process whole package in current directory.
		args = []string{"."}
	}
	// copy file utils
	_, currFile, _, ok := runtime.Caller(0)
	if !ok {
		log.Fatalf("No caller information")
	}

	// Parse the package once.
	var (
		dir string
		g   Generator
		err error
	)
	if *tpl == "pool.gotpl" {
		*importSync = true
	}
	templatePath, err := utils.GetExistingPathFor(*tpl, filepath.Dir(currFile))
	if err != nil {
		log.Fatalf("Could not find template: %s", err)
	}
	g.tpl, err = template.ParseFiles(templatePath)
	if err != nil {
		log.Fatalf("Could not parse template: %s", err)
	}

	if len(args) == 1 && isDirectory(args[0]) {
		dir = args[0]
		g.parsePackageDir(args[0])
	} else {
		dir = filepath.Dir(args[0])
		g.parsePackageFiles(args)
	}

	// Print the header and package clause.
	g.Printf("// Code generated by \"github.com/azr/generators/recycler %s\"; DO NOT EDIT\n", strings.Join(os.Args[1:], " "))
	g.Printf("\n")
	g.Printf("package %s", g.pkg.name)
	g.Printf("\n")
	if *importSync {
		g.Printf("import \"sync\"\n")
	}

	// Run generate for each type.
	for _, typeName := range types {
		g.generate(typeName, *size)
	}

	// Format the output.
	src := g.format()

	// Write to file.
	outputName := *output
	if outputName == "" {
		baseName := fmt.Sprintf("%s_recycler.go", types[0])
		outputName = filepath.Join(dir, strings.ToLower(baseName))
	}
	err = ioutil.WriteFile(outputName, src, 0644)
	if err != nil {
		log.Fatalf("writing output: %s", err)
	}
}

// isDirectory reports whether the named file is a directory.
func isDirectory(name string) bool {
	info, err := os.Stat(name)
	if err != nil {
		log.Fatal(err)
	}
	return info.IsDir()
}

// Generator holds the state of the analysis. Primarily used to buffer
// the output for format.Source.
type Generator struct {
	buf bytes.Buffer       // Accumulated output.
	pkg *Package           // Package we are scanning.
	tpl *template.Template // Template used for writing file.
}

func (g *Generator) Printf(format string, args ...interface{}) {
	fmt.Fprintf(&g.buf, format, args...)
}

// File holds a single parsed file and associated data.
type File struct {
	pkg  *Package  // Package to which this file belongs.
	file *ast.File // Parsed AST.
	// These fields are reset for each type being generated.
	typeName string // Name of the type.
	found    bool
}

type Package struct {
	dir      string
	name     string
	defs     map[*ast.Ident]types.Object
	files    []*File
	typesPkg *types.Package
}

// parsePackageDir parses the package residing in the directory.
func (g *Generator) parsePackageDir(directory string) {
	pkg, err := build.Default.ImportDir(directory, 0)
	if err != nil {
		log.Fatalf("cannot process directory %s: %s", directory, err)
	}
	var names []string
	names = append(names, pkg.GoFiles...)
	names = append(names, pkg.CgoFiles...)
	// TODO: Need to think about constants in test files. Maybe write type_string_test.go
	// in a separate pass? For later.
	// names = append(names, pkg.TestGoFiles...) // These are also in the "foo" package.
	names = append(names, pkg.SFiles...)
	names = prefixDirectory(directory, names)
	g.parsePackage(directory, names, nil)
}

// parsePackageFiles parses the package occupying the named files.
func (g *Generator) parsePackageFiles(names []string) {
	g.parsePackage(".", names, nil)
}

// prefixDirectory places the directory name on the beginning of each name in the list.
func prefixDirectory(directory string, names []string) []string {
	if directory == "." {
		return names
	}
	ret := make([]string, len(names))
	for i, name := range names {
		ret[i] = filepath.Join(directory, name)
	}
	return ret
}

// parsePackage analyzes the single package constructed from the named files.
// If text is non-nil, it is a string to be used instead of the content of the file,
// to be used for testing. parsePackage exits if there is an error.
func (g *Generator) parsePackage(directory string, names []string, text interface{}) {
	var files []*File
	var astFiles []*ast.File
	g.pkg = new(Package)
	fs := token.NewFileSet()
	for _, name := range names {
		if !strings.HasSuffix(name, ".go") {
			continue
		}
		parsedFile, err := parser.ParseFile(fs, name, text, 0)
		if err != nil {
			log.Fatalf("parsing package: %s: %s", name, err)
		}
		astFiles = append(astFiles, parsedFile)
		files = append(files, &File{
			file: parsedFile,
			pkg:  g.pkg,
		})
	}
	if len(astFiles) == 0 {
		log.Fatalf("%s: no buildable Go files", directory)
	}
	g.pkg.name = astFiles[0].Name.Name
	g.pkg.files = files
	g.pkg.dir = directory
	// Type check the package.
	g.pkg.check(fs, astFiles)
}

// check type-checks the package. The package must be OK to proceed.
func (pkg *Package) check(fs *token.FileSet, astFiles []*ast.File) {
	pkg.defs = make(map[*ast.Ident]types.Object)
	config := types.Config{
		FakeImportC: true,
		Importer:    importer.Default(),
	}
	info := &types.Info{
		Defs: pkg.defs,
	}
	typesPkg, err := config.Check(pkg.dir, fs, astFiles, info)
	if err != nil {
		log.Fatalf("checking package: %s", err)
	}
	pkg.typesPkg = typesPkg
}

// generate produces the String method for the named type.
func (g *Generator) generate(typeName string, size int) {
	found := false
	for _, file := range g.pkg.files {
		// Set the state for this run of the walker.
		file.typeName = typeName
		if file.file != nil {
			ast.Inspect(file.file, file.genDecl)
			if file.found {
				found = true
			}
		}
	}

	if found {
		g.build(typeName, size)
	} else {
		fmt.Printf("Type not found: %s", typeName)
	}
}

// format returns the gofmt-ed contents of the Generator's buffer.
func (g *Generator) format() []byte {
	src, err := format.Source(g.buf.Bytes())
	if err != nil {
		// Should never happen, but can arise when developing this code.
		// The user can compile the output to see the error.
		log.Printf("warning: internal error: invalid Go generated: %s", err)
		log.Printf("warning: compile the package to analyze the error")
		return g.buf.Bytes()
	}
	return src
}

// genDecl processes one declaration clause.
func (f *File) genDecl(node ast.Node) bool {
	decl, ok := node.(*ast.GenDecl)
	if !ok || decl.Tok != token.TYPE {
		// We only care about type declarations.
		return true
	}

	for _, spec := range decl.Specs {
		vspec := spec.(*ast.TypeSpec) // Guaranteed to succeed as this is Type.
		if vspec.Name.String() != f.typeName {
			continue
		}
		f.found = true
	}
	return false
}

// build generates the variables and String method for a single run of contiguous values.
func (g *Generator) build(typeName string, size int) {
	g.Printf("\n")

	g.tpl.Execute(&g.buf, struct {
		Type string
		Size int
	}{
		Type: typeName,
		Size: size,
	})
}