gobindlua.go

package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"go/ast"
	"go/types"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"
	"unicode"

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

	"github.com/ChrisTrenkamp/gobindlua/gobindlua/datatype"
	"github.com/ChrisTrenkamp/gobindlua/gobindlua/functiontype"
	"github.com/ChrisTrenkamp/gobindlua/gobindlua/gobindluautil"
	"github.com/ChrisTrenkamp/gobindlua/gobindlua/structfield"
)

func main() {
	workingDir := flag.String("d", "", "The Go source directory to generate the bindings from. Uses the current working directory if empty")
	structName := flag.String("s", "", "Generate the GopherLua bindings for the given struct.")
	metatableName := flag.String("t", "", "Generate a LuaMetatableType method that returns the given value.  Takes the snake_case form of -s if unspecified.")
	outFile := flag.String("o", "", "The output file.  Defaults to lua_structname.go if empty.  Use '-' for stdout.")

	flag.Parse()

	if flag.NArg() != 0 {
		log.Fatal("gobindlua does not accept arguments")
	}

	if *structName == "" {
		log.Fatal("-s parameter is required")
	}

	if *metatableName == "" {
		*metatableName = gobindluautil.SnakeCase(*structName)
	}

	if *workingDir == "" {
		wd, err := os.Getwd()
		if err != nil {
			log.Fatalf("error getting working directory: %s", err)
		}
		*workingDir = wd
	}

	if *outFile == "" {
		*outFile = "lua_" + *structName + ".go"
	}

	var out io.Writer
	pathToOutput := ""

	if *outFile == "-" {
		out = os.Stdout
	} else {
		pathToOutput = filepath.Join(*workingDir, *outFile)
		f, err := os.Create(pathToOutput)

		if err != nil {
			log.Fatalf("error opening file for writing: %s", err)
		}

		defer f.Close()

		out = f
	}

	gen := NewGenerator(*structName, *workingDir, *metatableName)
	if err := gen.GenerateSourceCode(out, pathToOutput); err != nil {
		log.Fatalf("error generating struct: %s", err)
	}
}

type Generator struct {
	structToGenerate string
	wd               string
	metatableName    string

	packageSource *packages.Package
	structObject  types.Object

	StaticFunctions []functiontype.FunctionType
	UserDataMethods []functiontype.FunctionType

	imports               map[string]string
	metatableInitFunction bytes.Buffer
	metatableFunctions    bytes.Buffer
	userDataFunctions     bytes.Buffer
}

func NewGenerator(structToGenerate, wd string, metatableName string) *Generator {
	return &Generator{
		structToGenerate: structToGenerate,
		wd:               wd,
		metatableName:    metatableName,
	}
}

func (g *Generator) GenerateSourceCode(out io.Writer, pathToOutput string) error {
	if err := g.loadSourcePackage(); err != nil {
		return fmt.Errorf("failed to load imported go packages: %s", err)
	}

	g.StaticFunctions = g.gatherConstructors()
	g.UserDataMethods = g.gatherReceivers()

	g.imports = make(map[string]string)
	g.imports["github.com/yuin/gopher-lua"] = "lua"
	g.imports["github.com/ChrisTrenkamp/gobindlua"] = ""

	g.addPackageFromFunctions(g.StaticFunctions)
	g.addPackageFromFunctions(g.UserDataMethods)

	for _, i := range g.GatherFields(true) {
		g.addPackage(i.DataType)
	}

	g.buildMetatableInitFunction()
	g.buildMetatableFunctions()
	g.builderUserDataFunctions()

	code := bytes.Buffer{}
	fmt.Fprintf(&code, "// Code generated by gobindlua; DO NOT EDIT.\n")
	fmt.Fprintf(&code, "package %s\n\nimport (\n", g.packageSource.Types.Name())
	for pkg, name := range g.imports {
		fmt.Fprintf(&code, "\t%s \"%s\"\n", name, pkg)
	}
	fmt.Fprintf(&code, ")\n")
	io.Copy(&code, &g.metatableInitFunction)
	io.Copy(&code, &g.metatableFunctions)
	io.Copy(&code, &g.userDataFunctions)

	originalCodeBytes := code.Bytes()

	if pathToOutput == "" {
		_, err := io.Copy(out, bytes.NewBuffer(originalCodeBytes))
		return err
	}

	formattedCode, err := imports.Process(pathToOutput, originalCodeBytes, nil)
	if err != nil {
		_, werr := io.Copy(out, bytes.NewBuffer(originalCodeBytes))
		return errors.Join(err, werr)
	}

	_, werr := io.Copy(out, bytes.NewBuffer(formattedCode))
	return werr
}

func (g *Generator) addPackageFromFunctions(f []functiontype.FunctionType) {
	for _, i := range f {
		for _, p := range i.Params {
			g.addPackage(p.DataType)
		}

		for _, p := range i.Ret {
			g.addPackage(p)
		}
	}
}

func (g *Generator) addPackage(d datatype.DataType) {
	if p := d.Package(); p != "" && p != g.packageSource.ID {
		g.imports[p] = ""
	}
}

func (g *Generator) loadSourcePackage() error {
	config := &packages.Config{
		Mode: packages.NeedSyntax | packages.NeedTypes | packages.NeedTypesInfo,
	}
	pack, err := packages.Load(config, g.wd)
	if err != nil {
		return err
	}

	if len(pack) == 1 {
		g.packageSource = pack[0]
	} else {
		return fmt.Errorf("packages.Load returned more than one package")
	}

	g.structObject = g.packageSource.Types.Scope().Lookup(g.structToGenerate)

	if g.structObject == nil {
		return fmt.Errorf("specified type not found")
	}

	if _, ok := g.structObject.Type().Underlying().(*types.Struct); !ok {
		return fmt.Errorf("specified type is not a struct")
	}

	return nil
}

func (g *Generator) StructToGenerate() string {
	return g.structObject.Name()
}

func (g *Generator) StructMetatableIdentifier() string {
	return gobindluautil.SnakeCase(g.StructToGenerate())
}

func (g *Generator) StructMetatableFieldsIdentifier() string {
	return gobindluautil.SnakeCase(g.StructToGenerate() + "Fields")
}

func (g *Generator) UserDataCheckFn() string {
	return "luaCheck" + g.StructToGenerate()
}

func (g *Generator) SourceUserDataAccess() string {
	return "luaAccess" + g.StructToGenerate()
}

func (g *Generator) SourceUserDataSet() string {
	return "luaSet" + g.StructToGenerate()
}

func (g *Generator) buildMetatableInitFunction() {
	templ := `
func (goType {{ .StructToGenerate }}) RegisterLuaType(L *lua.LState) {
	staticMethodsTable := L.NewTypeMetatable("{{ .StructMetatableIdentifier }}")
	L.SetGlobal("{{ .StructMetatableIdentifier }}", staticMethodsTable)
	{{ range $idx, $fn := .StaticFunctions -}}
		L.SetField(staticMethodsTable, "{{ $fn.LuaFnName }}", L.NewFunction({{ $fn.SourceFnName }}))
	{{ end }}
	fieldsTable := L.NewTypeMetatable(goType.LuaMetatableType())
	L.SetGlobal(goType.LuaMetatableType(), fieldsTable)
	L.SetField(fieldsTable, "__index", L.NewFunction({{ .SourceUserDataAccess }}))
	L.SetField(fieldsTable, "__newindex", L.NewFunction({{ .SourceUserDataSet }}))
}
`

	execTempl(&g.metatableInitFunction, g, templ)
}

func (g *Generator) buildMetatableFunctions() {
	for _, i := range g.StaticFunctions {
		generateLuaFunctionWrapper(&g.metatableFunctions, g, i)
	}
}

func (g *Generator) builderUserDataFunctions() {
	g.generateUserDataPredefinitions()
	g.generateStructAccessFunction()
	g.generateStructSetFunction()
	g.generateStructMethods()
}

func (g *Generator) generateUserDataPredefinitions() {
	templ := `
func (r *{{ .StructToGenerate }}) LuaMetatableType() string {
	return "{{ .StructMetatableFieldsIdentifier }}"
}

func {{ .UserDataCheckFn }}(param int, L *lua.LState) *{{ .StructToGenerate }} {
	ud := L.CheckUserData(param)
	if v, ok := ud.Value.(*{{ .StructToGenerate }}); ok {
		return v
	}
	L.ArgError(1, "{{ .StructToGenerate }} expected")
	return nil
}
`

	execTempl(&g.userDataFunctions, g, templ)
}

func (g *Generator) generateStructAccessFunction() {
	templ := `
func {{ .SourceUserDataAccess }}(L *lua.LState) int {
	{{- if gt (len (.GatherFields false)) 0 }}
	p1 := {{ .UserDataCheckFn }}(1, L)
	{{- end }}
	p2 := L.CheckString(2)

	switch p2 {
		{{- range $idx, $field := .GatherFields false }}
	case "{{ $field.LuaName }}":
		L.Push({{ $field.DataType.ConvertGoTypeToLua (printf "p1.%s" $field.FieldName) }})
		{{ end -}}

		{{- range $idx, $method := .UserDataMethods }}
	case "{{ $method.LuaFnName }}":
		L.Push(L.NewFunction({{ $method.SourceFnName }}))
		{{ end -}}
	}

	return 1
}
`

	execTempl(&g.userDataFunctions, g, templ)
}

func (g *Generator) generateStructSetFunction() {
	templ := `
func {{ .SourceUserDataSet }}(L *lua.LState) int {
	{{- if gt (len (.GatherFields false)) 0 }}
	p1 := {{ .UserDataCheckFn }}(1, L)
	{{- end }}
	p2 := L.CheckString(2)

	switch p2 {
		{{- range $idx, $field := .GatherFields false }}
	case "{{ $field.LuaName }}":
		{{ $field.DataType.ConvertLuaTypeToGo "ud" (printf "%s(3)" $field.DataType.LuaParamType) 3 }}
		p1.{{ $field.FieldName }} = {{ $field.DataType.ReferenceOrDereferenceForAssignmentToField }}ud
		{{ end -}}
	}

	return 1
}
`

	execTempl(&g.userDataFunctions, g, templ)
}

func (g *Generator) generateStructMethods() {
	for _, i := range g.UserDataMethods {
		generateLuaFunctionWrapper(&g.userDataFunctions, g, i)
	}
}

func generateLuaFunctionWrapper(out io.Writer, g *Generator, f functiontype.FunctionType) {
	type FunctionGenerator struct {
		*Generator
		*functiontype.FunctionType
	}

	templ := `
func {{ .FunctionType.SourceFnName }}(L *lua.LState) int {
	{{ if .FunctionType.Receiver -}}
		r := {{ .Generator.UserDataCheckFn }}(1, L)
	{{- end }}
	{{ range $idx, $param := .Params }}
		var p{{ $idx }} {{ $param.TemplateArg }}
	{{ end }}
	{{ range $idx, $param := .Params }}
		{
			{{ $param.ConvertLuaTypeToGo "ud" (printf "%s(%d)" $param.LuaParamType $param.ParamNum) $param.ParamNum }}
			p{{ $idx }} = {{ $param.ReferenceOrDereferenceForAssignmentToField }}ud
		}
	{{ end }}
	{{ .FunctionType.GenerateReturnValues "r" }} := {{ if .FunctionType.Receiver -}}r.{{ end }}{{ .FunctionType.ActualFnName  }}({{ .FunctionType.GenerateParamValues "p" }})

	{{ range $idx, $ret := .Ret -}}
		{{- if $ret.IsError -}}
			if r{{ $idx }} != nil {
				L.Error(lua.LString(r{{ $idx }}.Error()), 1)
			}
		{{- end -}}
	{{- end }}

	{{ range $idx, $ret := .Ret -}}
		{{- if not $ret.IsError -}}
			{{ $name := printf "r%d" $idx }}
			L.Push({{$ret.ConvertGoTypeToLua $name}})
		{{- end -}}
	{{- end }}

	return {{ .FunctionType.NumReturns }}
}
`

	execTempl(out, FunctionGenerator{g, &f}, templ)
}

func execTempl(out io.Writer, data any, templ string) {
	t := template.Must(template.New("").Parse(templ))
	err := t.Execute(out, data)

	if err != nil {
		panic(err)
	}
}

func (g *Generator) gatherConstructors() []functiontype.FunctionType {
	ret := make([]functiontype.FunctionType, 0)
	underylingStructType := g.structObject.Type().Underlying()
	constructorPrefix := "New" + g.structToGenerate

	for _, syn := range g.packageSource.Syntax {
		for _, dec := range syn.Decls {
			if fn, ok := dec.(*ast.FuncDecl); ok && fn.Type.Results != nil {
				fnName := fn.Name.Name

				for _, retType := range fn.Type.Results.List {
					retType := datatype.CreateDataTypeFromExpr(retType.Type, g.packageSource)

					if retType.Type.Underlying() == underylingStructType && strings.HasPrefix(fnName, constructorPrefix) {
						luaName := "New" + fnName[len(constructorPrefix):]
						luaName = gobindluautil.SnakeCase(luaName)
						sourceCodeName := "luaConstructor" + g.StructToGenerate() + fnName
						ret = append(ret, functiontype.CreateFunction(fn, false, luaName, sourceCodeName, g.packageSource))
						break
					}
				}
			}
		}
	}

	return ret
}

func (g *Generator) gatherReceivers() []functiontype.FunctionType {
	ret := make([]functiontype.FunctionType, 0)
	underylingStructType := g.structObject.Type().Underlying()

	for _, syn := range g.packageSource.Syntax {
		for _, dec := range syn.Decls {
			if fn, ok := dec.(*ast.FuncDecl); ok && fn.Recv != nil {
				fnName := fn.Name.Name

				for _, recType := range fn.Recv.List {
					recType := datatype.CreateDataTypeFromExpr(recType.Type, g.packageSource)

					if recType.Type.Underlying() == underylingStructType && unicode.IsUpper(rune(fnName[0])) {
						luaName := gobindluautil.SnakeCase(fnName)
						sourceCodeName := "luaMethod" + g.StructToGenerate() + fnName
						ret = append(ret, functiontype.CreateFunction(fn, true, luaName, sourceCodeName, g.packageSource))
						break
					}
				}
			}
		}
	}

	return ret
}

func (g *Generator) GatherFields(gatherPrivate bool) []structfield.StructField {
	ret := make([]structfield.StructField, 0)
	str := g.structObject.Type().Underlying().(*types.Struct)

	for i := 0; i < str.NumFields(); i++ {
		field := str.Field(i)

		if gatherPrivate || field.Exported() {
			ret = append(ret, structfield.CreateStructField(field, g.packageSource))
		}
	}

	return ret
}