tinymux/mux/modules/engine/jit_lua.cpp

/*! \file jit_lua.cpp
 * \brief CJITCompile COM class — Lua bytecode → native JIT compilation.
 *
 * Implements mux_IJITCompile.  Deserializes Lua 5.4 bytecode,
 * lowers through HIR/RV64/x86-64 pipeline, caches compiled programs.
 */

#include "copyright.h"
#include "autoconf.h"
#include "config.h"
#include "externs.h"

#include "dbt_compile.h"
#include "engine_api.h"
#include "lua_bytecode.h"
#include "hir_lower_lua.h"

#include <atomic>
#include <cstring>
#include <cstdio>
#include <unordered_map>
#include <vector>
#include <string>

// ---------------------------------------------------------------
// Compile cache
// ---------------------------------------------------------------

static std::unordered_map<uint64_t, compiled_program> s_lua_cache;
static std::atomic<uint64_t> s_next_key{1};

// ---------------------------------------------------------------
// Statistics
// ---------------------------------------------------------------

struct lua_jit_stats {
    uint64_t compile_ok;
    uint64_t compile_fail;
    uint64_t run_ok;
    uint64_t run_fail;
    uint64_t cache_hits;
    uint64_t invalidations;
};

static lua_jit_stats s_lua_jit_stats = {};

// ---------------------------------------------------------------
// Compile a Lua bytecode blob to a compiled_program.
// ---------------------------------------------------------------

static bool compile_lua_bytecode(const uint8_t *data, size_t len,
                                  compiled_program *out) {
    // Deserialize.
    lua_bc_chunk chunk;
    if (!lua_bc_load(data, len, &chunk)) {
        return false;
    }

    // Fast eligibility check — reject before allocating HIR/RV64 state.
    lua_bc_reject reason = lua_bc_eligible(&chunk.main);
    if (reason != LUA_BC_ELIGIBLE) {
        return false;
    }

    // Create HIR program and RV64 compiler state.
    hir_program *h = new hir_program;
    h->init();

    rv_compiler rc_state;

    // Lower Lua bytecode to HIR.
    int result = hir_lower_lua_proto(*h, rc_state, &chunk.main);
    if (result < 0) {
        delete h;
        return false;
    }

    // Always build block ranges (block_last starts at -1 until CFG is
    // computed).  Without this, single-block Lua programs skip every HIR
    // insn in hir_codegen and leave final_out=0 (#1309 empty folds).
    //
    hir_build_cfg(*h);
    if (h->n_blocks > 1) {
        hir_ssa_construct(*h);
        hir_optimize(*h);
    } else {
        // Single block: just constant folding.
        hir_const_fold(*h);
    }

    // Code generation: HIR → RV64.
    hir_codegen(*h, rc_state);

    // Copy code into guest memory at code_base — same as softcode
    // compile_expression.  Without this, code lives only in rc_state.code
    // while memory's code region stays zero; compact then materializes
    // zeros at entry_pc and the DBT spins forever at guest PC 0
    // (unknown opcode → exit_with_pc(same) → #1309 hang).
    //
    if (rc_state.code.size() * 4 > rv_compiler::CODE_LIMIT) {
        delete h;
        return false;
    }
    for (size_t i = 0; i < rc_state.code.size(); i++) {
        memcpy(rc_state.memory.data() + rc_state.code_base + i * 4,
               &rc_state.code[i], 4);
    }
    if (rc_state.out_exhausted || rc_state.pool_exhausted) {
        delete h;
        return false;
    }

    // Build compiled_program output.
    // Include rc_state.needs_jit: codegen may force runtime for results that
    // live only in registers (ITOA path), even when lowering saw no ecalls.
    //
    out->memory = std::move(rc_state.memory);
    out->memory_size = rv_compiler::MEM_SIZE;
    out->out_addr = rc_state.final_out;
    out->out_used = rc_state.out_pool;
    out->entry_pc = rc_state.code_base;
    out->code_size = rc_state.code.size() * 4;
    out->str_pool_end = rc_state.str_pool;
    out->fargs_pool_end = rc_state.fargs_pool;
    out->out_pool_end = rc_state.out_pool;
    out->ok = true;
    out->folds = h->folds;
    out->ecalls = h->ecalls;
    out->tier2_calls = 0;
    out->native_ops = h->native_ops;
    out->needs_jit = h->needs_jit || rc_state.needs_jit;

    // Classify CARGS/SUBST refs (mux.args[N] → emit_sref) so
    // run_cached_program populates the guest slots this program reads.
    // Same pass as softcode compile_expression.
    //
    out->subst_mask = 0;
    out->cargs_used = 0;
    for (uint64_t a : h->sref_addrs) {
        if (  a >= rv_compiler::CARGS_BASE
           && a < rv_compiler::CARGS_BASE
                  + static_cast<uint64_t>(rv_compiler::MAX_CARGS)
                    * rv_compiler::CARGS_SLOT) {
            int idx = static_cast<int>(
                (a - rv_compiler::CARGS_BASE) / rv_compiler::CARGS_SLOT);
            if (idx + 1 > out->cargs_used) {
                out->cargs_used = idx + 1;
            }
        } else if (  a >= rv_compiler::SUBST_BASE
                  && a < rv_compiler::SUBST_BASE
                         + static_cast<uint64_t>(rv_compiler::SUBST_COUNT)
                           * rv_compiler::SUBST_SLOT) {
            int slot = static_cast<int>(
                (a - rv_compiler::SUBST_BASE) / rv_compiler::SUBST_SLOT);
            out->subst_mask |= (UINT64_C(1) << slot);
        }
    }

    delete h;
    return true;
}

// ---------------------------------------------------------------
// CJITCompile COM class
// ---------------------------------------------------------------

class CJITCompile : public mux_IJITCompile
{
public:
    CJITCompile(void) : m_cRef(1) {}
    virtual ~CJITCompile() {}

    // mux_IUnknown
    MUX_RESULT QueryInterface(MUX_IID iid, void **ppv) override {
        if (mux_IID_IUnknown == iid) {
            *ppv = static_cast<mux_IUnknown *>(static_cast<mux_IJITCompile *>(this));
        } else if (IID_IJITCompile == iid) {
            *ppv = static_cast<mux_IJITCompile *>(this);
        } else {
            *ppv = nullptr;
            return MUX_E_NOINTERFACE;
        }
        AddRef();
        return MUX_S_OK;
    }

    uint32_t AddRef(void) override { return ++m_cRef; }
    uint32_t Release(void) override {
        uint32_t n = --m_cRef;
        if (0 == n) delete this;
        return n;
    }

    // mux_IJITCompile
    MUX_RESULT CompileLuaBytecode(const uint8_t *pData, size_t nData,
        uint64_t *pKey) override
    {
        if (nullptr == pData || nullptr == pKey) return MUX_E_INVALIDARG;

        // Persistent cache key: "lua:" + sha1(bytecodes).
        std::string cache_key = "lua:" + jit_sha1_hex(pData, nData);

        // Check SQLite cache first.
        compiled_program prog;
        bool from_cache = jit_load_from_sqlite(cache_key, prog);

        if (!from_cache) {
            // Cache miss — compile from bytecodes.
            if (!compile_lua_bytecode(pData, nData, &prog)) {
                s_lua_jit_stats.compile_fail++;
                *pKey = 0;
                return MUX_E_FAIL;
            }
            // Persist to SQLite while prog.memory still exists.
            jit_store_to_sqlite(cache_key, prog);
            // Compact: extract blobs, release 4MB memory.
            jit_compact_program(prog);
        }

        uint64_t key = s_next_key.fetch_add(1, std::memory_order_relaxed);
        s_lua_cache[key] = std::move(prog);
        *pKey = key;
        s_lua_jit_stats.compile_ok++;
        return MUX_S_OK;
    }

    MUX_RESULT RunCompiled(uint64_t key,
        dbref executor, dbref caller, dbref enactor,
        const UTF8 *pArgs[], int nArgs,
        UTF8 *pResult, size_t nResultMax, size_t *pnResultLen,
        void *pLuaState) override
    {
        auto it = s_lua_cache.find(key);
        if (it == s_lua_cache.end()) return MUX_E_NOTFOUND;

        s_lua_jit_stats.cache_hits++;

        // On a wall-clock alarm mid-run, run_cached_program writes
        // "#-1 CPU LIMITED" into pResult and returns true (handled) rather
        // than false.  That is deliberate: a false here returns MUX_E_FAIL,
        // and the Lua caller responds by re-running the whole chunk in the
        // Lua VM -- expensive work we must NOT do once the command is already
        // over its CPU budget.  So an alarm surfaces as a successful
        // CPU-LIMITED result, not a failover.
        //
        bool ok = run_cached_program(&it->second, executor, caller, enactor,
            pResult, nResultMax, pArgs, nArgs,
            EV_FCHECK | EV_EVAL, pLuaState);
        if (!ok) {
            s_lua_jit_stats.run_fail++;
            return MUX_E_FAIL;
        }

        if (pnResultLen) {
            *pnResultLen = strlen(reinterpret_cast<const char *>(pResult));
        }
        s_lua_jit_stats.run_ok++;
        return MUX_S_OK;
    }

    MUX_RESULT IsCompiled(uint64_t key, bool *pCompiled) override {
        if (nullptr == pCompiled) return MUX_E_INVALIDARG;
        *pCompiled = (s_lua_cache.find(key) != s_lua_cache.end());
        return MUX_S_OK;
    }

    MUX_RESULT Invalidate(uint64_t key) override {
        auto it = s_lua_cache.find(key);
        if (it != s_lua_cache.end()) {
            s_lua_cache.erase(it);
            s_lua_jit_stats.invalidations++;
        }
        return MUX_S_OK;
    }

private:
    uint32_t m_cRef;
};

// ---------------------------------------------------------------
// Factory creation function — called from engine_com.cpp.
// ---------------------------------------------------------------

MUX_RESULT jit_compile_create_instance(MUX_IID iid, void **ppv) {
    CJITCompile *pObj = nullptr;
    try { pObj = new CJITCompile; } catch (...) { ; }
    if (nullptr == pObj) return MUX_E_OUTOFMEMORY;

    MUX_RESULT mr = pObj->QueryInterface(iid, ppv);
    pObj->Release();
    return mr;
}