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;

typedef struct {
	unsigned char *buf;
	size_t size;
	size_t capacity;
} BytecodeBuffer;

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);
}

// Optimized bytecode writer with pre-allocated buffer
int buffered_bytecode_writer(lua_State *L, const void *p, size_t sz, void *ud) {
	BytecodeBuffer *buf = (BytecodeBuffer *)ud;

	// Grow buffer if needed (double size to avoid frequent reallocs)
	if (buf->size + sz > buf->capacity) {
		size_t new_capacity = buf->capacity;
		while (new_capacity < buf->size + sz) {
			new_capacity *= 2;
		}
		unsigned char *newbuf = realloc(buf->buf, new_capacity);
		if (newbuf == NULL) return 1;
		buf->buf = newbuf;
		buf->capacity = new_capacity;
	}

	memcpy(buf->buf + buf->size, p, sz);
	buf->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"
	"sync"
	"unsafe"
)

// bytecodeBuffer wraps []byte to avoid boxing allocations in sync.Pool
type bytecodeBuffer struct {
	data []byte
}

// Buffer pool for bytecode generation
var bytecodeBufferPool = sync.Pool{
	New: func() any {
		return &bytecodeBuffer{data: make([]byte, 0, 1024)}
	},
}

// 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)
	}

	// Always use C memory for dump operation to avoid cgo pointer issues
	cbuf := C.BytecodeBuffer{
		buf:      (*C.uchar)(C.malloc(1024)),
		size:     0,
		capacity: 1024,
	}
	if cbuf.buf == nil {
		return nil, fmt.Errorf("failed to allocate initial buffer")
	}

	// Dump the function to bytecode
	status := C.lua_dump(s.L, (*[0]byte)(unsafe.Pointer(C.buffered_bytecode_writer)), unsafe.Pointer(&cbuf))

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

	if status != 0 {
		C.free(unsafe.Pointer(cbuf.buf))
		return nil, fmt.Errorf("failed to dump bytecode: status %d", status)
	}

	// Copy to Go memory and free C buffer
	var result []byte
	if cbuf.size > 0 {
		result = C.GoBytes(unsafe.Pointer(cbuf.buf), C.int(cbuf.size))
	}
	C.free(unsafe.Pointer(cbuf.buf))

	return result, 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 := s.CreateLuaError(int(status), fmt.Sprintf("LoadBytecode(%s)", name))
		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
func (s *State) RunBytecodeWithResults(nresults int) error {
	status := C.lua_pcall(s.L, 0, C.int(nresults), 0)
	if status != 0 {
		err := s.CreateLuaError(int(status), fmt.Sprintf("RunBytecode(nresults=%d)", nresults))
		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))

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

	if status != 0 {
		err := s.CreateLuaError(int(status), fmt.Sprintf("LoadAndRunBytecode(%s)", name))
		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))

	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 := s.CreateLuaError(int(status), fmt.Sprintf("LoadAndRunBytecodeWithResults(%s, nresults=%d)", name, nresults))
		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
}