processor.go

package generator

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"path/filepath"
	"reflect"
	"strings"

	_ "golang.org/x/tools/go/gcimporter"
	"golang.org/x/tools/go/types"
)

const BaseDocument = "github.com/tyba/storable.Document"

type Processor struct {
	Path   string
	Ignore map[string]bool

	TypesPkg *types.Package
}

func NewProcessor(path string, ignore []string) *Processor {
	i := make(map[string]bool, 0)
	for _, file := range ignore {
		i[file] = true
	}

	return &Processor{
		Path:   path,
		Ignore: i,
	}
}

func (p *Processor) Do() (*Package, error) {
	files, err := p.getSourceFiles()
	if err != nil {
		return nil, err
	}

	p.TypesPkg, _ = p.parseSourceFiles(files)
	return p.ProcessTypesPkg()
}

func (p *Processor) ProcessTypesPkg() (*Package, error) {
	pkg := &Package{Name: p.TypesPkg.Name()}
	p.processPackage(pkg)

	return pkg, nil
}

func (p *Processor) getSourceFiles() ([]string, error) {
	pkg, err := build.Default.ImportDir(p.Path, 0)
	if err != nil {
		return nil, fmt.Errorf("cannot process directory %s: %s", p.Path, err)
	}

	var files []string
	files = append(files, pkg.GoFiles...)
	files = append(files, pkg.CgoFiles...)

	if len(files) == 0 {
		return nil, fmt.Errorf("%s: no buildable Go files", p.Path)
	}

	return joinDirectory(p.Path, files), nil
}

func (p *Processor) parseSourceFiles(filenames []string) (*types.Package, error) {
	var files []*ast.File
	fs := token.NewFileSet()
	for _, filename := range filenames {
		if _, ok := p.Ignore[filename]; ok {
			continue
		}

		file, err := parser.ParseFile(fs, filename, nil, 0)
		if err != nil {
			return nil, fmt.Errorf("parsing package: %s: %s", filename, err)
		}

		files = append(files, file)
	}

	config := types.Config{FakeImportC: true, Error: func(error) {}}
	info := &types.Info{}

	return config.Check(p.Path, fs, files, info)
}

func (p *Processor) processPackage(pkg *Package) {
	var newFuncs []*types.Func

	s := p.TypesPkg.Scope()
	for _, name := range s.Names() {
		fun := p.tryGetFunction(s.Lookup(name))
		if fun != nil {
			pkg.Functions = append(pkg.Functions, name)
			if strings.HasPrefix(fun.Name(), "new") {
				newFuncs = append(newFuncs, fun)
			}
		}

		str := p.tryGetStruct(s.Lookup(name).Type())
		if str == nil {
			continue
		}

		if m := p.processStruct(name, str); m != nil {
			pkg.Models = append(pkg.Models, m)
			m.CheckedNode = s.Lookup(name).Type().(*types.Named)
			m.Package = p.TypesPkg
		} else {
			pkg.Structs = append(pkg.Structs, name)
		}
	}

	for _, fun := range newFuncs {
		p.tryMatchNewFunc(pkg.Models, fun)
	}
}

func (p *Processor) tryMatchNewFunc(models []*Model, fun *types.Func) {
	modelName := fun.Name()[len("new"):]

	for _, m := range models {
		if m.Name != modelName {
			continue
		}

		sig := fun.Type().(*types.Signature)

		if sig.Recv() != nil {
			continue
		}

		res := sig.Results()
		for i := 0; i < res.Len(); i++ {
			if isTypeOrPtrTo(res.At(i).Type(), m.CheckedNode) {
				m.NewFunc = fun
				return
			}
		}
	}
}

func (p *Processor) tryGetFunction(typ types.Object) *types.Func {
	switch t := typ.(type) {
	case *types.Func:
		return t
	}

	return nil
}

func (p *Processor) tryGetStruct(typ types.Type) *types.Struct {
	switch t := typ.(type) {
	case *types.Named:
		return p.tryGetStruct(t.Underlying())
	case *types.Pointer:
		return p.tryGetStruct(t.Elem())
	case *types.Slice:
		return p.tryGetStruct(t.Elem())
	case *types.Map:
		return p.tryGetStruct(t.Elem())
	case *types.Struct:
		return t
	}

	return nil
}

func (p *Processor) processStruct(name string, s *types.Struct) *Model {
	m := NewModel(name)

	var base int
	if base, m.Fields = p.getFields(s); base == -1 {
		return nil
	}

	p.procesBaseField(m, m.Fields[base])
	p.findHooks(m)
	p.findStoreHooks(m)

	return m
}

func (p *Processor) getFields(s *types.Struct) (base int, fields []*Field) {
	base, fields = p.processFields(s, []*types.Struct{})
	return
}

// Returns which field index is an embedded storable.Document, or -1 if none.
func (p *Processor) processFields(s *types.Struct, done []*types.Struct) (base int, fields []*Field) {
	c := s.NumFields()

	base = -1
	fields = make([]*Field, 0)

	for i := 0; i < c; i++ {
		f := s.Field(i)
		if !f.Exported() {
			continue
		}

		t := reflect.StructTag(s.Tag(i))
		if f.Type().String() == BaseDocument {
			base = i
		}

		field := NewField(f.Name(), f.Type().Underlying().String(), t)
		field.CheckedNode = f
		str := p.tryGetStruct(f.Type())
		if f.Type().String() != BaseDocument && str != nil {
			field.Type = getStructType(f.Type())

			d := false
			for _, v := range done {
				if v == str {
					d = true
					break
				}
			}
			if !d {
				_, subfs := p.processFields(str, append(done, str))
				field.SetFields(subfs)
			}
		}

		fields = append(fields, field)
	}

	return base, fields
}

func getStructType(t types.Type) string {
	ts := t.String()
	if ts != "time.Time" && ts != "bson.ObjectId" {
		return "struct"
	}

	return ts
}

func (p *Processor) procesBaseField(m *Model, f *Field) {
	m.Collection = f.Tag.Get("collection")
}

func joinDirectory(directory string, files []string) []string {
	r := make([]string, len(files))
	for i, file := range files {
		r[i] = filepath.Join(directory, file)
	}

	return r
}

func (p *Processor) findHooks(m *Model) {
	modelType := types.NewPointer(p.TypesPkg.Scope().Lookup(m.Name).Type())
	m.Hooks = p.hookMethods(modelType, nil)

	for _, f := range m.Fields {
		p.findFieldHooks(f, []interface{}{modelType})
	}
}

func (p *Processor) findFieldHooks(f *Field, done []interface{}) {
	if f.CheckedNode == nil {
		return
	}

	for _, v := range done {
		if v == f.CheckedNode.Type() {
			return
		}
	}

	typ := f.CheckedNode.Type()
out:
	for {
		switch v := typ.(type) {
		case *types.Slice:
			typ = v.Elem()
		case *types.Map:
			typ = v.Elem()
		default:
			break out
		}
	}

	if _, ok := typ.(*types.Pointer); !ok {
		typ = types.NewPointer(typ)
	}

	f.Hooks = p.hookMethods(typ, nil)

	for _, f := range f.Fields {
		p.findFieldHooks(f, append(done, f.CheckedNode.Type()))
	}
}

func (p *Processor) hookMethods(t types.Type, storeOf *Model) []Hook {
	var hooks []Hook

	ms := types.NewMethodSet(t)

	actions := []HookAction{
		InsertHook,
		UpdateHook,
		SaveHook,
	}

	for _, action := range actions {
		hook := Hook{Before: false, Action: action}
		if p.hasHookMethod(ms, hook.MethodName(), storeOf) {
			hooks = append(hooks, hook)
		}
		hook.Before = true
		if p.hasHookMethod(ms, hook.MethodName(), storeOf) {
			hooks = append(hooks, hook)
		}
	}

	return hooks
}

func (p *Processor) hasHookMethod(ms *types.MethodSet, methodName string, storeOf *Model) bool {
	sel := ms.Lookup(p.TypesPkg, methodName)
	if sel == nil {
		return false
	}
	method, ok := sel.Obj().(*types.Func)
	if !ok {
		return false
	}
	sig, ok := method.Type().(*types.Signature)
	if !ok {
		return false
	}

	params := sig.Params()
	if storeOf != nil {
		if params == nil || params.Len() != 1 {
			return false
		}
		ty, ok := params.At(0).Type().(*types.Pointer)
		if !ok {
			return false
		}
		named, ok := ty.Elem().(*types.Named)
		if ok {
			if named.Obj().Name() != storeOf.Name+"Store" {
				return false
			}
		} else {
			if ty.Elem() != types.Typ[types.Invalid] {
				return false
			}
		}
	} else {
		if params != nil && params.Len() > 0 {
			return false
		}
	}

	ret := sig.Results()
	if ret == nil || ret.Len() != 1 || !isBuiltinError(ret.At(0).Type()) {
		return false
	}

	return true
}

func (p *Processor) findStoreHooks(m *Model) {
	modelType := types.NewPointer(p.TypesPkg.Scope().Lookup(m.Name).Type())
	m.StoreHooks = p.hookMethods(modelType, m)
}

func isBuiltinError(typ types.Type) bool {
	named, ok := typ.(*types.Named)
	if !ok {
		return false
	}

	return named.Obj().Name() == "error" && named.Obj().Parent() == types.Universe
}