LuaJIT-to-Go/functions.go

package luajit

/*
#include <lua.h>
#include <lauxlib.h>
#include <stdlib.h>

extern int goFunctionWrapper(lua_State* L);

// Helper function to access upvalues
static int get_upvalue_index(int i) {
    return lua_upvalueindex(i);
}
*/
import "C"
import (
	"fmt"
	"sync"
	"unsafe"
)

// GoFunction defines the signature for Go functions callable from Lua
type GoFunction func(*State) int

var (
	// functionRegistry stores all registered Go functions
	functionRegistry = struct {
		sync.RWMutex
		funcs map[unsafe.Pointer]GoFunction
	}{
		funcs: make(map[unsafe.Pointer]GoFunction),
	}
)

//export goFunctionWrapper
func goFunctionWrapper(L *C.lua_State) C.int {
	state := &State{L: L}

	// Get function pointer from the first upvalue
	ptr := C.lua_touserdata(L, C.get_upvalue_index(1))
	if ptr == nil {
		state.PushString("error: function pointer not found")
		return -1
	}

	functionRegistry.RLock()
	fn, ok := functionRegistry.funcs[ptr]
	functionRegistry.RUnlock()

	if !ok {
		state.PushString("error: function not found in registry")
		return -1
	}

	// Call the Go function
	return C.int(fn(state))
}

// PushGoFunction wraps a Go function and pushes it onto the Lua stack
func (s *State) PushGoFunction(fn GoFunction) error {
	// Allocate a pointer to use as the function key
	ptr := C.malloc(1)
	if ptr == nil {
		return fmt.Errorf("failed to allocate memory for function pointer")
	}

	// Register the function
	functionRegistry.Lock()
	functionRegistry.funcs[ptr] = fn
	functionRegistry.Unlock()

	// Push the pointer as lightuserdata (first upvalue)
	C.lua_pushlightuserdata(s.L, ptr)

	// Create closure with the C wrapper and the upvalue
	C.lua_pushcclosure(s.L, (*[0]byte)(C.goFunctionWrapper), 1)

	return nil
}

// RegisterGoFunction registers a Go function as a global Lua function
func (s *State) RegisterGoFunction(name string, fn GoFunction) error {
	if err := s.PushGoFunction(fn); err != nil {
		return err
	}

	s.SetGlobal(name)
	return nil
}

// UnregisterGoFunction removes a global function
func (s *State) UnregisterGoFunction(name string) {
	s.PushNil()
	s.SetGlobal(name)
}

// Cleanup frees all function pointers and clears the registry
func (s *State) Cleanup() {
	functionRegistry.Lock()
	defer functionRegistry.Unlock()

	// Free all allocated pointers
	for ptr := range functionRegistry.funcs {
		C.free(ptr)
		delete(functionRegistry.funcs, ptr)
	}
}