path: root/lua.go

//go:build puc

package main

/*
#cgo LDFLAGS: -llua
#include <stdlib.h>
#include <stdint.h>
#include <lua.h>
#include <lualib.h>
#include <lauxlib.h>

extern int luna_run_go_func(uintptr_t);

static inline int luna_run(lua_State *l) {
	return luna_run_go_func(lua_tonumber(l, lua_upvalueindex(1)));
}

static inline void luna_push_function(lua_State *l, uintptr_t f) {
	lua_pushinteger(l, f);
	lua_pushcclosure(l, luna_run, 1);
}
*/
import "C"
import (
	"errors"
	"fmt"
	"reflect"
	"runtime/cgo"
	"slices"
	"time"
	"unsafe"
)

type LuaRef C.longlong
// Lua is a wrapper around C Lua state with several conveniences.
type Lua struct {
	l *C.lua_State
}

//export luna_run_go_func
// luna_run_go_func calls a Go function from C/Lua.
func luna_run_go_func(f C.uintptr_t) C.int {
	fn := cgo.Handle(f).Value().(func() int)
	return C.int(fn())
}

// Start opens the Lua context with all built-in libraries.
func (l *Lua) Start() {
	l.l = C.luaL_newstate()
	C.luaL_openlibs(l.l)
}

// Close closes the Lua context.
func (l *Lua) Close() {
	C.lua_close(l.l)
}

// Require loads and executes the file pushing results onto the Lua stack.
func (l *Lua) Require(file string) error {
	err := l.LoadFile(file)
	if err != nil {
		return errors.New("could not open the file:\n" + err.Error())
	}
	err = l.PCall(0, C.LUA_MULTRET)
	if err != nil {
		return errors.New("could not execute the file:\n" + err.Error())
	}
	return nil
}

// PCall calls a function with the stack index (-nargs-1) expecting nresults
// number of results.
func (l *Lua) PCall(nargs int, nresults int) error {
	res := C.lua_pcallk(l.l, C.int(nargs), C.int(nresults), 0, 0, nil)
	if res != C.LUA_OK {
		errMsg := l.ToString(-1)
		l.Pop(1)
		return errors.New(errMsg)
	}
	return nil
}

// LoadFile loads the file code in the current Lua context.
func (l *Lua) LoadFile(file string) error {
	cstr := C.CString(file)
	defer C.free(unsafe.Pointer(cstr))
	res := C.luaL_loadfilex(l.l, cstr, nil)
	if res != C.LUA_OK {
		errMsg := l.ToString(-1)
		l.Pop(1)
		return errors.New(errMsg)
	}
	return nil
}

// StackLen returns the length of the Lua stack.
func (l *Lua) StackLen() int {
	return int(C.lua_gettop(l.l))
}

// ToInt converts a stack value with stack index into a string.
func (l *Lua) ToString(index int) string {
	return C.GoString(C.lua_tolstring(l.l, C.int(index), nil))
}

// ToInt converts a stack value with stack index -1 into a number.
func (l *Lua) ToInt(index int) int {
	return int(C.lua_tonumberx(l.l, C.int(index), nil))
}

// PopToRef pops the item (stack offset -1) from the stack and then stores it
// in the registry.
func (l *Lua) PopToRef() LuaRef {
	return LuaRef(C.luaL_ref(l.l, C.LUA_REGISTRYINDEX))
}

// Pop pops N items from the Lua stack.
func (l *Lua) Pop(n int) {
	C.lua_settop(l.l, C.int(-n - 1))
}

// PushNumber pushes tnil onto the Lua stack.
func (l *Lua) PushNil() {
	C.lua_pushnil(l.l)
}

// PushNumber pushes the number onto the Lua stack.
func (l *Lua) PushNumber(num int) {
	C.lua_pushnumber(l.l, C.double(num))
}

// PushNumber pushes the number onto the Lua stack.
func (l *Lua) PushFloatNumber(num float64) {
	C.lua_pushnumber(l.l, C.double(num))
}

// PushString pushes the string onto the Lua stack.
func (l *Lua) PushString(str string) {
	cstr := C.CString(str)
	defer C.free(unsafe.Pointer(cstr))
	C.lua_pushstring(l.l, cstr)
}

// PushGoFunction pushes a Go function onto the Lua stack so it's possible to
// call it from Lua. You may need to delete returned handle with .Delete()
// method later.
func (l *Lua) PushGoFunction(f func (l *Lua) int) cgo.Handle {
	h := cgo.NewHandle(func () int {
		return f(l)
	})
	C.luna_push_function(l.l, C.uintptr_t(h))
	return h
}

// CreateTable pushes a new Lua table onto the stack.
func (l *Lua) CreateTable(len int) {
	C.lua_createtable(l.l, 0, C.int(len))
}

// SetTableItem adds a value (stack offset -1) to a table (offset -2) under
// the specified key.
func (l *Lua) SetTableItem(key string) {
	cstr := C.CString(key)
	defer C.free(unsafe.Pointer(cstr))
	C.lua_setfield(l.l, -2, C.CString(key))
}

// PushAny pushes value v onto the stack.
func (l *Lua) PushAny(v any) error {
	switch v.(type) {
	case string:
		v, _ := v.(string)
		l.PushString(v)
	case func (l *Lua) int:
		v, _ := v.(func (l *Lua) int)
		l.PushGoFunction(v)
	case int:
		v, _ := v.(int)
		l.PushNumber(v)
	case int64:
		v, _ := v.(int64)
		l.PushNumber(int(v))
	case float64:
		v, _ := v.(float64)
		l.PushFloatNumber(v)
	case map[string]any:
		v, _ := v.(map[string]any)
		l.PushObject(v)
	case []any:
		v, _ := v.([]any)
		l.PushArray(v)
	case time.Time:
		v, _ := v.(time.Time)
		l.PushString(v.Format(time.DateTime))
	default:
		return fmt.Errorf("unsupported value type: %T", v)
	}
	return nil
}

// PushObject recursively pushes string->any Go table onto the stack.
func (l *Lua) PushObject(table map[string]any) error {
	l.CreateTable(len(table))
	for k, v := range table {
		err := l.PushAny(v)
		if err != nil {
			return err
		}
		l.SetTableItem(k)
	}
	return nil
}

// PushArray recursively pushes an array of Go values onto the stack.
func (l *Lua) PushArray(array []any) error {
	l.CreateTable(len(array))
	for k, v := range array {
		l.PushNumber(k + 1)
		err := l.PushAny(v)
		if err != nil {
			return err
		}
		C.lua_settable(l.l, C.int(-3))
	}
	return nil
}

// PushFromRef pushes a value from registry ref onto the stack.
func (l *Lua) PushFromRef(ref LuaRef) {
	C.lua_rawgeti(l.l, C.LUA_REGISTRYINDEX, C.longlong(ref));
}

// Type returns type of the value sitting at n index on the stack.
func (l *Lua) Type(n int) int {
	return int(C.lua_type(l.l, C.int(n)))
}

// IsNil checks if the stack contains nil under the given index.
func (l *Lua) IsNil(index int) bool {
	return C.lua_type(l.l, C.int(index)) == C.LUA_TNIL
}

// IsString checks if the stack contains a string under the given index.
func (l *Lua) IsString(index int) bool {
	return C.lua_isstring(l.l, C.int(index)) == 1
}

// IsNumber checks if the stack contains a number (both int and float) under
// the given index.
func (l *Lua) IsNumber(index int) bool {
	return C.lua_isnumber(l.l, C.int(index)) == 1
}

// IsFunction checks if the stack contains a function under the given index.
func (l *Lua) IsFunction(index int) bool {
	return C.lua_type(l.l, C.int(index)) == C.LUA_TFUNCTION
}

// IsTable checks if the stack contains a table under the given index.
func (l *Lua) IsTable(index int) bool {
	return C.lua_type(l.l, C.int(index)) == C.LUA_TTABLE
}

// Next advances a table iterator at the stack index -2.
func (l *Lua) Next() bool {
	return C.lua_next(l.l, -2) != 0
}

// LoadString loads and executes the code in the current Lua context.
func (l *Lua) LoadString(code string) error {
	cstr := C.CString(code)
	defer C.free(unsafe.Pointer(cstr))
	if C.luaL_loadstring(l.l, cstr) != C.LUA_OK {
		errMsg := l.ToString(-1)
		l.Pop(1)
		return errors.New(errMsg)
	}
	err := l.PCall(0, 0)
	if err != nil {
		return err
	}
	l.Pop(l.StackLen())
	return nil
}

// Scan scans values from the Lua stack into vars according to their types.
func (l *Lua) Scan(vars ...any) error {
	slices.Reverse(vars)
	for i, v := range vars {
		t := reflect.TypeOf(v)
		// unwrap pointers
		for t.Kind() == reflect.Ptr {
			t = t.Elem()
		}
		tk := t.Kind()
		if t.Name() == "LuaRef" {
			if l.IsFunction(-1) == false {
				return fmt.Errorf(
					"passed arg #%d must be function",
					len(vars)-i,
				)
			}
			*v.(*LuaRef) = l.PopToRef()
		} else if t.String() == "cgo.Handle" {
			if l.IsNumber(-1) == false {
				return fmt.Errorf(
					"passed arg #%d must be Go handler",
					len(vars)-i,
				)
			}
			*v.(*cgo.Handle) = cgo.Handle(uintptr(l.ToInt(-1)))
			l.Pop(1)
		} else if tk == reflect.String {
			if l.IsString(-1) == false {
				return fmt.Errorf(
					"passed arg #%d must be string",
					len(vars)-i,
				)
			}
			*v.(*string) = l.ToString(-1)
			l.Pop(1)
		} else if tk == reflect.Int {
			if l.IsNumber(-1) == false {
				return fmt.Errorf(
					"passed arg #%d must be number",
					len(vars)-i,
				)
			}
			*v.(*int) = l.ToInt(-1)
			l.Pop(1)
		} else if tk == reflect.Map && t.Key().Kind() == reflect.String {
			// TODO: should be possible to use maps with any types
			// of value, not only strings.
			vm, _ := v.(*map[string]string)
			l.PushNil()
			for l.Next() {
				if !l.IsString(-1) {
					return fmt.Errorf(
						"map arg #%d must only have string values",
						len(vars)-i,
					)
				}
				if !l.IsString(-2) {
					return fmt.Errorf(
						"map arg #%d must only have string keys",
						len(vars)-i,
					)
				}
				v := l.ToString(-1)
				l.Pop(1)
				// We must not pop the item key from the stack
				// because otherwise C.lua_next won't work
				// properly.
				k := l.ToString(-1)
				(*vm)[k] = v
			}
			l.Pop(1)
		} else if tk == reflect.Slice &&
			t.Elem().Kind() == reflect.Interface {
			if !l.IsTable(-1) {
				return fmt.Errorf(
					"passed arg #%d must be a table",
					len(vars)-i,
				)
			}
			va, _ := v.(*[]any)
			l.PushNil()
			for l.Next() {
				var v any = nil
				if l.IsString(-1) {
					v = l.ToString(-1)
				} else if l.IsNumber(-1) {
					v = l.ToInt(-1)
				} else if l.IsNil(-1) {
					v = nil
				} else {
					return fmt.Errorf("unknown value in array")
				}
				l.Pop(1)
				(*va) = append((*va), v)
			}
			l.Pop(1)
		} else {
			return fmt.Errorf("unknown var type: %s", t)
		}
	}
	return nil
}

// RestoreStackFunc remembers the Lua stack size and then restores it when a
// returned function is called. It's a helper function to avoid stack leakage.
func (l *Lua) RestoreStackFunc() func () {
	before := l.StackLen()
	return func () {
		after := l.StackLen()
		diff := after - before
		if diff == 0 {
			return
		} else if diff < 0 {
			msg := fmt.Sprintf(
				"too many stack pops: len before: %d, after: %d\n",
				before,
				after,
			)
			panic(msg)
		}
		C.lua_settop(l.l, C.int(before))
	}
}

// SetGlobal sets a global value at the -1 stack index with the name.
func (l *Lua) SetGlobal(name string) {
	cstr := C.CString(name)
	defer C.free(unsafe.Pointer(cstr))
	C.lua_setglobal(l.l, cstr)
}

// GetGlobal gets a global value under the name.
func (l *Lua) GetGlobal(name string) {
	cstr := C.CString(name)
	defer C.free(unsafe.Pointer(cstr))
	C.lua_getglobal(l.l, cstr)
}