LuaJIT-to-Go/bytecode.go

package luajit

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

typedef struct {
    const unsigned char *buf;
    size_t size;
    const char *name;
} BytecodeReader;

const char *bytecode_reader(lua_State *L, void *ud, size_t *size) {
    BytecodeReader *r = (BytecodeReader *)ud;
    (void)L;  // unused
    if (r->size == 0) return NULL;
    *size = r->size;
    r->size = 0;  // Only read once
    return (const char *)r->buf;
}

int load_bytecode(lua_State *L, const unsigned char *buf, size_t len, const char *name) {
    BytecodeReader reader = {buf, len, name};
    return lua_load(L, bytecode_reader, &reader, name);
}

// Direct bytecode dumping without intermediate buffer - more efficient
int direct_bytecode_writer(lua_State *L, const void *p, size_t sz, void *ud) {
    void **data = (void **)ud;
    size_t current_size = (size_t)data[1];
    void *newbuf = realloc(data[0], current_size + sz);
    if (newbuf == NULL) return 1;

    memcpy((unsigned char*)newbuf + current_size, p, sz);
    data[0] = newbuf;
    data[1] = (void*)(current_size + sz);
    return 0;
}

// Combined load and run bytecode in a single call
int load_and_run_bytecode(lua_State *L, const unsigned char *buf, size_t len,
                         const char *name, int nresults) {
    BytecodeReader reader = {buf, len, name};
    int status = lua_load(L, bytecode_reader, &reader, name);
    if (status != 0) return status;

    return lua_pcall(L, 0, nresults, 0);
}
*/
import "C"
import (
	"fmt"
	"unsafe"
)

// CompileBytecode compiles a Lua chunk to bytecode without executing it
func (s *State) CompileBytecode(code string, name string) ([]byte, error) {
	if err := s.LoadString(code); err != nil {
		return nil, fmt.Errorf("failed to load string: %w", err)
	}

	// Use a simpler direct writer with just two pointers
	data := [2]unsafe.Pointer{nil, nil}

	// Dump the function to bytecode
	status := C.lua_dump(s.L, (*[0]byte)(unsafe.Pointer(C.direct_bytecode_writer)), unsafe.Pointer(&data))
	if status != 0 {
		return nil, fmt.Errorf("failed to dump bytecode: status %d", status)
	}

	// Get result
	var bytecode []byte
	if data[0] != nil {
		// Create Go slice that references the C memory
		length := uintptr(data[1])
		bytecode = C.GoBytes(data[0], C.int(length))
		C.free(data[0])
	}

	s.Pop(1) // Remove the function from stack

	return bytecode, nil
}

// LoadBytecode loads precompiled bytecode without executing it
func (s *State) LoadBytecode(bytecode []byte, name string) error {
	if len(bytecode) == 0 {
		return fmt.Errorf("empty bytecode")
	}

	cname := C.CString(name)
	defer C.free(unsafe.Pointer(cname))

	// Load the bytecode
	status := C.load_bytecode(
		s.L,
		(*C.uchar)(unsafe.Pointer(&bytecode[0])),
		C.size_t(len(bytecode)),
		cname,
	)

	if status != 0 {
		err := &LuaError{
			Code:    int(status),
			Message: s.ToString(-1),
		}
		s.Pop(1) // Remove error message
		return err
	}

	return nil
}

// RunBytecode executes previously loaded bytecode with 0 results
func (s *State) RunBytecode() error {
	return s.RunBytecodeWithResults(0)
}

// RunBytecodeWithResults executes bytecode and keeps nresults on the stack
// Use LUA_MULTRET (-1) to keep all results
func (s *State) RunBytecodeWithResults(nresults int) error {
	status := C.lua_pcall(s.L, 0, C.int(nresults), 0)
	if status != 0 {
		err := &LuaError{
			Code:    int(status),
			Message: s.ToString(-1),
		}
		s.Pop(1) // Remove error message
		return err
	}
	return nil
}

// LoadAndRunBytecode loads and executes bytecode in a single CGO transition
func (s *State) LoadAndRunBytecode(bytecode []byte, name string) error {
	if len(bytecode) == 0 {
		return fmt.Errorf("empty bytecode")
	}

	cname := C.CString(name)
	defer C.free(unsafe.Pointer(cname))

	// Use combined load and run function
	status := C.load_and_run_bytecode(
		s.L,
		(*C.uchar)(unsafe.Pointer(&bytecode[0])),
		C.size_t(len(bytecode)),
		cname,
		0, // No results
	)

	if status != 0 {
		err := &LuaError{
			Code:    int(status),
			Message: s.ToString(-1),
		}
		s.Pop(1) // Remove error message
		return err
	}

	return nil
}

// LoadAndRunBytecodeWithResults loads and executes bytecode, preserving results
func (s *State) LoadAndRunBytecodeWithResults(bytecode []byte, name string, nresults int) error {
	if len(bytecode) == 0 {
		return fmt.Errorf("empty bytecode")
	}

	cname := C.CString(name)
	defer C.free(unsafe.Pointer(cname))

	// Use combined load and run function
	status := C.load_and_run_bytecode(
		s.L,
		(*C.uchar)(unsafe.Pointer(&bytecode[0])),
		C.size_t(len(bytecode)),
		cname,
		C.int(nresults),
	)

	if status != 0 {
		err := &LuaError{
			Code:    int(status),
			Message: s.ToString(-1),
		}
		s.Pop(1) // Remove error message
		return err
	}

	return nil
}

// CompileAndRun compiles and immediately executes Lua code
func (s *State) CompileAndRun(code string, name string) error {
	// Skip bytecode step for small scripts - direct execution is faster
	if len(code) < 1024 {
		return s.DoString(code)
	}

	bytecode, err := s.CompileBytecode(code, name)
	if err != nil {
		return fmt.Errorf("compile error: %w", err)
	}

	if err := s.LoadAndRunBytecode(bytecode, name); err != nil {
		return fmt.Errorf("execution error: %w", err)
	}

	return nil
}