generator.go

package generator

import (
	"bytes"
	"fmt"
	"go/types"
	"log"
	"sort"
	"strconv"
	"strings"

	"github.com/subtle-byte/mockigo/internal/generator/path_trie"
	"github.com/subtle-byte/mockigo/internal/util"
	"golang.org/x/tools/go/packages"
)

type Config struct {
	TargetPkgDirPath string
	Targets          Targets
	OutPkgName       func(inspectedPkgName string) string
	OutFilePath      string
	OutPublic        bool
	GoGenCmd         string
}

type Targets struct {
	Include    bool
	Exceptions map[string]struct{}
}

func Generate(cfg Config) error {
	pkgs, err := packages.Load(&packages.Config{
		Mode: packages.NeedName | packages.NeedFiles |
			packages.NeedTypes | packages.NeedTypesInfo,
		Dir:        cfg.TargetPkgDirPath,
		Tests:      false,
		BuildFlags: nil, // TODO?
	}, ".")
	if err != nil {
		return fmt.Errorf("load packages: %w", err)
	}
	if len(pkgs) == 1 && len(pkgs[0].GoFiles) == 0 {
		log.Printf("WARN: no packages found\n")
		return nil
	}
	for _, pkg := range pkgs {
		goFiles := pkg.GoFiles
		if len(goFiles) == 0 {
			continue
		}
		err = generateForPackage(pkg.Types, cfg)
		if err != nil {
			return fmt.Errorf("generate for package %s: %w", pkg.PkgPath, err)
		}
	}
	return nil
}

// targetDesc identifies either interface or function.
type targetDesc struct {
	Name       string
	TypeParams *types.TypeParamList
	Methods    []*types.Func
}

func findTargetByIdent(pkg *types.Package, ident string) *targetDesc {
	obj := pkg.Scope().Lookup(ident)
	typeNameObj, ok := obj.(*types.TypeName)
	if !ok {
		return nil
	}
	var typeParams *types.TypeParamList
	if namedType, ok := obj.Type().(*types.Named); ok {
		typeParams = namedType.TypeParams()
	}
	methods := []*types.Func(nil)
	switch underlying := typeNameObj.Type().Underlying().(type) {
	case *types.Interface:
		for i := 0; i < underlying.NumMethods(); i++ {
			methods = append(methods, underlying.Method(i))
		}
	case *types.Signature:
		methods = []*types.Func{types.NewFunc(0, pkg, "Execute", underlying)}
	default:
		return nil
	}
	return &targetDesc{
		Name:       ident,
		TypeParams: typeParams,
		Methods:    methods,
	}
}

func generateForPackage(pkg *types.Package, cfg Config) error {
	buf := &bytes.Buffer{}
	w := &writer{buf: buf}

	w.Println("// Code generated by mockigo. DO NOT EDIT.")
	w.Println()
	if cfg.GoGenCmd != "" {
		w.Println("//go:generate " + cfg.GoGenCmd)
		w.Println()
	}
	mocksPkg := cfg.OutPkgName(pkg.Name())
	w.Println("package " + mocksPkg)

	targets := []targetDesc(nil)

	for _, ident := range pkg.Scope().Names() {
		_, isException := cfg.Targets.Exceptions[ident]
		delete(cfg.Targets.Exceptions, ident)
		if isException == cfg.Targets.Include {
			continue
		}
		target := findTargetByIdent(pkg, ident)
		if target == nil {
			if isException {
				return fmt.Errorf("target %q is not interface or function", ident)
			}
			continue
		}
		targets = append(targets, *target)
	}
	if len(cfg.Targets.Exceptions) > 0 {
		target := ""
		for ident := range cfg.Targets.Exceptions {
			target = ident
			break
		}
		return fmt.Errorf("target %q is not found", target)
	}

	if len(targets) == 0 {
		return nil
	}

	generatingToSameDir, err := util.FileInDir(cfg.TargetPkgDirPath, cfg.OutFilePath)
	if err != nil {
		return err
	}
	importNames, pkgToImportName := generateImports(w, targets, pkg, !generatingToSameDir || mocksPkg != pkg.Name())
	w.Println()
	w.Println("var _ = match.Any[int]")

	for _, target := range targets {
		generateForInterface(w, target, cfg.OutPublic, pkgToImportName, importNames)
	}

	if w.err != nil {
		return fmt.Errorf("writing file: %w", w.err)
	}
	return writeToFile(w.buf.Bytes(), cfg.OutFilePath)
}

func generateImports(
	w *writer, targets []targetDesc, pkgOfTargets *types.Package, importPkgOfTargets bool,
) (importNames map[string]struct{}, pkgToImportName func(pkg *types.Package) string) {
	trie := path_trie.New()

	trie.LoadPath(path_trie.Path{"mock"}, "github.com/subtle-byte/mockigo/mock")
	trie.LoadPath(path_trie.Path{"match"}, "github.com/subtle-byte/mockigo/match")

	maxConsecutiveUnderscores := 0
	walkingPkgQualifier := func(pkg *types.Package) string {
		maxConsecutiveUnderscores = util.CountMaxConsecutiveUnderscores(pkg.Path(), maxConsecutiveUnderscores)
		splittedPath := strings.Split(pkg.Path(), "/")
		for i := range splittedPath {
			splittedPath[i] = strings.ReplaceAll(splittedPath[i], "-", "")
		}
		trie.LoadPath(splittedPath, pkg.Path())
		return ""
	}

	for _, target := range targets {
		for _, method := range target.Methods {
			signature := method.Type().(*types.Signature)
			for _, tuple := range [2]*types.Tuple{signature.Params(), signature.Results()} {
				for i := 0; i < tuple.Len(); i++ {
					types.TypeString(tuple.At(i).Type(), walkingPkgQualifier)
				}
			}
		}
		if target.TypeParams != nil {
			for i := 0; i < target.TypeParams.Len(); i++ {
				types.TypeString(target.TypeParams.At(i).Constraint(), walkingPkgQualifier)
			}
		}
	}

	importPathToName := map[string]string{}
	importNames = map[string]struct{}{}

	for _, reducedPath := range trie.ReducedPaths() {
		importPath := reducedPath.Meta.(string)
		if importPath == pkgOfTargets.Path() && !importPkgOfTargets {
			continue
		}
		name := strings.Join(reducedPath.Rest, strings.Repeat("_", maxConsecutiveUnderscores+1))
		importNames[name] = struct{}{}
		importPathToName[importPath] = name
	}

	sortedImports := make([]string, 0, len(importPathToName))
	for path, name := range importPathToName {
		sortedImports = append(sortedImports, fmt.Sprintf("import %s %q", name, path))
	}
	sort.Strings(sortedImports)
	w.Println()
	for _, aImport := range sortedImports {
		w.Println(aImport)
	}

	pkgToImportName = func(pkg *types.Package) string {
		if pkg.Path() == pkgOfTargets.Path() && !importPkgOfTargets {
			return ""
		}
		name, ok := importPathToName[pkg.Path()]
		if !ok {
			return "<internal error>"
		}
		return name
	}

	return
}

func genericFormats(typeParams *types.TypeParamList, pkgToImportName func(pkg *types.Package) string) (onTarget, inReceiver string) {
	if typeParams == nil {
		return "", ""
	}
	identAndType := make([]string, typeParams.Len())
	idents := make([]string, typeParams.Len())
	for i := 0; i < typeParams.Len(); i++ {
		tp := typeParams.At(i)
		ident := tp.Obj().Name()
		typeStr := types.TypeString(tp.Constraint(), pkgToImportName)
		identAndType[i] = ident + " " + typeStr
		idents[i] = ident
	}
	onTarget = "[" + strings.Join(identAndType, ", ") + "]"
	inReceiver = "[" + strings.Join(idents, ", ") + "]"
	return
}

func generateForInterface(w *writer, target targetDesc, public bool, pkgToImportName func(pkg *types.Package) string, forbiddenIdents map[string]struct{}) {
	typeParamsOnTarget, typeParamsInReceiver := genericFormats(target.TypeParams, pkgToImportName)

	w.Println()
	mockName := util.SetCaseForFirstLetter(target.Name, true) + "Mock"
	w.Println("type ", mockName, typeParamsOnTarget, " struct {\n\tmock *mock.Mock\n}")
	w.Println()
	w.Println("func ", util.SetCaseForFirstLetter("New"+mockName, public), typeParamsOnTarget, "(t mock.Testing) *", mockName, typeParamsInReceiver, " {")
	w.Println("\tt.Helper()")
	w.Println("\treturn &", mockName, typeParamsInReceiver, "{mock: mock.NewMock(t)}")
	w.Println("}")
	w.Println()
	expecterName := "_" + mockName + "_Expecter"
	w.Println("type ", expecterName, typeParamsOnTarget, " struct {\n\tmock *mock.Mock\n}")
	expecterName += typeParamsInReceiver
	w.Println()
	w.Println("func (_mock *", mockName, typeParamsInReceiver, ") EXPECT() ", expecterName, " {")
	w.Println("\t return ", expecterName, "{mock: _mock.mock}")
	w.Println("}")

	for _, method := range target.Methods {
		generateForMethod(w, mockName, typeParamsOnTarget, typeParamsInReceiver, expecterName, method, pkgToImportName, forbiddenIdents)
	}
}

func isNilable(aType types.Type) bool {
	if underlying, ok := aType.(*types.Named); ok {
		aType = underlying.Underlying()
	}
	switch aType.(type) {
	case *types.Pointer, *types.Array, *types.Slice, *types.Map, *types.Chan, *types.Signature, *types.Interface:
		return true
	default:
		return false
	}
}

func generateForMethod(w *writer, mockName, typeParamsOnTarget, typeParamsInReceiver, expecterName string, method *types.Func, pkgToImportName func(pkg *types.Package) string, forbiddenNames map[string]struct{}) {
	signature := method.Type().(*types.Signature)
	variadic := signature.Variadic()

	inFormats := NewTupleFormatter(signature.Params(), variadic, pkgToImportName).Format("_a", forbiddenNames)
	outFormats := NewTupleFormatter(signature.Results(), false, pkgToImportName).Format("_r", forbiddenNames)

	callStruct := "_" + mockName + "_" + method.Name() + "_Call"
	w.Println()
	w.Println("type ", callStruct, typeParamsOnTarget, " struct {\n\t*mock.Call\n}")
	callStruct += typeParamsInReceiver
	w.Println("\nfunc (_mock *", mockName, typeParamsInReceiver, ") ", method.Name(), "(", inFormats.NamedParams, ") (", outFormats.RawParams, ") {")
	w.Println("\t_mock.mock.T.Helper()")
	w.Print(inFormats.VariadicArgsEval)
	results := "_results := "
	if signature.Results().Len() == 0 {
		results = ""
	}
	if !variadic {
		w.Println("\t", results, "_mock.mock.Called(\"", method.Name(), "\", ", inFormats.Args, ")")
	} else {
		w.Println("\t", results, "_mock.mock.Called(\"", method.Name(), "\", ", inFormats.VariadicArgs, ")")
	}
	if signature.Results().Len() != 0 {
		returnArgs := []string(nil)
		for i := 0; i < signature.Results().Len(); i++ {
			resultVar := signature.Results().At(i)
			varName := "_r" + strconv.Itoa(i)
			returnArgs = append(returnArgs, varName)
			typeStr := types.TypeString(resultVar.Type(), pkgToImportName)
			if typeStr == "error" {
				w.Println("\t", varName, " := _results.Error(", strconv.Itoa(i), ")")
			} else if !isNilable(resultVar.Type()) {
				w.Println("\t", varName, " := _results.Get(", strconv.Itoa(i), ").(", typeStr, ")")
			} else {
				w.Println("\tvar ", varName, " ", typeStr)
				w.Println("\tif _got := _results.Get(", strconv.Itoa(i), "); _got != nil {")
				w.Println("\t\t", varName, " = _got.(", typeStr, ")")
				w.Println("\t}")
			}
		}
		if len(returnArgs) != 0 {
			w.Println("\treturn ", strings.Join(returnArgs, ", "))
		}
	}
	w.Println("}")

	w.Println()
	w.Println("func (_expecter ", expecterName, ") ", method.Name(), "(", inFormats.NamedArgedParams, ") ", callStruct, " {")
	w.Print(inFormats.VariadicArgsMatchersEval)
	if !variadic {
		w.Println("\treturn ", callStruct, `{Call: _expecter.mock.ExpectCall("`, method.Name(), `", `, inFormats.ArgsMatchers, ")}")
	} else {
		w.Println("\treturn ", callStruct, `{Call: _expecter.mock.ExpectCall("`, method.Name(), `", `, inFormats.VariadicArgs, ")}")
	}
	w.Println("}")

	w.Println()
	w.Println("func (_call ", callStruct, ") Return(", outFormats.NamedParams, ") ", callStruct, " {")
	w.Println("\t_call.Call.Return(", outFormats.Args, ")")
	w.Println("\treturn _call")
	w.Println("}")

	w.Println("")
	runReturnFuncType := "func(" + inFormats.RawParams + ") (" + outFormats.RawParams + ")"
	w.Println("func (_call ", callStruct, ") RunReturn(f ", runReturnFuncType, ") ", callStruct, " {")
	w.Println("\t_call.Call.RunReturn(f)")
	w.Println("\treturn _call")
	w.Println("}")
}