vj-codegen 0.1.0

//! Task 24 (`docs/superpowers/plans/2026-09-23-wave-A-codegen.md`, ledger
//! CG-062, CG-064, CG-124): the production JIT backend -- a small, safe
//! wrapper around `llvm-sys` 221's ORC v2 `LLJIT` C API
//! (`llvm_sys::orc2`/`llvm_sys::orc2::lljit`).
//!
//! **BINDING CORRECTION 4:** inkwell 0.10 ships no `orc2` module at all, so
//! this crate's first hand-written FFI surface lives here. `JitBackend` owns
//! one shared `LLJIT`/main `JITDylib` pair that accepts modules from
//! multiple threads concurrently, with no external `Mutex` -- ORC v2 is
//! explicitly designed for concurrent module addition (`docs/design/
//! codegen.md` §4), which is the actual safety argument behind this file's
//! `unsafe impl Send`/`Sync`, not merely "it compiles". Every `unsafe` block
//! below carries a `// SAFETY:` comment (`tests/jit_backend.rs`'s
//! `unsafe_blocks_in_orc_module_are_all_safety_commented` enforces this
//! mechanically).
//!
//! **Deviation from the plan's literal API list:** the plan names
//! `LLVMOrcThreadSafeContextGetContext` as the way `add_module_from_bitcode`
//! recovers a raw `LLVMContextRef` to parse bitcode into inside a fresh
//! `ThreadSafeContext`. That symbol does not exist anywhere in the real
//! `llvm-sys = "221"` generated bindings or the LLVM 22.1.8 `llvm-c/Orc.h`
//! header this workspace builds against (confirmed by inspecting both,
//! per Task 1's cross-check discipline) -- `LLVMOrcOpaqueThreadSafeContext`
//! is a fully opaque type with only `Create`/`CreateFromLLVMContext`/
//! `Dispose` in its C API surface, no accessor back to the wrapped
//! `LLVMContextRef`. This module instead creates its own plain
//! `LLVMContextRef` first (via `inkwell::context::Context::create()`, so
//! bitcode parsing can reuse `inkwell`'s existing, already-tested
//! `Module::parse_bitcode_from_buffer`), parses the blob into it as a
//! completely ordinary (not-yet-thread-safe) `Context`+`Module` pair while
//! it is still exclusively owned by this function, and only *afterwards*
//! wraps that pair as an ORC `ThreadSafeContext`/`ThreadSafeModule` via
//! `LLVMOrcCreateNewThreadSafeContextFromLLVMContext` (whose documented
//! precondition -- "must not be associated with any existing
//! ThreadSafeContext" -- a freshly-created, never-yet-wrapped `Context`
//! trivially satisfies) and `LLVMOrcCreateNewThreadSafeModule`. This is the
//! same "fresh Context, used, then handed off, never shared" shape the plan
//! specifies, achieved through the API LLVM 22 actually exposes rather than
//! the one the plan's text (written against an older/aspirational API
//! sketch) named.
use std::ffi::{CStr, CString};
use std::ptr;

use inkwell::context::Context;
use inkwell::memory_buffer::MemoryBuffer;
use inkwell::module::Module;
use inkwell::targets::{InitializationConfig, Target};

use llvm_sys::error::{
    LLVMConsumeError, LLVMDisposeErrorMessage, LLVMErrorRef, LLVMGetErrorMessage,
};
use llvm_sys::orc2::lljit::{
    LLVMOrcCreateLLJIT, LLVMOrcDisposeLLJIT, LLVMOrcLLJITAddLLVMIRModule,
    LLVMOrcLLJITAddObjectFile, LLVMOrcLLJITGetMainJITDylib, LLVMOrcLLJITLookup,
    LLVMOrcLLJITMangleAndIntern, LLVMOrcLLJITRef,
};
use llvm_sys::orc2::{
    LLVMJITEvaluatedSymbol, LLVMJITSymbolFlags, LLVMJITSymbolGenericFlags, LLVMOrcAbsoluteSymbols,
    LLVMOrcCSymbolMapPair, LLVMOrcCreateNewThreadSafeContextFromLLVMContext,
    LLVMOrcCreateNewThreadSafeModule, LLVMOrcDisposeMaterializationUnit,
    LLVMOrcDisposeThreadSafeContext, LLVMOrcExecutorAddress, LLVMOrcJITDylibDefine,
    LLVMOrcJITDylibRef, LLVMOrcThreadSafeContextRef,
};
use llvm_sys::prelude::LLVMContextRef;

use crate::error::LoweringError;

/// Owns one ORC v2 `ThreadSafeContext` reference. `Drop` releases this
/// struct's own share of the underlying (shared-ownership,
/// reference-counted) `ThreadSafeContext` data; any `ThreadSafeModule`
/// created from it beforehand keeps a separate share alive independently
/// (`llvm-c/Orc.h`: "Other references (e.g. from ThreadSafeModules) will
/// keep the underlying data alive as long as it is needed").
struct ThreadSafeContext(LLVMOrcThreadSafeContextRef);

impl ThreadSafeContext {
    /// Wraps `ctx` -- a raw `LLVMContextRef` the caller exclusively owns and
    /// has not registered with any existing `ThreadSafeContext` -- as a new
    /// ORC v2 `ThreadSafeContext`, transferring ownership of `ctx` to it.
    fn from_raw_owned(ctx: LLVMContextRef) -> Self {
        // SAFETY: the caller (this module's only call site, in
        // `JitBackend::add_module_from_bitcode`) passes a freshly created
        // `inkwell::context::Context`'s raw pointer, immediately after
        // `mem::forget`-ing that wrapper so nothing else disposes it -- `ctx`
        // is therefore non-null, exclusively owned by this call, and not
        // associated with any existing `ThreadSafeContext`, exactly
        // `LLVMOrcCreateNewThreadSafeContextFromLLVMContext`'s documented
        // precondition. Per that same doc, this call takes ownership of
        // `ctx`; we must not (and do not) call `LLVMContextDispose` on it
        // ourselves after this point.
        let raw = unsafe { LLVMOrcCreateNewThreadSafeContextFromLLVMContext(ctx) };
        ThreadSafeContext(raw)
    }

    fn as_raw(&self) -> LLVMOrcThreadSafeContextRef {
        self.0
    }
}

impl Drop for ThreadSafeContext {
    fn drop(&mut self) {
        // SAFETY: `self.0` was produced by `LLVMOrcCreateNewThreadSafeContextFromLLVMContext`
        // in `from_raw_owned` and has not been disposed before (this struct
        // has no `Clone` impl and is not `Copy`, so there is exactly one
        // `ThreadSafeContext` value, hence exactly one call to
        // `LLVMOrcDisposeThreadSafeContext` for it, here). Per `llvm-c/
        // Orc.h`, this only releases this struct's own share of the
        // (possibly still shared, e.g. with a `ThreadSafeModule`) underlying
        // data -- never a use-after-free of data another reference still
        // needs.
        unsafe { LLVMOrcDisposeThreadSafeContext(self.0) };
    }
}

/// CG-124's `pub struct JitBackend { jit: LLJIT }` -- one shared ORC v2
/// `LLJIT` plus its main `JITDylib`, accepting bitcode-serialized modules
/// from any thread with no external synchronization.
pub struct JitBackend {
    jit: LLVMOrcLLJITRef,
    main_jd: LLVMOrcJITDylibRef,
}

// SAFETY: ORC v2's `LLJIT` is documented by LLVM as safe for concurrent
// `addLLVMIRModule`/lookup calls against one shared instance from multiple
// threads (`docs/design/codegen.md` §4: "ORC v2 is explicitly designed for
// concurrent module addition") -- that is the actual safety argument for
// moving a `JitBackend` to another thread, not merely that the raw pointer
// field happens to compile as `Send`.
unsafe impl Send for JitBackend {}
// SAFETY: see the comment on `impl Send for JitBackend` above -- the same
// ORC v2 concurrency guarantee is exactly what justifies sharing `&
// JitBackend` across threads too; every method below that touches `self.jit`
// (`add_module_from_bitcode`, `get_fn`) takes `&self`, never `&mut self`.
unsafe impl Sync for JitBackend {}

impl JitBackend {
    /// Creates one fresh `LLJIT` instance with its default target-machine
    /// configuration (a null `LLVMOrcLLJITBuilderRef` is documented, `llvm-c/
    /// LLJIT.h`, to mean "use a default-constructed LLJITBuilder").
    pub fn new() -> Result<Self, LoweringError> {
        Target::initialize_native(&InitializationConfig::default()).map_err(|e| LoweringError {
            message: format!("failed to initialize native target: {e}"),
            at: "orc::JitBackend::new".to_string(),
        })?;

        let mut jit: LLVMOrcLLJITRef = ptr::null_mut();
        // `SV_SIM_JIT_OPT` unset gives null (the default builder).
        let builder = jit_builder_from_env()?;
        // SAFETY: `&mut jit` is a valid, uniquely-owned out-pointer for the
        // duration of this call; `builder` is either null (documented in
        // `llvm-c/LLJIT.h` to use a default-constructed `LLJITBuilder`) or a
        // fresh builder from `jit_builder_from_env` that nothing else holds;
        // `LLVMOrcCreateLLJIT` takes ownership of a non-null `Builder`
        // regardless of outcome.
        let err = unsafe { LLVMOrcCreateLLJIT(&mut jit, builder) };
        check_error(err, "LLVMOrcCreateLLJIT")?;
        if jit.is_null() {
            return Err(LoweringError {
                message: "LLVMOrcCreateLLJIT reported success but returned a null LLJIT"
                    .to_string(),
                at: "orc::JitBackend::new".to_string(),
            });
        }

        // SAFETY: `jit` is non-null (checked above), freshly created by
        // `LLVMOrcCreateLLJIT`, and not yet disposed -- `LLVMOrcLLJITGetMainJITDylib`
        // is a plain accessor with no further preconditions.
        let main_jd = unsafe { LLVMOrcLLJITGetMainJITDylib(jit) };

        Ok(JitBackend { jit, main_jd })
    }

    /// CG-124/CG-062/CG-064: parses `bitcode` into a freshly created,
    /// exclusively-owned `Context` (matching `session.rs`'s "the collector
    /// parses each blob into its own fresh context" contract), wraps it as
    /// an ORC v2 `ThreadSafeContext`/`ThreadSafeModule`, and adds it to this
    /// `JitBackend`'s shared main `JITDylib`. Safe to call concurrently from
    /// multiple threads against one shared `&JitBackend` with no external
    /// mutex (CG-062).
    pub fn add_module_from_bitcode(&self, bitcode: &[u8]) -> Result<(), LoweringError> {
        self.add_module_from_bitcode_inner(bitcode)
    }

    /// [`Self::add_module_from_bitcode`] for a big design: splits the module
    /// into one slice per compile thread ([`crate::jit_split`]), code-generates
    /// the slices to relocatable objects concurrently (the `LLJIT` C API has
    /// no compile-thread pool, and its one shared target machine is not
    /// thread-safe, so each thread owns its machine), and links the objects
    /// into the main `JITDylib`, which resolves their cross-references. The
    /// symbols and each function's code match the unsplit module's. Small
    /// designs, `SV_SIM_JIT_THREADS=1` and modules the splitter declines take
    /// the plain path.
    pub fn add_module_split(&self, bitcode: &[u8]) -> Result<(), LoweringError> {
        use std::collections::HashSet;
        let threads = compile_threads()?;
        if threads < 2 || bitcode.len() < SPLIT_MIN_BITCODE_BYTES {
            return self.add_module_from_bitcode(bitcode);
        }
        let at = "orc::JitBackend::add_module_split";
        let err = |message: String| LoweringError {
            message,
            at: at.to_string(),
        };
        fn parse<'c>(bitcode: &[u8], ctx: &'c Context) -> Result<Module<'c>, LoweringError> {
            let buf = MemoryBuffer::create_from_memory_range_copy(bitcode, "jit_split");
            Module::parse_bitcode_from_buffer(&buf, ctx).map_err(|e| LoweringError {
                message: format!("failed to parse bitcode: {e}"),
                at: "orc::JitBackend::add_module_split".to_string(),
            })
        }
        // Every worker parses the whole module before stripping it to its
        // slice, and a UVM-sized module is gigabytes per parse (nine parses at
        // once exceeded a 30 GB cap on uvm_monitor_env). The plan is built
        // first (one parse), then only [`parse_permits`] workers hold an
        // unstripped copy at once; code generation after the strip runs on all
        // of them.
        let t_plan = crate::timing::start();
        let planned = {
            let ctx = discarding_names_context();
            parse(bitcode, &ctx)
                .ok()
                .and_then(|module| crate::jit_split::plan(&module, threads))
        };
        release_freed_memory();
        crate::timing::end(
            t_plan,
            "split plan (parse + call graph)",
            &format!(
                "{threads} threads, {:?} bins",
                planned.as_ref().map(|p| p.bins.len())
            ),
        );
        let Some(plan) = planned else {
            return self.add_module_from_bitcode(bitcode);
        };
        let gate = (
            std::sync::Mutex::new(parse_permits(bitcode.len(), threads)),
            std::sync::Condvar::new(),
        );
        let plan = &plan;
        let gate = &gate;
        let objects: Vec<Result<Option<Vec<u8>>, LoweringError>> = std::thread::scope(|s| {
            let handles: Vec<_> = (0..plan.bins.len())
                .map(|bin| {
                    std::thread::Builder::new()
                        .stack_size(JIT_THREAD_STACK_BYTES)
                        .spawn_scoped(s, move || {
                            let ctx = discarding_names_context();
                            let owned: HashSet<String> = plan.bins[bin].iter().cloned().collect();
                            let module = {
                                let (permits, cv) = gate;
                                let mut n = permits.lock().unwrap_or_else(|e| e.into_inner());
                                while *n == 0 {
                                    n = cv.wait(n).unwrap_or_else(|e| e.into_inner());
                                }
                                *n -= 1;
                                drop(n);
                                let t = crate::timing::start();
                                let module = parse(bitcode, &ctx);
                                crate::timing::end(t, &format!("bin {bin}: parse bitcode"), "");
                                if let Ok(module) = &module {
                                    let t = crate::timing::start();
                                    crate::jit_split::strip_to_bin(
                                        module,
                                        bin,
                                        &owned,
                                        &plan.shared_globals,
                                    );
                                    crate::timing::end(
                                        t,
                                        &format!("bin {bin}: strip"),
                                        &format!("{} owned functions", owned.len()),
                                    );
                                }
                                release_freed_memory();
                                *permits.lock().unwrap_or_else(|e| e.into_inner()) += 1;
                                cv.notify_one();
                                module?
                            };
                            let t = crate::timing::start();
                            let obj = module_to_object(&module);
                            crate::timing::end(
                                t,
                                &format!("bin {bin}: codegen"),
                                &format!("{} bytes", obj.as_ref().map_or(0, Vec::len)),
                            );
                            drop(module);
                            release_freed_memory();
                            obj.map(Some)
                        })
                        .map_err(|e| err(format!("failed to spawn a JIT compile thread: {e}")))
                })
                .collect();
            handles
                .into_iter()
                .map(|h| {
                    h?.join()
                        .unwrap_or_else(|_| Err(err("a JIT compile thread panicked".to_string())))
                })
                .collect()
        });
        // In bin order, so the link order (hence the layout) is deterministic.
        for object in objects {
            if let Some(object) = object? {
                self.add_object(&object)?;
            }
        }
        Ok(())
    }

    /// Adds a relocatable object file to the main `JITDylib`.
    fn add_object(&self, object: &[u8]) -> Result<(), LoweringError> {
        let name = c"jit_split_object";
        // SAFETY: `object`/`name` are valid for the call; the function copies
        // the range into a new buffer owned by the returned reference.
        let buf = unsafe {
            llvm_sys::core::LLVMCreateMemoryBufferWithMemoryRangeCopy(
                object.as_ptr().cast(),
                object.len(),
                name.as_ptr(),
            )
        };
        // SAFETY: `self.jit`/`self.main_jd` are valid for the lifetime of
        // `self`; `buf` is a fresh buffer whose ownership this call takes
        // unconditionally (`llvm-c/LLJIT.h`), so it is not used again.
        let err = unsafe { LLVMOrcLLJITAddObjectFile(self.jit, self.main_jd, buf) };
        check_error(err, "LLVMOrcLLJITAddObjectFile")
    }

    fn add_module_from_bitcode_inner(&self, bitcode: &[u8]) -> Result<(), LoweringError> {
        let ctx = discarding_names_context();
        let buf =
            MemoryBuffer::create_from_memory_range_copy(bitcode, "jit_add_module_from_bitcode");
        let module = Module::parse_bitcode_from_buffer(&buf, &ctx).map_err(|e| LoweringError {
            message: format!("failed to parse bitcode into a fresh context: {e}"),
            at: "orc::JitBackend::add_module_from_bitcode".to_string(),
        })?;

        let raw_module = module.as_mut_ptr();
        let raw_ctx = ctx.raw();
        // SAFETY: both raw pointers have already been read out above; `ctx`
        // and `module`'s ownership is about to transfer to ORC
        // (`ThreadSafeContext::from_raw_owned` takes ownership of `raw_ctx`
        // below, and `LLVMOrcCreateNewThreadSafeModule` takes ownership of
        // `raw_module`). Forgetting these inkwell wrappers -- rather than
        // letting their normal `Drop` impls run, which would call
        // `LLVMContextDispose`/`LLVMDisposeModule` -- is exactly what avoids
        // a double-free of the same underlying LLVM objects once ORC frees
        // them (via `JitBackend`'s own `Drop`, transitively).
        std::mem::forget(module);
        std::mem::forget(ctx);

        let tsctx = ThreadSafeContext::from_raw_owned(raw_ctx);
        // SAFETY: `raw_module` was parsed into `raw_ctx` above (now owned,
        // via `tsctx`, by ORC) and this call site holds sole ownership of it
        // (its inkwell wrapper was `mem::forget`-ten above without running
        // `LLVMDisposeModule`). `LLVMOrcCreateNewThreadSafeModule` takes
        // ownership of `raw_module` and takes its own new shared reference
        // to `tsctx`'s underlying data (so dropping our local `tsctx` value
        // below, once we are done with it, only releases our reference, not
        // this new one).
        let tsm = unsafe { LLVMOrcCreateNewThreadSafeModule(raw_module, tsctx.as_raw()) };
        if tsm.is_null() {
            return Err(LoweringError {
                message: "LLVMOrcCreateNewThreadSafeModule returned null".to_string(),
                at: "orc::JitBackend::add_module_from_bitcode".to_string(),
            });
        }

        // SAFETY: `self.jit`/`self.main_jd` are valid for the lifetime of
        // `self` (created in `new`, disposed only in `JitBackend`'s `Drop`);
        // `tsm` is a uniquely-owned, not-yet-disposed `ThreadSafeModule` we
        // just created. `LLVMOrcLLJITAddLLVMIRModule` "transfers ownership
        // of the TSM argument to the LLJIT instance" unconditionally (`llvm-
        // c/LLJIT.h`: "should not be disposed of or referenced once this
        // function returns", with no success/error qualifier) -- so `tsm`
        // must not, and is not, referenced or disposed again on either path
        // below.
        let err = unsafe { LLVMOrcLLJITAddLLVMIRModule(self.jit, self.main_jd, tsm) };
        check_error(err, "LLVMOrcLLJITAddLLVMIRModule")
    }

    /// Review fix (r1): defines `mangled` as an absolute symbol in this
    /// `JitBackend`'s main `JITDylib`, resolving to the plain, real function
    /// address `addr` (`some_fn as *const () as usize` -- the exact shape
    /// `pipeline.rs::RealBindings` already carries). This is `JitBackend`'s
    /// equivalent of `jit::Jit::add_global_mapping` (BINDING CORRECTION 4's
    /// blocking gap: the legacy `ExecutionEngine` path had no ORC
    /// counterpart) -- it is what lets a bitcode-compiled kernel's
    /// `extern "C"` accessor/`nba_enqueue`/`sv_rt_*` declarations resolve to
    /// real Rust functions instead of dangling as undefined externals.
    ///
    /// Uses `LLVMOrcLLJITMangleAndIntern` (not the bare
    /// `LLVMOrcExecutionSessionIntern`) so the interned name is mangled by
    /// the exact same target-data-layout rules the object-linking layer
    /// applies to every symbol reference inside a module added via
    /// `add_module_from_bitcode`/`LLVMOrcLLJITAddLLVMIRModule` -- an
    /// unmangled intern here would silently fail to satisfy a compiled
    /// module's external reference on any target whose mangling adds a
    /// prefix (e.g. a leading `_`).
    pub fn define_absolute_symbol(&self, mangled: &str, addr: usize) -> Result<(), LoweringError> {
        let name = CString::new(mangled).map_err(|e| LoweringError {
            message: format!("symbol name {mangled:?} contains an interior NUL: {e}"),
            at: "orc::JitBackend::define_absolute_symbol".to_string(),
        })?;
        // SAFETY: `self.jit` is valid for the lifetime of `self`; `name` is
        // a valid, NUL-terminated C string owned by this stack frame for
        // the duration of the call. `LLVMOrcLLJITMangleAndIntern` returns a
        // new, already-retained `LLVMOrcSymbolStringPoolEntryRef` the
        // caller owns (`llvm-c/Orc.h`'s `LLVMOrcExecutionSessionIntern`
        // doc, which this shares the ownership contract with) -- ownership
        // of that reference is transferred to `LLVMOrcAbsoluteSymbols`
        // below, so this function never calls
        // `LLVMOrcReleaseSymbolStringPoolEntry` on it itself.
        let interned = unsafe { LLVMOrcLLJITMangleAndIntern(self.jit, name.as_ptr()) };
        if interned.is_null() {
            return Err(LoweringError {
                message: format!("LLVMOrcLLJITMangleAndIntern returned null for {mangled:?}"),
                at: "orc::JitBackend::define_absolute_symbol".to_string(),
            });
        }

        let flags = LLVMJITSymbolFlags {
            GenericFlags: LLVMJITSymbolGenericFlags::LLVMJITSymbolGenericFlagsExported as u8,
            TargetFlags: 0,
        };
        let mut pair = LLVMOrcCSymbolMapPair {
            Name: interned,
            Sym: LLVMJITEvaluatedSymbol {
                Address: addr as LLVMOrcExecutorAddress,
                Flags: flags,
            },
        };
        // SAFETY: `&mut pair` points to one valid, uniquely-owned
        // `LLVMOrcCSymbolMapPair` for the duration of this call, with a
        // non-null, just-interned `Name`. Per `llvm-c/Orc.h`:
        // "`LLVMOrcAbsoluteSymbols` takes ownership of the elements of the
        // Syms array. The Name fields...are taken to have been retained for
        // this function" -- exactly the retained reference
        // `LLVMOrcLLJITMangleAndIntern` handed back above, so this call
        // consumes it; this function must not (and does not) release
        // `interned` itself. "The client is still responsible for deleting
        // the Sym array itself" -- `pair` is a plain stack value, so no
        // separate free is needed for the array itself.
        let mu = unsafe { LLVMOrcAbsoluteSymbols(&mut pair, 1) };
        if mu.is_null() {
            return Err(LoweringError {
                message: format!("LLVMOrcAbsoluteSymbols returned null for {mangled:?}"),
                at: "orc::JitBackend::define_absolute_symbol".to_string(),
            });
        }

        // SAFETY: `self.main_jd` is valid for the lifetime of `self`; `mu`
        // is a non-null, uniquely-owned `MaterializationUnit` just created
        // above and not yet disposed. Per `llvm-c/Orc.h`: "If this operation
        // succeeds then JITDylib JD will take ownership of MU. If the
        // operation fails then ownership remains with the caller who should
        // call LLVMOrcDisposeMaterializationUnit" -- so on the success path
        // below `mu` must not be disposed again (and is not); the error
        // path disposes it exactly once, matching that documented contract.
        let err = unsafe { LLVMOrcJITDylibDefine(self.main_jd, mu) };
        if !err.is_null() {
            // SAFETY: `mu` was NOT taken over by `LLVMOrcJITDylibDefine`
            // on this failing path (per the doc quoted above), so it is
            // still this function's sole responsibility to dispose of; `mu`
            // is non-null and has not been disposed before.
            unsafe { LLVMOrcDisposeMaterializationUnit(mu) };
            return check_error(err, "LLVMOrcJITDylibDefine");
        }
        Ok(())
    }

    /// Looks up `mangled` in this `JitBackend`'s JIT'd symbols, returning
    /// `Some(fn_ptr)` cast to the caller-specified function-pointer type `F`
    /// when found, `None` for an unknown name. The compiled function
    /// pointer has no remaining dependency on whichever `Context` originally
    /// parsed its defining module (CG-064) -- that `Context` was already
    /// fully consumed/forgotten inside `add_module_from_bitcode`, long
    /// before this call.
    ///
    /// Review fix (r1): this is now `unsafe` -- the caller asserts, via `F`,
    /// that `mangled` names a function with EXACTLY that Rust-level ABI
    /// signature (argument count, argument types, and return type, not
    /// merely `size_of::<F>() == size_of::<usize>()`). `transmute_copy` only
    /// panics on a *size* mismatch; a same-size-but-wrong-shape `F` (e.g.
    /// requesting `extern "C" fn(i32, i32) -> i32` for a symbol that is
    /// actually `extern "C" fn() -> i32`) passes that check silently and
    /// then invokes real undefined behavior when the returned value is
    /// called. This mirrors `inkwell::execution_engine::ExecutionEngine::
    /// get_function`'s own `pub unsafe fn` safety contract (the library's
    /// existing precedent for a generic typed-lookup API over a JIT).
    ///
    /// # Safety
    ///
    /// The caller must ensure `F` is exactly the ABI signature of the
    /// function named `mangled` actually has once JIT-compiled (calling
    /// convention, argument types/count, and return type must all match).
    pub unsafe fn get_fn<F: Copy>(&self, mangled: &str) -> Option<F> {
        let name = CString::new(mangled).ok()?;
        let mut addr: LLVMOrcExecutorAddress = 0;
        let t_lookup = crate::timing::start();
        // SAFETY: `self.jit` is valid for the lifetime of `self`; `name` is
        // a valid, NUL-terminated C string owned by this stack frame for the
        // duration of the call; `&mut addr` is a valid, uniquely-owned
        // out-pointer.
        let err = unsafe { LLVMOrcLLJITLookup(self.jit, &mut addr, name.as_ptr()) };
        crate::timing::end_if_over(t_lookup, 0.5, "JIT lookup (lazy codegen)", mangled);
        if !err.is_null() {
            // SAFETY: `err` is non-null here, so `LLVMConsumeError` is the
            // documented, well-formed way to free it. An unresolved symbol
            // is an ordinary `None` from this method's perspective (matching
            // the existing `Option`-returning shape of `jit.rs::get_raw`),
            // not a reportable error.
            unsafe { LLVMConsumeError(err) };
            return None;
        }
        if addr == 0 {
            return None;
        }
        // SAFETY: the caller asserts, via the type parameter `F`, that
        // `mangled` names a function with exactly that Rust-level ABI
        // signature -- this mirrors the existing transmute pattern in `jit.rs
        // ::get_fn2_i32`/`get_process_fn`. `transmute_copy` (rather than
        // `mem::transmute`) is used because `F`'s size is not known to equal
        // `usize` at this generic definition site -- only at each
        // monomorphized call site, where `transmute_copy`'s runtime size
        // check (debug builds) / caller obligation (release) applies.
        Some(unsafe { std::mem::transmute_copy::<usize, F>(&(addr as usize)) })
    }
}

impl Drop for JitBackend {
    fn drop(&mut self) {
        // SAFETY: `self.jit` was created by `LLVMOrcCreateLLJIT` in `new`
        // and is disposed exactly once, here: `JitBackend` has no `Clone`
        // impl, so there is exactly one owner and exactly one `Drop` call.
        // Disposing it frees all memory it manages, including every module
        // added via `add_module_from_bitcode` and the `main_jd`/`tsctx`
        // shares those modules held -- `self.main_jd` itself needs no
        // separate disposal (owned by the `LLJIT`, per `llvm-c/Orc.h`).
        let err = unsafe { LLVMOrcDisposeLLJIT(self.jit) };
        if !err.is_null() {
            // SAFETY: `err` is non-null here, so `LLVMConsumeError` is a
            // well-formed way to free it. `Drop` is not fallible, so a
            // disposal-time error is discarded rather than propagated or
            // panicking inside a destructor.
            unsafe { LLVMConsumeError(err) };
        }
    }
}

/// Converts a non-null `LLVMErrorRef` into a `LoweringError` (consuming and
/// freeing `err` in the process, per `LLVMGetErrorMessage`'s documented
/// contract); a null `err` (LLVM's success sentinel) maps to `Ok(())`.
fn check_error(err: LLVMErrorRef, context: &str) -> Result<(), LoweringError> {
    if err.is_null() {
        return Ok(());
    }
    // SAFETY: `err` is non-null (checked above). `LLVMGetErrorMessage` is
    // LLVM's documented way to both read and consume/free an `LLVMErrorRef`
    // in one call, so `err` itself needs no separate dispose call on either
    // branch below.
    let msg_ptr = unsafe { LLVMGetErrorMessage(err) };
    if msg_ptr.is_null() {
        return Err(LoweringError {
            message: format!("{context} failed with no error message"),
            at: "orc::JitBackend".to_string(),
        });
    }
    // SAFETY: `msg_ptr` is non-null (checked above), a valid NUL-terminated
    // C string returned by `LLVMGetErrorMessage`, and remains valid until
    // `LLVMDisposeErrorMessage` is called on it immediately below.
    let message = unsafe { CStr::from_ptr(msg_ptr) }
        .to_string_lossy()
        .into_owned();
    // SAFETY: `msg_ptr` was returned by `LLVMGetErrorMessage` above and has
    // not been freed yet; this is its documented matching deallocator.
    unsafe { LLVMDisposeErrorMessage(msg_ptr) };
    Err(LoweringError {
        message: format!("{context}: {message}"),
        at: "orc::JitBackend".to_string(),
    })
}

/// Environment knob for the number of threads that code-generate a big design
/// in parallel ([`JitBackend::add_module_split`]). Unset means the host's
/// available parallelism capped at [`MAX_COMPILE_THREADS`]; `1` disables the
/// split. The machine code and the simulation results are identical at every
/// value, so the knob only trades start-up time against CPU use.
pub const JIT_THREADS_ENV: &str = "SV_SIM_JIT_THREADS";

/// Cap on the default compile-thread count: past this the per-function size
/// skew of real designs (one or two huge processes) leaves threads idle.
const MAX_COMPILE_THREADS: usize = 8;

/// Designs whose bitcode is smaller are compiled whole: the split's extra
/// parses and thread start-up cost more than it saves.
const SPLIT_MIN_BITCODE_BYTES: usize = 200_000;

/// Hands the allocator's freed-but-retained pages back to the OS.
///
/// glibc keeps one arena per thread, and a module the lowering thread
/// dropped (or a parse a worker finished with) stays resident in that
/// arena, invisible to the other threads' allocations: a UVM-sized design
/// holds ~15 GB that way while every worker then allocates its own
/// parse on top. Trimming after each large free keeps the peak at the live
/// set. No-op on non-glibc targets.
pub fn release_freed_memory() {
    #[cfg(all(target_os = "linux", target_env = "gnu"))]
    {
        // SAFETY: the declaration matches glibc's `int malloc_trim(size_t)`.
        unsafe extern "C" {
            fn malloc_trim(pad: usize) -> i32;
        }
        // SAFETY: `malloc_trim` takes no pointers and is thread-safe.
        unsafe {
            malloc_trim(0);
        }
    }
}

/// Environment knob for the number of full bitcode parses of a huge module
/// that may be live at once during the split ([`JitBackend::add_module_split`]).
/// Unset sizes it from the memory budget ([`parse_permits`]).
pub const JIT_PARSES_ENV: &str = "SV_SIM_JIT_PARSES";

/// Parsed-module memory per bitcode byte, a conservative estimate (a parse
/// of the 1.06 GB UVM bitcode holds ~8 GB until the worker strips it).
const PARSE_BYTES_PER_BITCODE_BYTE: u64 = 10;

/// The memory this process may still use: the smaller of the host's
/// `MemAvailable` and the cgroup limit headroom, `None` when neither reads.
fn memory_budget() -> Option<u64> {
    let kb = |text: &str, key: &str| -> Option<u64> {
        text.lines()
            .find_map(|l| l.strip_prefix(key))
            .and_then(|r| r.split_whitespace().next())
            .and_then(|n| n.parse::<u64>().ok())
    };
    let avail = std::fs::read_to_string("/proc/meminfo")
        .ok()
        .and_then(|t| kb(&t, "MemAvailable:"))
        .map(|k| k * 1024);
    let cg = std::fs::read_to_string("/proc/self/cgroup")
        .ok()
        .and_then(|t| {
            t.lines()
                .find_map(|l| l.strip_prefix("0::").map(str::to_owned))
        })
        .and_then(|path| {
            let dir = format!("/sys/fs/cgroup{}", path.trim());
            let max = std::fs::read_to_string(format!("{dir}/memory.max")).ok()?;
            let max = max.trim().parse::<u64>().ok()?;
            let cur = std::fs::read_to_string(format!("{dir}/memory.current"))
                .ok()?
                .trim()
                .parse::<u64>()
                .ok()?;
            Some(max.saturating_sub(cur))
        });
    match (avail, cg) {
        (Some(a), Some(c)) => Some(a.min(c)),
        (a, c) => a.or(c),
    }
}

/// How many split compile threads may parse (and strip) the full module at
/// once: as many as half of the [`memory_budget`] holds at
/// [`PARSE_BYTES_PER_BITCODE_BYTE`] each, clamped to `1..=threads`.
fn parse_permits(bitcode_len: usize, threads: usize) -> usize {
    if let Some(n) = std::env::var(JIT_PARSES_ENV)
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .filter(|&n| n >= 1)
    {
        return n.min(threads);
    }
    let Some(budget) = memory_budget() else {
        return 1;
    };
    permits_for_budget(budget, bitcode_len, threads)
}

/// [`parse_permits`] for a known `budget`: only half of it is lent to the
/// parses, because the bins already past their strip keep their codegen state
/// resident (several GB each on the 1 GB UVM bitcode) while the later bins
/// parse. Three quarters admitted 2 parses at a 30 GB cap and the sweep's
/// `uvm_agent_active` was cgroup-killed on the third pair of bins.
fn permits_for_budget(budget: u64, bitcode_len: usize, threads: usize) -> usize {
    let per = (bitcode_len as u64)
        .saturating_mul(PARSE_BYTES_PER_BITCODE_BYTE)
        .max(1);
    ((budget / 2 / per) as usize).clamp(1, threads)
}

/// Stack for a compile thread; LLVM's instruction selection recurses deeply on
/// the very large straight-line functions the lowering can produce.
const JIT_THREAD_STACK_BYTES: usize = 64 << 20;

/// [`JIT_THREADS_ENV`], parsed; an unrecognised value fails loudly.
fn compile_threads() -> Result<usize, LoweringError> {
    match std::env::var_os(JIT_THREADS_ENV) {
        None => Ok(std::thread::available_parallelism()
            .map_or(1, std::num::NonZero::get)
            .min(MAX_COMPILE_THREADS)),
        Some(v) => match v.to_str().and_then(|s| s.parse::<usize>().ok()) {
            Some(n) if n >= 1 => Ok(n),
            _ => Err(LoweringError {
                message: format!("{JIT_THREADS_ENV}={v:?}: expected a thread count of at least 1"),
                at: "orc::compile_threads".to_string(),
            }),
        },
    }
}

/// Environment knob naming the JIT backend's code-generation optimization
/// level (`0` = FastISel/no regalloc tuning, `1`, `2` = LLVM's default,
/// `3`). Unset means the `LLJIT` default (level 2). The lowered IR and the
/// simulation results are identical at every level (only the machine code
/// differs), so the knob trades JIT time against kernel speed: on the bench
/// suite level 0 cuts JIT start-up by 26% (picorv32) to 40% (lanes) and costs
/// 18% on a kernel-heavy run (fourstate), so it is opt-in.
pub const JIT_OPT_ENV: &str = "SV_SIM_JIT_OPT";

/// [`JIT_OPT_ENV`] as an optimization level, `None` when unset.
fn opt_level_from_env() -> Result<Option<inkwell::OptimizationLevel>, LoweringError> {
    use inkwell::OptimizationLevel;
    let Some(v) = std::env::var_os(JIT_OPT_ENV) else {
        return Ok(None);
    };
    match v.to_str() {
        Some("0") => Ok(Some(OptimizationLevel::None)),
        Some("1") => Ok(Some(OptimizationLevel::Less)),
        Some("2") => Ok(Some(OptimizationLevel::Default)),
        Some("3") => Ok(Some(OptimizationLevel::Aggressive)),
        _ => Err(LoweringError {
            message: format!(
                "{JIT_OPT_ENV}={v:?}: expected 0, 1, 2 or 3 (JIT code-generation optimization level)"
            ),
            at: "orc::jit_builder_from_env".to_string(),
        }),
    }
}

/// Code-generates `module` to a relocatable object on a target machine of its
/// own (LLVM target machines are not shareable across threads), configured as
/// the `LLJIT` default machine is for ELF/JITLink: host CPU and features, PIC,
/// small code model, `SV_SIM_JIT_OPT` (default level 2).
fn module_to_object(module: &Module<'_>) -> Result<Vec<u8>, LoweringError> {
    use inkwell::OptimizationLevel;
    use inkwell::targets::{CodeModel, FileType, RelocMode, TargetMachine};
    let err = |message: String| LoweringError {
        message,
        at: "orc::module_to_object".to_string(),
    };
    let level = opt_level_from_env()?.unwrap_or(OptimizationLevel::Default);
    let triple = TargetMachine::get_default_triple();
    let target = Target::from_triple(&triple)
        .map_err(|e| err(format!("host triple {:?}: {e}", triple.as_str())))?;
    let cpu = TargetMachine::get_host_cpu_name();
    let features = TargetMachine::get_host_cpu_features();
    let tm = target
        .create_target_machine(
            &triple,
            cpu.to_str().unwrap_or(""),
            features.to_str().unwrap_or(""),
            level,
            RelocMode::PIC,
            CodeModel::Small,
        )
        .ok_or_else(|| err(format!("no target machine for {:?}", triple.as_str())))?;
    module.set_triple(&triple);
    module.set_data_layout(&tm.get_target_data().get_data_layout());
    let buf = tm
        .write_to_memory_buffer(module, FileType::Object)
        .map_err(|e| err(format!("object emission failed: {e}")))?;
    Ok(buf.as_slice().to_vec())
}

/// Builds the `LLJITBuilder` for [`JIT_OPT_ENV`], or null (the default
/// builder) when the variable is unset. An unrecognised value fails loudly.
fn jit_builder_from_env() -> Result<llvm_sys::orc2::lljit::LLVMOrcLLJITBuilderRef, LoweringError> {
    use inkwell::targets::{CodeModel, RelocMode, TargetMachine};
    let err = |message: String| LoweringError {
        message,
        at: "orc::jit_builder_from_env".to_string(),
    };
    let Some(level) = opt_level_from_env()? else {
        return Ok(ptr::null_mut());
    };
    let triple = TargetMachine::get_default_triple();
    let target = Target::from_triple(&triple)
        .map_err(|e| err(format!("host triple {:?}: {e}", triple.as_str())))?;
    let cpu = TargetMachine::get_host_cpu_name();
    let features = TargetMachine::get_host_cpu_features();
    // Same host CPU/feature pair the default `LLJIT` target machine uses, so
    // only the optimization level differs from the unset case.
    let tm = target
        .create_target_machine(
            &triple,
            cpu.to_str().unwrap_or(""),
            features.to_str().unwrap_or(""),
            level,
            RelocMode::Default,
            CodeModel::JITDefault,
        )
        .ok_or_else(|| {
            err(format!(
                "no target machine for host triple {:?}",
                triple.as_str()
            ))
        })?;
    // SAFETY: `tm` is a live target machine; `CreateFromTargetMachine` takes
    // ownership of the raw handle, so `tm` is forgotten to avoid a double
    // dispose. The JTMB is then owned by the builder (`SetJITTargetMachine
    // Builder`), which `LLVMOrcCreateLLJIT` consumes in `JitBackend::new`.
    unsafe {
        let jtmb =
            llvm_sys::orc2::LLVMOrcJITTargetMachineBuilderCreateFromTargetMachine(tm.as_mut_ptr());
        std::mem::forget(tm);
        let b = llvm_sys::orc2::lljit::LLVMOrcCreateLLJITBuilder();
        llvm_sys::orc2::lljit::LLVMOrcLLJITBuilderSetJITTargetMachineBuilder(b, jtmb);
        Ok(b)
    }
}

/// A fresh context for machine-code generation from bitcode. Lowering names
/// local values (the wide-temporary pool keys on alloca names), but once the
/// module is bitcode nothing reads those names again: dropping them at parse
/// time saves the symbol-table memory and work on every instruction (the
/// UVM-sized modules hold tens of millions). Global (function) names are
/// never discarded, so symbol resolution is unchanged.
fn discarding_names_context() -> Context {
    let ctx = Context::create();
    // SAFETY: `ctx.raw()` is the live context just created.
    unsafe { llvm_sys::core::LLVMContextSetDiscardValueNames(ctx.raw(), 1) };
    ctx
}

#[cfg(test)]
mod permit_tests {
    use super::*;

    /// A bitcode too big for any budget still gets one permit; a tiny one
    /// never gets more than one per thread.
    #[test]
    fn parse_permits_clamps() {
        assert_eq!(parse_permits(usize::MAX / 64, 8), 1);
        assert!(parse_permits(1, 8) <= 8);
        assert!(parse_permits(1, 8) >= 1);
    }

    /// The 1.06 GB UVM bitcode under a 30 GB cap gets a single parse (two
    /// parses plus the resident codegen of finished bins exceeded the cap);
    /// a roomy host still parses in parallel.
    #[test]
    fn parse_permits_uvm_under_a_30g_cap() {
        const GB: u64 = 1 << 30;
        let uvm = 1_058_271_860;
        assert_eq!(permits_for_budget(27 * GB, uvm, 8), 1);
        assert_eq!(permits_for_budget(57 * GB, uvm, 8), 2);
        assert_eq!(permits_for_budget(139 * GB, uvm, 8), 7);
    }
}