parse.go

package parse

import (
	"bufio"
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/scanner"
	"go/token"
	"io"
	"os"
	"strings"
	"unicode"

	"golang.org/x/tools/imports"
)

var header = []byte(`

// This file was automatically generated by genny.
// Any changes will be lost if this file is regenerated.
// see https://github.com/cheekybits/genny

`)

var (
	packageKeyword = []byte("package")
	importKeyword  = []byte("import")
	openBrace      = []byte("(")
	closeBrace     = []byte(")")
	genericPackage = "generic"
	genericType    = "generic.Type"
	genericNumber  = "generic.Number"
	linefeed       = "\r\n"
)
var unwantedLinePrefixes = [][]byte{
	[]byte("//go:generate genny "),
}

func subIntoLiteral(lit, typeTemplate, specificType string) string {
	if lit == typeTemplate {
		return specificType
	}
	if !strings.Contains(lit, typeTemplate) {
		return lit
	}
	specificLg := wordify(specificType, true)
	specificSm := wordify(specificType, false)
	result := strings.Replace(lit, typeTemplate, specificLg, -1)
	if strings.HasPrefix(result, specificLg) && !isExported(lit) {
		return strings.Replace(result, specificLg, specificSm, 1)
	}
	return result
}

func subTypeIntoComment(line, typeTemplate, specificType string) string {
	var subbed string
	for _, w := range strings.Fields(line) {
		subbed = subbed + subIntoLiteral(w, typeTemplate, specificType) + " "
	}
	return subbed
}

// Does the heavy lifting of taking a line of our code and
// sbustituting a type into there for our generic type
func subTypeIntoLine(line, typeTemplate, specificType string) string {
	src := []byte(line)
	var s scanner.Scanner
	fset := token.NewFileSet()
	file := fset.AddFile("", fset.Base(), len(src))
	s.Init(file, src, nil, scanner.ScanComments)
	output := ""
	for {
		_, tok, lit := s.Scan()
		if tok == token.EOF {
			break
		} else if tok == token.COMMENT {
			subbed := subTypeIntoComment(lit, typeTemplate, specificType)
			output = output + subbed + " "
		} else if tok.IsLiteral() {
			subbed := subIntoLiteral(lit, typeTemplate, specificType)
			output = output + subbed + " "
		} else {
			output = output + tok.String() + " "
		}
	}
	return output
}

// typeSet looks like "KeyType: int, ValueType: string"
func generateSpecific(filename string, in io.ReadSeeker, typeSet map[string]string) ([]byte, error) {

	// ensure we are at the beginning of the file
	in.Seek(0, os.SEEK_SET)

	// parse the source file
	fs := token.NewFileSet()
	file, err := parser.ParseFile(fs, filename, in, 0)
	if err != nil {
		return nil, &errSource{Err: err}
	}

	// make sure every generic.Type is represented in the types
	// argument.
	for _, decl := range file.Decls {
		switch it := decl.(type) {
		case *ast.GenDecl:
			for _, spec := range it.Specs {
				ts, ok := spec.(*ast.TypeSpec)
				if !ok {
					continue
				}
				switch tt := ts.Type.(type) {
				case *ast.SelectorExpr:
					if name, ok := tt.X.(*ast.Ident); ok {
						if name.Name == genericPackage {
							if _, ok := typeSet[ts.Name.Name]; !ok {
								return nil, &errMissingSpecificType{GenericType: ts.Name.Name}
							}
						}
					}
				}
			}
		}
	}

	in.Seek(0, os.SEEK_SET)

	var buf bytes.Buffer

	comment := ""
	scanner := bufio.NewScanner(in)
	for scanner.Scan() {

		line := scanner.Text()

		// does this line contain generic.Type?
		if strings.Contains(line, genericType) || strings.Contains(line, genericNumber) {
			comment = ""
			continue
		}

		for t, specificType := range typeSet {
			if strings.Contains(line, t) {
				newLine := subTypeIntoLine(line, t, specificType)
				line = newLine
			}
		}

		if comment != "" {
			buf.WriteString(makeLine(comment))
			comment = ""
		}

		// is this line a comment?
		// TODO: should we handle /* */ comments?
		if strings.HasPrefix(line, "//") {
			// record this line to print later
			comment = line
			continue
		}

		// write the line
		buf.WriteString(makeLine(line))
	}

	// write it out
	return buf.Bytes(), nil
}

// Generics parses the source file and generates the bytes replacing the
// generic types for the keys map with the specific types (its value).
func Generics(filename, pkgName string, in io.ReadSeeker, typeSets []map[string]string) ([]byte, error) {

	totalOutput := header

	for _, typeSet := range typeSets {

		// generate the specifics
		parsed, err := generateSpecific(filename, in, typeSet)
		if err != nil {
			return nil, err
		}

		totalOutput = append(totalOutput, parsed...)

	}

	// clean up the code line by line
	packageFound := false
	insideImportBlock := false
	var cleanOutputLines []string
	scanner := bufio.NewScanner(bytes.NewReader(totalOutput))
	for scanner.Scan() {

		// end of imports block?
		if insideImportBlock {
			if bytes.HasSuffix(scanner.Bytes(), closeBrace) {
				insideImportBlock = false
			}
			continue
		}

		if bytes.HasPrefix(scanner.Bytes(), packageKeyword) {
			if packageFound {
				continue
			} else {
				packageFound = true
			}
		} else if bytes.HasPrefix(scanner.Bytes(), importKeyword) {
			if bytes.HasSuffix(scanner.Bytes(), openBrace) {
				insideImportBlock = true
			}
			continue
		}

		// check all unwantedLinePrefixes - and skip them
		skipline := false
		for _, prefix := range unwantedLinePrefixes {
			if bytes.HasPrefix(scanner.Bytes(), prefix) {
				skipline = true
				continue
			}
		}

		if skipline {
			continue
		}

		cleanOutputLines = append(cleanOutputLines, makeLine(scanner.Text()))
	}

	cleanOutput := strings.Join(cleanOutputLines, "")

	output := []byte(cleanOutput)
	var err error

	// change package name
	if pkgName != "" {
		output = changePackage(bytes.NewReader([]byte(output)), pkgName)
	}
	// fix the imports
	output, err = imports.Process(filename, output, nil)
	if err != nil {
		return nil, &errImports{Err: err}
	}

	return output, nil
}

func makeLine(s string) string {
	return fmt.Sprintln(strings.TrimRight(s, linefeed))
}

// isAlphaNumeric gets whether the rune is alphanumeric or _.
func isAlphaNumeric(r rune) bool {
	return r == '_' || unicode.IsLetter(r) || unicode.IsDigit(r)
}

// wordify turns a type into a nice word for function and type
// names etc.
func wordify(s string, exported bool) string {
	s = strings.TrimRight(s, "{}")
	s = strings.TrimLeft(s, "*&")
	s = strings.Replace(s, ".", "", -1)
	if !exported {
		return s
	}
	return strings.ToUpper(string(s[0])) + s[1:]
}

func changePackage(r io.Reader, pkgName string) []byte {
	var out bytes.Buffer
	sc := bufio.NewScanner(r)
	done := false

	for sc.Scan() {
		s := sc.Text()

		if !done && strings.HasPrefix(s, "package") {
			parts := strings.Split(s, " ")
			parts[1] = pkgName
			s = strings.Join(parts, " ")
			done = true
		}

		fmt.Fprintln(&out, s)
	}
	return out.Bytes()
}

func isExported(lit string) bool {
	if len(lit) == 0 {
		return false
	}
	return unicode.IsUpper(rune(lit[0]))
}