generate_types.go

package path

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	"github.com/iheartradio/cog/cfs"

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

const genHead = `// DO NOT EDIT.
//
// File autogenerated by cog-path.

package %s

//gocovr:skip-file

import "github.com/iheartradio/cog/encoding/path"

`

type genProgram struct {
	*loader.Program
	genVisitor
	genDir   bool
	pkg      *loader.PackageInfo
	pkgPath  string // For the package being generated
	absPath  string // Abs path of fileOrDir
	types    map[string]*genType
	imported map[string]bool
}

type genType struct {
	types.Type
	file       string
	localName  string
	fullName   string
	generating bool // If this type is being generated on this run
}

type genVisitor interface {
	io.Writer

	funcOpen(name string)
	funcClose()

	// Handle a static path
	staticPath(tag string)

	// Function name for this thing
	marshalerName() string

	// Emit a call to an un/marshaler
	marshaler(name string)

	// Emit the proper array var
	arrayVar(name string, t types.Type) string

	// This is a byte array: use the fast version
	byteArray(name string)

	// A basic value to handle
	basic(name, kind string)
}

type genMarshal struct {
	*bytes.Buffer
}

type genUnmarshal struct {
	*bytes.Buffer
	vars *bytes.Buffer
	name string
	tagI int
}

const (
	defaultName = "v"
	fileSuffix  = "_cog_path"
)

var marshalerType, unmarshalerType *types.Interface

func newProgram(fileOrDir string) (*genProgram, error) {
	stat, err := os.Stat(fileOrDir)
	if err != nil {
		return nil, err
	}

	absPath := fileOrDir
	if !filepath.IsAbs(absPath) {
		absPath, err = filepath.Abs(fileOrDir)
		if err != nil {
			return nil, err
		}
	}

	conf := loader.Config{
		ParserMode:          parser.AllErrors,
		TypeCheckFuncBodies: func(string) bool { return false },
	}

	pkgPath, err := cfs.ImportPath(absPath, stat.IsDir())
	if err != nil {
		return nil, err
	}

	if stat.IsDir() {
		fs, err := ioutil.ReadDir(fileOrDir)
		if err != nil {
			return nil, err
		}

		for _, f := range fs {
			os.Remove(genFileName(filepath.Join(fileOrDir, f.Name())))
		}

		conf.Import(pkgPath)
	} else {
		os.Remove(genFileName(fileOrDir))

		fileOrDir = filepath.Base(fileOrDir)
		if strings.HasSuffix(fileOrDir, "_test.go") {
			conf.ImportWithTests(pkgPath)
		} else {
			conf.Import(pkgPath)
		}
	}

	conf.Import("github.com/iheartradio/cog/encoding/path")

	if err != nil {
		return nil, err
	}

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

	p := genProgram{
		Program:  prog,
		genDir:   stat.IsDir(),
		pkg:      prog.Package(pkgPath),
		pkgPath:  pkgPath,
		absPath:  absPath,
		types:    map[string]*genType{},
		imported: map[string]bool{},
	}

	return &p, nil
}

func genFileName(file string) (out string) {
	outFile := cfs.DropExt(file)
	if strings.HasSuffix(outFile, "_test") {
		out = outFile[:len(outFile)-5] + fileSuffix + "_test.go"
	} else {
		out = outFile + fileSuffix + ".go"
	}

	return
}

// GenerateFrom generates MarshalPath and UnmarshalPath methods for all types
// found in the given file or directory. Marshal-implementation files are named
// as "$BASE_FILE_cog_path.go".
func GenerateFrom(fileOrDir string) error {
	p, err := newProgram(fileOrDir)
	if err != nil {
		return err
	}

	p.importTypes(p.pkgPath)

	var files []string

	if p.genDir { // Generating for entire directory
		for _, t := range p.types {
			t.generating = true
		}

		fs, err := ioutil.ReadDir(fileOrDir)
		if err != nil {
			return err
		}

		for _, f := range fs {
			if !f.IsDir() {
				files = append(files, filepath.Join(p.absPath, f.Name()))
			}
		}

	} else { // Generating for 1 file
		files = append(files, p.absPath)
	}

	for _, f := range files {
		var types []*genType

		for _, t := range p.types {
			if t.file == f {
				t.generating = true
				types = append(types, t)
			}
		}

		err = p.generate(f, types)
		if err != nil {
			break
		}
	}

	return err
}

func (p *genProgram) generate(file string, ts []*genType) error {
	outPath := genFileName(file)
	of, err := os.Create(outPath)
	if err == nil {
		defer of.Close()
	}

	vars := &bytes.Buffer{}
	cont := &bytes.Buffer{}
	if err == nil {
		fmt.Fprintf(cont, genHead, p.pkg.Pkg.Name())

		for _, t := range ts {
			p.genVisitor = genMarshal{cont}
			err = p.visitAll(defaultName, t)

			if err == nil {
				p.genVisitor = &genUnmarshal{
					Buffer: cont,
					vars:   vars,
					name:   strings.Title(t.localName),
				}

				err = p.visitAll(defaultName, t)
			}

			if err != nil {
				break
			}
		}
	}

	if err == nil && vars.Len() > 0 {
		fmt.Fprintf(cont, "var(\n%s\n)", vars.Bytes())
	}

	if err == nil {
		var out []byte
		out, err = imports.Process(outPath, cont.Bytes(), nil)
		if err == nil {
			_, err = of.Write(out)
		}
	}

	return err
}

func (p *genProgram) visitAll(name string, t *genType) (err error) {
	p.funcOpen(t.localName)
	err = p.visit(name, t.Type)
	p.funcClose()

	return
}

func (p *genProgram) visit(name string, t types.Type) error {
	switch t := t.(type) {
	case *types.Struct:
		return p.visitStruct(name, t)

	case *types.Array:
		return p.visitArray(name, t)

	case *types.Basic:
		return p.visitBasic(name, t)

	case *types.Named:
		return p.visitNamed(name, t)

	default:
		return fmt.Errorf("unrecognized type: %v", reflect.TypeOf(t))
	}
}

func (p *genProgram) visitStruct(name string, t *types.Struct) (err error) {
	n := t.NumFields()

	for i := 0; err == nil && i < n; i++ {
		f := t.Field(i)

		if f.Type().String() == "github.com/iheartradio/cog/encoding/path.Static" {
			path := reflect.StructTag(t.Tag(i)).Get("path")
			p.staticPath(path)
			continue
		}

		if !f.Exported() {
			continue
		}

		name := fmt.Sprintf("%s.%s", name, f.Name())
		err = p.visit(name, f.Type())
	}

	return
}

func (p *genProgram) visitNamed(name string, t *types.Named) (err error) {
	gt := p.getType(t)
	if gt.generating {
		p.marshaler(name)
		return
	}

	fnName := p.marshalerName()
	n := t.NumMethods()
	for i := 0; err == nil && i < n; i++ {
		if t.Method(i).Name() == fnName {
			p.marshaler(name)
			return
		}
	}

	// If it doesn't have a marshal function, just emit inline
	p.visit(name, t.Underlying())

	return
}

func (p *genProgram) visitArray(name string, t *types.Array) (err error) {
	et := t.Elem()

	if b, ok := et.(*types.Basic); ok && (b.Kind() == types.Int8 || b.Kind() == types.Uint8) {
		p.byteArray(name)
	} else {
		fmt.Fprintf(p,
			"	for i := 0; s.Err == nil && i < %d; i++ {\n",
			t.Len())

		name = p.arrayVar(name, et)
		err = p.visit(name, et)
		fmt.Fprintf(p, "	}\n")
	}

	return
}

func (p *genProgram) visitBasic(name string, t *types.Basic) (err error) {
	kind := ""
	switch t.Kind() {
	case types.Bool:
		kind = "bool"
	case types.Int8:
		kind = "int8"
	case types.Int16:
		kind = "int16"
	case types.Int32:
		kind = "int32"
	case types.Int64:
		kind = "int64"
	case types.Uint8:
		kind = "uint8"
	case types.Uint16:
		kind = "uint16"
	case types.Uint32:
		kind = "uint32"
	case types.Uint64:
		kind = "uint64"
	case types.Float32:
		kind = "float32"
	case types.Float64:
		kind = "float64"
	case types.Complex64:
		kind = "complex64"
	case types.Complex128:
		kind = "complex128"
	case types.String:
		kind = "string"
	default:
		err = fmt.Errorf("unrecognized basic kind: %v", t.Kind())
	}

	if err == nil {
		p.basic(name, kind)
	}

	return
}

func (p *genProgram) getType(t *types.Named) *genType {
	name := t.String()
	p.importTypes(cfs.DropExt(name))
	return p.types[name]
}

func (p *genProgram) importTypes(pkgPath string) {
	if p.imported[pkgPath] {
		return
	}

	pkg := p.Package(pkgPath)

	for _, astf := range pkg.Files {
		file := p.Fset.File(astf.Pos()).Name()
		p.importTypesFromAst(pkg, file, astf)
	}

	p.imported[pkgPath] = true
}

func (p *genProgram) importTypesFromAst(
	pkg *loader.PackageInfo,
	file string,
	astf *ast.File) {

	for _, decl := range astf.Decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok {
			continue
		}

		for _, spec := range genDecl.Specs {
			typeSpec, ok := spec.(*ast.TypeSpec)
			if !ok {
				continue
			}

			switch typeSpec.Type.(type) {
			case *ast.StructType, *ast.ArrayType, *ast.Ident, *ast.SelectorExpr:
				name := fmt.Sprintf("%s.%s",
					pkg.Pkg.Path(),
					typeSpec.Name.Name)

				p.types[name] = &genType{
					Type:      pkg.Types[typeSpec.Type].Type,
					file:      file,
					localName: typeSpec.Name.Name,
					fullName:  name,
				}
			}
		}
	}
}

func (g genMarshal) funcOpen(name string) {
	fmt.Fprintf(g,
		"func (v %s) MarshalPath(s path.Encoder) path.Encoder {\n",
		name)
}

func (g genMarshal) funcClose() {
	g.WriteString("	return s\n}\n\n")
}

func (g genMarshal) staticPath(tag string) {
	fmt.Fprintf(g,
		"	s.B = append(s.B, \"%s\"...)\n"+
			"	s = s.EmitSep()\n",
		tag)
}

func (g genMarshal) marshalerName() string {
	return "MarshalPath"
}

func (g genMarshal) marshaler(name string) {
	fmt.Fprintf(g, "	s = %s.MarshalPath(s)\n", name)
}

func (g genMarshal) arrayVar(name string, t types.Type) string {
	fmt.Fprintf(g,
		"		av := %s[i]\n",
		name)
	return "av"
}

func (g genMarshal) byteArray(name string) {
	fmt.Fprintf(g,
		"	s.B = append(s.B, %s[:]...)\n"+
			"	s = s.EmitSep()\n",
		name)
}

func (g genMarshal) basic(name, kind string) {
	fmt.Fprintf(g, "	s = s.Emit%s(%s(%s))\n",
		strings.Title(kind),
		kind,
		name)
}

func (g *genUnmarshal) funcOpen(name string) {
	fmt.Fprintf(g,
		"func (v *%s) UnmarshalPath(s path.Decoder) path.Decoder {\n",
		name)
}

func (g *genUnmarshal) funcClose() {
	g.WriteString("	return s\n}\n\n")
}

func (g *genUnmarshal) staticPath(tag string) {
	tagVarName := fmt.Sprintf("cogTag%s%d", g.name, g.tagI)
	g.tagI++

	g.guard("s = s.ExpectTagBytes(%s)", tagVarName)

	fmt.Fprintf(g.vars, "	%s = []byte(\"%s\")\n",
		tagVarName,
		tag)
}

func (g *genUnmarshal) marshalerName() string {
	return "UnmarshalPath"
}

func (g *genUnmarshal) marshaler(name string) {
	g.guard("s = %s.UnmarshalPath(s)", name)
}

func (g *genUnmarshal) arrayVar(name string, t types.Type) string {
	if _, ok := t.(*types.Basic); ok {
		name = fmt.Sprintf("%s[i]", name)
	} else {
		fmt.Fprintf(g,
			"		av := &%s[i]\n",
			name)
		name = "av"
	}

	return name
}

func (g *genUnmarshal) byteArray(name string) {
	g.guard("s = s.ExpectByteArray(%s[:])", name)
}

func (g *genUnmarshal) basic(name, kind string) {
	addrOf := "&"
	if name == defaultName {
		addrOf = ""
	}

	g.guard("s = s.Expect%s((*%s)(%s%s))",
		strings.Title(kind),
		kind,
		addrOf,
		name)
}

func (g *genUnmarshal) guard(stmt string, args ...interface{}) {
	stmt = strings.Join(strings.Split(stmt, "\n"), "\n\t\t")
	fmt.Fprintf(g,
		fmt.Sprintf(
			"	if s.Err == nil {\n"+
				"		%s\n"+
				"	}\n",
			stmt),
		args...)
}