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;

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

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

typedef struct {
    unsigned char *buf;
    size_t len;
    size_t capacity;
} BytecodeWriter;

static int bytecode_writer(lua_State *L, const void *p, size_t sz, void *ud) {
    BytecodeWriter *w = (BytecodeWriter *)ud;
    unsigned char *newbuf;
    (void)L;  // unused

    // Check if we need to reallocate
    if (w->len + sz > w->capacity) {
        size_t new_capacity = w->capacity * 2;
        if (new_capacity < w->len + sz) {
            new_capacity = w->len + sz;
        }

        newbuf = (unsigned char *)realloc(w->buf, new_capacity);
        if (newbuf == NULL) return 1;

        w->buf = newbuf;
        w->capacity = new_capacity;
    }

    memcpy(w->buf + w->len, p, sz);
    w->len += sz;
    return 0;
}

// Wrapper function that calls lua_dump with bytecode_writer
static int dump_lua_function(lua_State *L, BytecodeWriter *w) {
    return lua_dump(L, bytecode_writer, w);
}
*/
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)
	}

	// Set up writer with initial capacity
	var writer C.BytecodeWriter
	writer.buf = nil
	writer.len = 0
	writer.capacity = 0

	// Initial allocation with a reasonable size
	const initialSize = 4096
	writer.buf = (*C.uchar)(C.malloc(initialSize))
	if writer.buf == nil {
		s.Pop(1) // Remove the loaded function
		return nil, fmt.Errorf("failed to allocate memory for bytecode")
	}
	writer.capacity = initialSize

	// Dump the function to bytecode
	err := s.safeCall(func() C.int {
		return C.dump_lua_function(s.L, (*C.BytecodeWriter)(unsafe.Pointer(&writer)))
	})

	// Copy bytecode to Go slice regardless of the result
	bytecode := C.GoBytes(unsafe.Pointer(writer.buf), C.int(writer.len))

	// Clean up
	C.free(unsafe.Pointer(writer.buf))
	s.Pop(1) // Remove the function from stack

	if err != nil {
		return nil, fmt.Errorf("failed to dump bytecode: %w", err)
	}

	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
	err := s.safeCall(func() C.int {
		return C.load_bytecode(
			s.L,
			(*C.uchar)(unsafe.Pointer(&bytecode[0])),
			C.size_t(len(bytecode)),
			cname,
		)
	})

	if err != nil {
		return fmt.Errorf("failed to load bytecode: %w", 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 {
	return s.safeCall(func() C.int {
		return C.lua_pcall(s.L, 0, C.int(nresults), 0)
	})
}

// LoadAndRunBytecode loads and executes bytecode
func (s *State) LoadAndRunBytecode(bytecode []byte, name string) error {
	if err := s.LoadBytecode(bytecode, name); err != nil {
		return err
	}
	return s.RunBytecode()
}

// LoadAndRunBytecodeWithResults loads and executes bytecode, preserving results
func (s *State) LoadAndRunBytecodeWithResults(bytecode []byte, name string, nresults int) error {
	if err := s.LoadBytecode(bytecode, name); err != nil {
		return err
	}
	return s.RunBytecodeWithResults(nresults)
}

// CompileAndRun compiles and immediately executes Lua code
func (s *State) CompileAndRun(code string, name string) error {
	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
}