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_chunk(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;
} BytecodeWriter;

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

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

    memcpy(newbuf + w->len, p, sz);
    w->buf = newbuf;
    w->len += sz;
    return 0;
}
*/
import "C"
import (
	"fmt"
	"unsafe"
)

func (s *State) compileBytecodeUnsafe(code string, name string) ([]byte, error) {
	// First load the string but don't execute it
	ccode := C.CString(code)
	defer C.free(unsafe.Pointer(ccode))

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

	if C.luaL_loadstring(s.L, ccode) != 0 {
		err := &LuaError{
			Code:    int(C.lua_status(s.L)),
			Message: s.ToString(-1),
		}
		s.Pop(1)
		return nil, fmt.Errorf("failed to load string: %w", err)
	}

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

	// Dump the function to bytecode
	if C.lua_dump(s.L, (*[0]byte)(C.bytecode_writer), unsafe.Pointer(&writer)) != 0 {
		if writer.buf != nil {
			C.free(unsafe.Pointer(writer.buf))
		}
		s.Pop(1)
		return nil, fmt.Errorf("failed to dump bytecode")
	}

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

	// Clean up
	if writer.buf != nil {
		C.free(unsafe.Pointer(writer.buf))
	}
	s.Pop(1) // Remove the function

	return bytecode, nil
}

func (s *State) loadBytecodeUnsafe(bytecode []byte, name string) error {
	cname := C.CString(name)
	defer C.free(unsafe.Pointer(cname))

	// Load the bytecode
	status := C.load_bytecode_chunk(
		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)
		return fmt.Errorf("failed to load bytecode: %w", err)
	}

	// Execute the loaded chunk
	if err := s.safeCall(func() C.int {
		return C.lua_pcall(s.L, 0, 0, 0)
	}); err != nil {
		return fmt.Errorf("failed to execute bytecode: %w", err)
	}

	return nil
}

// CompileBytecode compiles a Lua chunk to bytecode without executing it
func (s *State) CompileBytecode(code string, name string) ([]byte, error) {
	if s.safeStack {
		return stackGuardValue[[]byte](s, func() ([]byte, error) {
			return s.compileBytecodeUnsafe(code, name)
		})
	}
	return s.compileBytecodeUnsafe(code, name)
}

// LoadBytecode loads precompiled bytecode and executes it
func (s *State) LoadBytecode(bytecode []byte, name string) error {
	if s.safeStack {
		return stackGuardErr(s, func() error {
			return s.loadBytecodeUnsafe(bytecode, name)
		})
	}
	return s.loadBytecodeUnsafe(bytecode, name)
}

// Helper function to compile and immediately load/execute bytecode
func (s *State) CompileAndLoad(code string, name string) error {
	bytecode, err := s.CompileBytecode(code, name)
	if err != nil {
		return fmt.Errorf("compile error: %w", err)
	}

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

	return nil
}