chaingen.go

package chaingen

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/token"
	"go/types"
	"io"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	"golang.org/x/tools/go/packages"
)

type Chaingen struct {
	opts Options
}

func New(opts Options) Chaingen {
	return Chaingen{
		opts: opts,
	}
}

// Builder represents a struct type that is considered to be a builder.
// Builder has at least one method that returns altered Builder copy (chaining method).
// Every non-chaining method is considered as finalizer
type Builder struct {
	PkgPath          string
	Package          *packages.Package
	Type             *types.Named
	Struct           *types.Struct
	Methods          []Method
	Children         []*BuilderRef
	MethodNames      map[string]*Method
	File             *ast.File
	FilePath         string
	Rendered         bool
	GeneratedMethods []Method
	Depth            int
}

type BuilderRef struct {
	Name    string
	Tag     string
	Builder *Builder
}

type Import struct {
	Alias   string
	Package *types.Package
}

type Method struct {
	Name         string
	Alias        string
	Scope        *types.Scope
	Pos          token.Pos
	Variadic     bool
	Exported     bool
	Builder      *Builder
	BuilderField string

	Recv          MethodParam
	Params        []MethodParam
	Results       []MethodParam
	Prefixes      []string
	FinalizerName string
}

func (m Method) String() string {
	return m.Builder.Type.Obj().Name() + "." + m.Name
}

type MethodParam struct {
	Name  string
	Type  types.Type
	Named *types.Named
}

func NewMethod(builder *Builder, f *types.Func, sig *types.Signature) Method {
	m := Method{
		Name:     f.Name(),
		Alias:    f.Name(),
		Pos:      f.Pos(),
		Scope:    f.Scope(),
		Variadic: sig.Variadic(),
		Builder:  builder,
		Exported: f.Exported(),
		Recv: MethodParam{
			Name:  sig.Recv().Name(),
			Type:  sig.Recv().Type(),
			Named: builderType(sig.Recv().Type()),
		},
	}
	for i := 0; i < sig.Params().Len(); i++ {
		param := sig.Params().At(i)
		m.Params = append(m.Params, MethodParam{
			Name: param.Name(),
			Type: param.Type(),
		})
	}
	for i := 0; i < sig.Results().Len(); i++ {
		param := sig.Results().At(i)
		m.Results = append(m.Results, MethodParam{
			Name: param.Name(),
			Type: param.Type(),
		})
	}
	return m
}

func (b *Builder) ReceiverName() string {
	if len(b.Methods) > 0 {
		return b.Methods[0].Recv.Name
	}
	return strings.ToLower(b.Type.Obj().Name()[0:1])
}

func (b *Builder) ReceiverType() string {
	s := strings.Builder{}
	s.WriteString(b.Type.Obj().Name())
	tps := b.Type.TypeParams()
	var tpss []string
	for i := 0; i < tps.Len(); i++ {
		tp := tps.At(i)
		tpss = append(tpss, tp.String())
	}
	if len(tpss) > 0 {
		s.WriteString("[" + strings.Join(tpss, ", ") + "]")
	}
	return s.String()
}

func (b *Builder) Ref(builder *Builder) *BuilderRef {
	for _, child := range b.Children {
		if child.Builder == builder {
			return child
		}
	}
	return nil
}

func (b *Builder) RenderChainMethod(file *File, method Method) {
	var inputParams []string
	var callParams []string
	for i, param := range method.Params {
		last := i == len(method.Params)-1
		if last && method.Variadic {
			inputParams = append(inputParams, param.Name+" ..."+file.TypeIdentifier(param.Type.(*types.Slice).Elem()))
			callParams = append(callParams, param.Name+"...")
		} else {
			inputParams = append(inputParams, param.Name+" "+file.TypeIdentifier(param.Type))
			callParams = append(callParams, param.Name)
		}
	}
	file.L()
	doc := method.Doc()
	if doc != nil {
		for _, line := range doc.List {
			file.L(line.Text)
		}
	}
	ref := b.ReceiverName() + "." + b.Ref(method.Builder).Name
	file.L(`func (` + b.ReceiverName() + ` ` + b.ReceiverType() + `) ` + method.Alias + `(` + strings.Join(inputParams, ", ") + `) ` + b.ReceiverType() + " {")
	file.L("\t" + ref + ` = ` + ref + `.` + method.Name + `(` + strings.Join(callParams, ", ") + `)`)
	file.L("\treturn " + b.ReceiverName())
	file.L("}")
}

func (b *Builder) RenderFinalizer(file *File, method Method) {
	var inputParams []string
	var callParams []string
	for i, param := range method.Params {
		last := i == len(method.Params)-1
		if last && method.Variadic {
			inputParams = append(inputParams, param.Name+" ..."+file.TypeIdentifier(param.Type.(*types.Slice).Elem()))
			callParams = append(callParams, param.Name+"...")
		} else {
			inputParams = append(inputParams, param.Name+" "+file.TypeIdentifier(param.Type))
			callParams = append(callParams, param.Name)
		}
	}
	var outputParams []string
	for _, param := range method.Results {
		typeName := file.TypeIdentifier(param.Type)
		if typeName == "" {
			// TODO: error
			return
		}
		name := param.Name
		if name != "" {
			name += " "
		}
		outputParams = append(outputParams, name+typeName)
	}
	outputParamsStr := strings.Join(outputParams, ", ")
	if len(outputParams) > 1 {
		outputParamsStr = "(" + outputParamsStr + ")"
	}
	file.L()
	doc := method.Doc()
	if doc != nil {
		for _, line := range doc.List {
			file.L(line.Text)
		}
	}
	ref := b.ReceiverName() + "." + b.Ref(method.Builder).Name
	file.L(`func (` + b.ReceiverName() + ` ` + b.ReceiverType() + `) ` + method.Alias + `(` + strings.Join(inputParams, ", ") + `) ` + outputParamsStr + " {")
	file.P("\t")
	if len(outputParams) > 0 {
		file.P("return ")
	}
	if method.FinalizerName != "" {
		for _, prefix := range method.Prefixes {
			file.L(prefix)
		}
		file.P(b.ReceiverName(), ".", method.FinalizerName, "(")
	}
	file.P(ref, `.`, method.Name, `(`, strings.Join(callParams, ", "), `)`)
	if method.FinalizerName != "" {
		file.P(")")
	}
	file.L("}")
}

func (m Method) IsChaining() bool {
	if len(m.Results) != 1 {
		return false
	}
	return m.Results[0].Type == m.Builder.Type
}

func (m Method) IsFinalizer() bool {
	if m.FinalizerName != "" {
		return true
	}
	return !m.IsChaining()
}

const generatedPrefix = "Code generated by chaingen. DO NOT EDIT."

func (m Method) Doc() *ast.CommentGroup {
	pkg := m.Builder.Package
	methodPos := pkg.Fset.Position(m.Pos)
	for _, file := range pkg.Syntax {
		for _, cg := range file.Comments {
			commentPos := pkg.Fset.Position(cg.End())
			if commentPos.Filename == methodPos.Filename && commentPos.Line == methodPos.Line-1 {
				if m.Name != m.Alias {
					firstLine := cg.List[0].Text
					if strings.HasPrefix(firstLine, "// "+m.Name+" ") {
						firstLine = "// " + m.Alias + " " + firstLine[len(m.Name)+4:]
						cg.List[0].Text = firstLine
					}
				}
				return cg
			}
		}
	}
	return nil
}

type Options struct {
	Src           string
	TypeName      string
	Recursive     bool
	FileSuffix    string
	ErrOnConflict bool
	StructTag     string
	BuildTag      string
}

type File struct {
	Builders []*Builder
	Package  *packages.Package
	File     *ast.File
	Path     string
	BuildTag string

	Imports       map[string]Import
	ImportAliases map[string]*Import
	Body          bytes.Buffer
}

func (f *File) Render(w io.Writer) error {
	output := "//go:build !" + f.BuildTag + "\n"
	output += "// +build !" + f.BuildTag + "\n\n"
	output += "// " + generatedPrefix + "\n\n"
	output += "package " + f.Package.Name + "\n\n"
	if len(f.Imports) > 0 {
		output += "import (\n"
		for _, i := range f.Imports {
			alias := ""
			if i.Alias != i.Package.Name() {
				alias = i.Alias + " "
			}
			output += "\t" + alias + `"` + i.Package.Path() + `"` + "\n"
		}
		output += ")\n"
	}
	_, err := w.Write([]byte(output + f.Body.String()))
	return err
}

func (f *File) P(s ...string) {
	f.Body.WriteString(strings.Join(s, ""))
}

func (f *File) L(s ...string) {
	f.P(append(s, "\n")...)
}

func (f *File) TypeIdentifier(typ types.Type) string {
	switch t := typ.(type) {
	case *types.Named:
		pkg := t.Obj().Pkg()
		pkgPath := ""
		if pkg != nil && pkg.Path() != f.Package.PkgPath {
			pkgPath = f.PackageIdentifier(pkg) + "."
		}

		typeArgs := ""
		if ta := t.TypeArgs(); ta != nil {
			var args []string
			for i := 0; i < ta.Len(); i++ {
				args = append(args, f.TypeIdentifier(ta.At(i)))
			}
			typeArgs = "[" + strings.Join(args, ", ") + "]"
		}
		return pkgPath + t.Obj().Name() + typeArgs
	case *types.Slice:
		return "[]" + f.TypeIdentifier(t.Elem())
	case *types.Map:
		return "map[" + f.TypeIdentifier(t.Key()) + "]" + f.TypeIdentifier(t.Elem())
	case *types.Pointer:
		return "*" + f.TypeIdentifier(t.Elem())
	default:
		return t.String()
	}
}

func (f *File) PackageIdentifier(pkg *types.Package) string {
	i, ok := f.Imports[pkg.Path()]
	if ok {
		return i.Alias
	}
	i = Import{
		Alias:   pkg.Name(),
		Package: pkg,
	}
	f.Imports[pkg.Path()] = i
	suffix := 0
	for {
		_, ok := f.ImportAliases[i.Alias]
		if !ok {
			break
		}
		suffix++
		i.Alias = fmt.Sprintf("%s%d", pkg.Name(), suffix)
	}
	f.ImportAliases[i.Alias] = &i
	return i.Alias
}

func (c Chaingen) Render(builders map[*types.Named]*Builder) (map[string]*File, error) {
	files := make(map[string]*File)
	for _, builder := range builders {
		if builder.Depth > 0 {
			continue
		}
		err := c.render(files, builder)
		if err != nil {
			return nil, err
		}
	}
	return files, nil
}

func (c Chaingen) evalTag(tag string, methods []Method, builderMethods []Method) ([]Method, error) {
	if tag == "" || tag == "*" {
		return methods, nil
	}
	if tag == "-" {
		return nil, nil
	}
	pool := make(map[string]Method, len(methods))
	for _, method := range methods {
		pool[method.Alias] = method
	}
	modifiers := strings.Split(tag, ",")
	for _, modifier := range modifiers {
		if len(modifier) == 0 {
			continue
		}
		parts := strings.Split(modifier, "=")
		switch {
		case modifier == "*":
			break
		case modifier[0] == '-':
			delete(pool, modifier[1:])
		case strings.HasPrefix(modifier, "fin("):
			selector := parts[0][4 : len(parts[0])-1]
			glob := NewGlob(selector)
			newPool := make(map[string]Method, len(pool))
			for _, method := range pool {
				if glob.Match(method.Recv.Named.Obj().Name(), method.Alias) {
					for _, m := range builderMethods {
						if m.Name == parts[1] {
							compatible := m.Variadic
							if !compatible && len(m.Params) == len(method.Results) {
								compatible = true
								for i := 0; i < len(m.Params); i++ {
									if m.Params[i].Type.String() != method.Results[i].Type.String() {
										compatible = false
										break
									}
								}
							}
							if compatible {
								method.FinalizerName = parts[1]
								method.Results = m.Results
							}
							break
						}
					}
				}
				newPool[method.Alias] = method
			}
			pool = newPool
		case strings.HasPrefix(modifier, "pre("):
			selector := parts[0][4 : len(parts[0])-1]
			glob := NewGlob(selector)
			newPool := make(map[string]Method, len(pool))
			for _, method := range pool {
				if glob.Match(method.Recv.Named.Obj().Name(), method.Alias) {
					method.Prefixes = append(method.Prefixes, parts[1])
				}
				newPool[method.Alias] = method
			}
			pool = newPool
		default:
			left := NewGlob(parts[0])
			var right *Glob
			if len(parts) > 1 {
				rightGlob := NewGlob(parts[1])
				right = &rightGlob
			}
			newPool := make(map[string]Method, len(pool))
			for _, method := range pool {
				if left.Match(method.Recv.Named.Obj().Name(), method.Alias) {
					method.Alias = left.Replace(method.Alias, right)
				}
				newPool[method.Alias] = method
			}
			pool = newPool
		}
	}
	result := make([]Method, 0, len(pool))
	for _, original := range methods {
		for _, altered := range pool {
			if original.Pos == altered.Pos {
				result = append(result, altered)
				break
			}
		}
	}
	return result, nil
}

type Glob struct {
	TypeName string
	Const    string
	Prefix   string
	Suffix   string
	Middle   string
}

func NewGlob(glob string) Glob {
	dot := strings.Index(glob, ".")
	typeName := ""
	if dot >= 0 {
		typeName = glob[:dot]
		glob = glob[dot+1:]
	}
	if glob == "*" {
		return Glob{
			TypeName: typeName,
		}
	}
	asterisk := strings.Index(glob, "*")
	if asterisk < 0 {
		return Glob{
			TypeName: typeName,
			Const:    glob,
		}
	}
	if asterisk == 0 && glob[len(glob)-1] == '*' {
		return Glob{
			TypeName: typeName,
			Middle:   glob[1 : len(glob)-1],
		}
	}
	return Glob{
		TypeName: typeName,
		Prefix:   glob[:asterisk],
		Suffix:   glob[asterisk+1:],
	}
}

func (g Glob) Match(typeName string, str string) bool {
	if g.TypeName != "" && g.TypeName != typeName {
		return false
	}
	if g.Const != "" {
		return str == g.Const
	}
	if g.Middle != "" {
		return strings.Contains(str, g.Middle)
	}
	match := true
	if g.Prefix != "" {
		match = match && len(str) > len(g.Prefix) && strings.HasPrefix(str, g.Prefix)
	}
	if g.Suffix != "" {
		match = match && len(str) > len(g.Suffix) && strings.HasSuffix(str, g.Suffix)
	}
	return match
}

func (g Glob) Replace(methodName string, right *Glob) string {
	left := g
	switch {
	case right == nil:
		return methodName
	case right.Const != "":
		return right.Const
	case left.Middle != "":
		pos := strings.Index(methodName, g.Middle)
		leftMatchPrefix := methodName[:pos]
		leftMatchSuffix := methodName[pos+len(g.Middle):]
		if right.Middle != "" {
			return leftMatchPrefix + right.Middle + leftMatchSuffix
		}
		return right.Prefix + leftMatchPrefix + leftMatchSuffix + right.Suffix
	default:
		match := methodName[len(left.Prefix):]
		match = match[:len(match)-len(left.Suffix)]
		if right.Middle != "" {
			if len(left.Prefix) == 0 {
				return match + right.Middle
			}
			return right.Middle + match
		}
		return right.Prefix + match + right.Suffix
	}
}

func (c Chaingen) render(files map[string]*File, builder *Builder) error {
	file, ok := files[builder.FilePath]
	if !ok {
		file = &File{
			Package:       builder.Package,
			File:          builder.File,
			BuildTag:      c.opts.BuildTag,
			Path:          builder.FilePath,
			Imports:       map[string]Import{},
			ImportAliases: map[string]*Import{},
		}
		files[builder.FilePath] = file
	}
	file.Builders = append(file.Builders, builder)
	if builder.Rendered {
		return nil
	}
	builder.Rendered = true
	for _, child := range builder.Children {
		var methods []Method
		if c.opts.Recursive {
			err := c.render(files, child.Builder)
			if err != nil {
				return err
			}
			methods = append(methods, child.Builder.GeneratedMethods...)
		}
		methods = append(methods, child.Builder.Methods...)

		methods, err := c.evalTag(child.Tag, methods, builder.Methods)
		if err != nil {
			return err
		}
		for _, m := range methods {
			if !m.Exported && m.Builder.PkgPath != builder.PkgPath {
				continue
			}
			if conflict, ok := builder.MethodNames[m.Alias]; ok {
				if c.opts.ErrOnConflict {
					return fmt.Errorf("method naming conflict for %s.%s: %s and %s", builder.Type.Obj().Name(), m.Alias, conflict.String(), m.String())
				}
				continue
			}
			switch {
			case m.IsChaining():
				builder.RenderChainMethod(file, m)
				builder.MethodNames[m.Alias] = &m
				generated := m
				generated.Name = generated.Alias
				generated.Results = []MethodParam{
					{
						Type: builder.Type,
					},
				}
				generated.Builder = builder
				builder.GeneratedMethods = append(builder.GeneratedMethods, generated)
			case m.IsFinalizer():
				builder.RenderFinalizer(file, m)
				builder.MethodNames[m.Alias] = &m
				generated := m
				generated.Name = generated.Alias
				generated.Builder = builder
				builder.GeneratedMethods = append(builder.GeneratedMethods, generated)
			}
		}
	}

	return nil
}

func (c Chaingen) Generate() error {
	if c.opts.Src == "" {
		return fmt.Errorf("source dir is not set")
	}
	if c.opts.FileSuffix == "" {
		return fmt.Errorf("file suffix is not set")
	}
	pkgs, err := packages.Load(&packages.Config{
		Dir:        c.opts.Src,
		Mode:       packages.NeedName | packages.NeedSyntax | packages.NeedTypes | packages.NeedTypesInfo | packages.NeedDeps | packages.NeedImports | packages.NeedModule,
		BuildFlags: []string{"-tags=" + c.opts.BuildTag},
	})
	if err != nil {
		return fmt.Errorf("error loading Go packages from %s: %w", c.opts.Src, err)
	}

	var errors []string
	for _, p := range pkgs {
		for _, imp := range p.Imports {
			for _, err := range imp.Errors {
				errors = append(errors, err.Error())
			}
		}
		for _, err := range p.Errors {
			errors = append(errors, err.Error())
		}
	}
	if len(errors) > 0 {
		return fmt.Errorf("errors occurred loading source code:\n%s\n", strings.Join(errors, "\n"))
	}

	found := make(map[*types.Named]*packages.Package)
	if c.opts.TypeName == "" {
		for _, p := range pkgs {
			for _, name := range p.Types.Scope().Names() {
				typ, _ := toStruct(p.Types.Scope().Lookup(name))
				if typ != nil {
					found[typ] = p
				}
			}
		}
		if len(found) == 0 {
			return fmt.Errorf("unable to find structs in %s", c.opts.Src)
		}
	} else {
		names := strings.Split(c.opts.TypeName, ",")
		for _, p := range pkgs {
			for _, name := range names {
				typ, _ := toStruct(p.Types.Scope().Lookup(name))
				if typ != nil {
					found[typ] = p
				}
			}
		}
		if len(found) == 0 {
			return fmt.Errorf("unable to find struct type %q in %s", c.opts.TypeName, c.opts.Src)
		}
	}

	builders := map[*types.Named]*Builder{}
	for typ, pkg := range found {
		err := c.NewBuilder(builders, pkg, typ)
		if err != nil {
			return fmt.Errorf("error creating builder: %w", err)
		}
	}

	files, err := c.Render(builders)
	if err != nil {
		return fmt.Errorf("error generating code: %w", err)
	}
	for _, file := range files {
		if file.Body.Len() == 0 {
			continue
		}
		buf := bytes.Buffer{}
		err = file.Render(&buf)
		if err != nil {
			return err
		}
		src := file.Path
		dest := src[:len(src)-3] + c.opts.FileSuffix
		formatted, err := format.Source(buf.Bytes())
		if err != nil {
			log.Printf("error formatting file: %s", err.Error())
			formatted = buf.Bytes()
		}
		err = os.WriteFile(dest, formatted, 0755)
		if err != nil {
			return err
		}
		rel, _ := filepath.Rel(c.opts.Src, dest)
		log.Printf("generated file: %s", rel)
	}
	return nil
}

func (c Chaingen) NewBuilder(builders map[*types.Named]*Builder, pkg *packages.Package, n *types.Named) error {
	_, err := c.newBuilder(builders, pkg, n, 0)
	return err
}

func (c Chaingen) newBuilder(builders map[*types.Named]*Builder, pkg *packages.Package, typ *types.Named, depth int) (*Builder, error) {
	if b, ok := builders[typ]; ok {
		return b, nil
	}
	builder := &Builder{
		Depth:       depth,
		Package:     pkg,
		PkgPath:     pkg.PkgPath,
		Type:        typ,
		MethodNames: map[string]*Method{},
	}
	if s, ok := typ.Underlying().(*types.Struct); ok {
		builder.Struct = s
	}
	builders[typ] = builder

	builderFile := pkg.Fset.File(builder.Type.Obj().Pos())
	for _, file := range pkg.Syntax {
		astFile := pkg.Fset.File(file.Pos())
		if astFile.Name() == builderFile.Name() {
			builder.File = file
			builder.FilePath = astFile.Name()
			break
		}
	}
	if builder.File == nil {
		return nil, fmt.Errorf("could not discover file for builder %q", typ.Obj().Name())
	}

	for i := 0; i < typ.NumMethods(); i++ {
		fun := typ.Method(i)
		sig, ok := fun.Type().(*types.Signature)
		if !ok {
			continue
		}
		m := NewMethod(builder, fun, sig)
		builder.Methods = append(builder.Methods, m)
	}
	if builder.Struct == nil {
		return builder, nil
	}

	for i := 0; i < builder.Struct.NumFields(); i++ {
		field := builder.Struct.Field(i)
		tag, _ := reflect.StructTag(builder.Struct.Tag(i)).Lookup(c.opts.StructTag)
		typ := builderType(field.Type())
		if typ == nil {
			continue
		}
		name := typ.Obj().Name()
		if !field.Embedded() {
			name = field.Name()
		}
		childPkg := pkg
		pkgPath := typ.Obj().Pkg().Path()
		if pkgPath != pkg.PkgPath {
			childPkg = pkg.Imports[pkgPath]
		}
		child, err := c.newBuilder(builders, childPkg, typ, depth+1)
		if err != nil {
			return nil, fmt.Errorf("error creating builder for field %s: %w", field.Name(), err)
		}
		builder.Children = append(builder.Children, &BuilderRef{
			Name:    name,
			Tag:     tag,
			Builder: child,
		})
	}

	return builder, nil
}

func toStruct(obj types.Object) (*types.Named, *types.Struct) {
	if obj == nil {
		return nil, nil
	}
	typName, ok := obj.(*types.TypeName)
	if !ok {
		return nil, nil
	}
	typ, ok := typName.Type().(*types.Named)
	if !ok {
		return nil, nil
	}
	str, ok := typ.Underlying().(*types.Struct)
	if !ok {
		return nil, nil
	}
	return typ, str
}

func builderType(typ types.Type) *types.Named {
	if n, ok := typ.(*types.Named); ok {
		_, ok := n.Underlying().(*types.Struct)
		if ok {
			return n
		}
		_, ok = n.Underlying().(*types.Slice)
		if ok {
			return n
		}
	}
	if p, ok := typ.Underlying().(*types.Pointer); ok {
		return builderType(p.Elem())
	}
	return nil
}