main.go

// This program generates Go code with wrappers for exported methods of a global var.
//
// Usage:
//
//	Generates functions for exported methods on MyStruct
//	//go:generate gen-singleton-functions github.com/my/package/path MyStruct
//
//	Generates functions for exported methods on *MyStruct
//	//go:generate gen-singleton-functions github.com/my/package/path *MyStruct
package main

import (
	"go/types"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/dave/jennifer/jen"
	"golang.org/x/tools/go/packages"
)

var rePkgReferencedType = regexp.MustCompile(`([^*\[\]]+)\.[^.]+`)

//nolint:intrange
func main() {

	//nolint:mnd
	if len(os.Args) != 3 {
		log.Fatal("expected 2 arguments -- <package name> <type name>")
	}

	packageName := os.Args[1]
	pkg := loadPackage(packageName)

	typeName := os.Args[2]
	typeTarget, pointer := strings.CutPrefix(typeName, "*")
	typeObject := pkg.Types.Scope().Lookup(typeTarget)
	if typeObject == nil {
		log.Fatalf("unable to find type %s", typeTarget)
	}

	varName := "global" + typeTarget
	isSetVarName := varName + "Set"

	f := jen.NewFile(typeObject.Pkg().Name())

	f.PackageComment("Code generated by generator, DO NOT EDIT.")
	f.PackageComment("")
	f.PackageComment("This file provides the wrappers for exported methods of a global var `" + varName + "`.")

	imports := make(map[string]struct{})
	funcs := jen.Empty()

	typeObjectMeta := typeObject.Type()
	if pointer {
		typeObjectMeta = types.NewPointer(typeObjectMeta)
	}
	methods := types.NewMethodSet(typeObjectMeta)
	for i := 0; i < methods.Len(); i++ {
		method := methods.At(i).Obj()

		// Skip private methods
		if !method.Exported() {
			continue
		}

		signature := method.Type().(*types.Signature)
		signatureParams := signature.Params()

		params := make([]jen.Code, 0, signatureParams.Len())
		paramNames := make([]jen.Code, 0, signatureParams.Len())

		// Get imported packages from params and prepare param names and types
		for j := 0; j < signatureParams.Len(); j++ {
			param := signatureParams.At(j)
			paramType := param.Type().String()

			// Remove current package prefix
			paramType = strings.ReplaceAll(paramType, packageName+".", "")

			// Remove prefixes of imported packages
			paramType = extractImports(paramType, imports)

			params = append(params, jen.Id(param.Name()).Id(paramType))
			paramNames = append(paramNames, jen.Id(param.Name()))
		}

		results := signature.Results()
		returnTypes := make([]string, 0, results.Len())

		// Get imported packages from returns and parepare return types
		for j := 0; j < results.Len(); j++ {
			resultType := results.At(j).Type().String()

			// Remove current package prefix
			resultType = strings.ReplaceAll(resultType, packageName+".", "")

			// Remove prefixes of imported packages
			resultType = extractImports(resultType, imports)

			returnTypes = append(returnTypes, resultType)
		}

		var returnType string
		var lastStatement jen.Code = jen.Id(varName).Dot(method.Name()).Call(paramNames...)

		if len(returnTypes) > 0 {
			returnType = "(" + strings.Join(returnTypes, ", ") + ")"
			lastStatement = jen.Return(lastStatement)
		}

		panicMsg := varName + " must be set before calling " + method.Name() + "()"
		var globalVarCheck jen.Code
		if pointer {
			globalVarCheck = jen.If(jen.Id(varName).Op("==").Id("nil")).Block(
				jen.Panic(jen.Lit(panicMsg)),
			)
		} else {
			globalVarCheck = jen.If(jen.Op("!").Id(isSetVarName)).Block(
				jen.Panic(jen.Lit(panicMsg)),
			)
		}

		funcs.Comment(method.Name() + " calls `" + varName + "." + method.Name() + "`.")
		funcs.Line()
		funcs.Func().Id(method.Name()).Params(
			params...,
		).Id(returnType).Block(
			globalVarCheck,
			jen.Line(),
			lastStatement,
		)
		funcs.Line()
	}

	if len(imports) > 0 {
		importsToAdd := jen.Empty()
		for importPackage := range imports {
			importsToAdd.Add(
				jen.Lit(importPackage),
				jen.Line(),
			)
		}

		f.Comment("Import missing dependencies")
		f.Add(jen.Id("import").Parens(importsToAdd))
	}

	f.Id("var").Id(varName).Id(typeName)

	var setBoolVarCode jen.Code
	if !pointer {
		f.Id("var").Id(isSetVarName).Id("bool")
		setBoolVarCode = jen.Id(isSetVarName).Op("=").Id("true")
	}

	f.Func().Id("SetGlobal"+typeTarget).Params(jen.Id("v").Id(typeName)).Block(
		jen.Id(varName).Op("=").Id("v"),
		setBoolVarCode,
	)

	f.Add(funcs)

	filename := os.Getenv("GOFILE")
	ext := filepath.Ext(filename)
	filename, _ = strings.CutSuffix(filename, ext)
	filename = filename + "_gen.go"

	if err := f.Save(filename); err != nil {
		log.Fatal(err)
	}
}

func loadPackage(path string) *packages.Package {
	cfg := &packages.Config{Mode: packages.NeedTypes}
	pkgs, err := packages.Load(cfg, path)
	if err != nil {
		log.Fatalf("loading packages for inspection: %v", err)
	}
	if packages.PrintErrors(pkgs) > 0 {
		log.Fatalf("errors: %v", packages.PrintErrors(pkgs))
	}

	return pkgs[0]
}

// Remove imported packages prefixes and accumulate missing imports.
//
//	  github.com/somelib/name.Map -> name.Map
//		[]github.com/somelib/name.Item -> []name.Item
func extractImports(paramType string, imports map[string]struct{}) string {
	if strings.Index(paramType, ".") < 0 {
		return paramType
	}

	submatches := rePkgReferencedType.FindStringSubmatch(paramType)

	imports[submatches[1]] = struct{}{}

	for packageName := range imports {
		paramType = strings.ReplaceAll(paramType, filepath.Dir(packageName)+"/", "")
	}

	return paramType
}