Session D: Interpreter ↔ codegen JIT dispatch wiring

End-to-end: parse OMC source → compile to bytecode → JIT eligible
fns in dual-band mode → register dispatch hook on Interpreter →
JIT-compiled native code runs in place of tree-walk for those fns.

The mechanism is proven via integration tests that drive the full
pipeline. The CLI-level `OMC_HBIT_JIT=1` env var is one small
follow-up away (extracting omnimcode-cli to a new package — see
deferred Session D.5 task; the workspace-cycle constraint prevents
adding codegen as a dep of core's existing binary).

Codegen API additions:
- `JitContext::jit_module(&Module) -> HashMap<String, JittedFn>`:
  walks user fns, attempts dual-band lowering on each, returns
  the successes as arity-tagged raw fn pointers. Failed fns are
  silently skipped — and crucially, their partial LLVM IR is now
  erased from the module so JIT finalization for the rest of the
  module doesn't crash on broken IR (`function.delete()` after
  the lowering errors).
- `JittedFn { arity, fn_ptr: *const () }` with `call(&[i64])`
  dispatcher for arities 0..=4. Uses transmute to cast the raw
  pointer to the right `unsafe extern "C" fn(...) -> i64`
  signature based on arity.
- `extract_raw_fn_ptr` helper: typed `engine.get_function::<F>`
  per arity, then `JitFunction::into_raw()` for storage.
  (Discovered the bug: `engine.get_function_address` returned
  garbage where the typed `get_function::<F>` produces correct
  fn pointers — sticking with the typed path.)

Codegen lowerer fix (applies to BOTH scalar and dual-band):
- `bind_params_into_locals()` now runs at fn entry. The OMC
  bytecode compiler emits `Op::LoadVar("x")` for parameter access
  in fn bodies (treating params as already-bound locals); the
  bytecode VM and tree-walk both pre-populate these bindings
  before executing the body. Without this fix, the JIT was
  reading uninitialized allocas — returning pointer-looking
  garbage instead of the param value (the symptom that broke
  the first integration test before it was diagnosed).

Interpreter additions:
- `JitDispatch = Rc<dyn Fn(&str, &[Value]) -> Option<Result<Value, String>>>`
  type alias for the dispatch closure.
- `Interpreter::jit_dispatch: Option<JitDispatch>` field.
- `Interpreter::set_jit_dispatch(Some(hook))` setter (the
  embedder API the future omnimcode-cli will call).
- `invoke_user_function_at` consults the hook BEFORE running the
  tree-walk body. `Some(_)` return short-circuits tree-walk;
  `None` falls through. Marshalling Value↔i64 happens in the
  hook closure (decoupled from core).

Tests (5 new integration tests, all passing):
- jit_dispatch_routes_simple_int_fn: `double(21) == 42` via JIT
- jit_module_returns_callable_fn_directly: isolation — call JIT'd
  fn without going through Interpreter (proves jit_module's
  fn-ptr extraction)
- jit_dispatch_matches_tree_walk_factorial: factorial(10) == 3.6M,
  matches tree-walk-only run
- jit_dispatch_matches_tree_walk_sum_loop: while loop with two
  locals matches across both code paths
- jit_dispatch_falls_through_on_unsupported_fn: a string-using fn
  silently skipped by JIT (delete() of partial IR), int fn still
  routes through JIT, both produce correct outputs

Workspace state: 19/19 codegen tests green, 149/149 omnimcode-core
Rust unit tests green, 18/18 OMC harmonic-lib tests green.

Deferred to Session D.5: extract omnimcode-cli, wire OMC_HBIT_JIT
to populate set_jit_dispatch from inside main.rs.

Deferred to Session E: AVX-512 path (`<8 x i64>` carrier with
explicit LLVM intrinsics), benchmark harness measuring actual
speedup vs tree-walk and bytecode VM.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: RandomCoder-lab

omnimcode-codegen/src/dual_band.rs

@@ -103,10 +103,41 @@ impl<'ctx, 'a> DualBandLowerer<'ctx, 'a> {
 
         self.collect_leaders()?;
         self.collect_cleanup_pops();
+        self.bind_params_into_locals()?;
         self.emit_body()?;
         Ok(self.function)
     }
 
+    /// Bind each fn parameter into a named local-variable slot. The
+    /// OMC bytecode compiler emits `LoadVar("x")` for parameter access
+    /// in fn bodies; we mirror what the bytecode VM does at fn entry
+    /// and pre-populate each parameter into a `<2 x i64>` alloca slot
+    /// keyed by the parameter name. β = α at entry (matched bands);
+    /// later sessions add explicit phi-shadow ops that diverge β.
+    fn bind_params_into_locals(&mut self) -> Result<(), CodegenError> {
+        for (i, pname) in self.f.params.clone().iter().enumerate() {
+            let param = self
+                .function
+                .get_nth_param(i as u32)
+                .ok_or_else(|| format!("hbit bind_params: no param at slot {}", i))?;
+            let iv = match param {
+                BasicValueEnum::IntValue(iv) => iv,
+                _ => {
+                    return Err(format!(
+                        "hbit bind_params: non-int param at slot {}",
+                        i
+                    ))
+                }
+            };
+            let v = self.splat(iv, &format!("{}_init", pname))?;
+            let slot = self.get_or_create_slot(pname)?;
+            self.builder
+                .build_store(slot, v)
+                .map_err(|e| format!("hbit bind_params store {}: {}", pname, e))?;
+        }
+        Ok(())
+    }
+
     fn collect_leaders(&mut self) -> Result<(), CodegenError> {
         let mut leaders: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
         leaders.insert(0);

omnimcode-codegen/src/lib.rs

@@ -56,6 +56,71 @@ pub struct JitContext<'ctx> {
 /// Error type for codegen failures. Keeps it simple — just a String.
 pub type CodegenError = String;
 
+/// A successfully JIT'd OMC function, presented as an arity-tagged
+/// raw function pointer. Callable via `JittedFn::call(args)` for
+/// the supported arities (0..=4); larger arities should be folded
+/// down via a future uniform-arg-array calling convention.
+///
+/// SAFETY: the underlying machine code is owned by the
+/// `JitContext`/`ExecutionEngine` that produced this struct. Calling
+/// after that JitContext is dropped is undefined behavior. In the
+/// current Session D design, the main CLI keeps the JitContext
+/// alive for the entire program duration (Box::leak), so the
+/// invariant holds for normal use.
+#[derive(Clone, Copy, Debug)]
+pub struct JittedFn {
+    pub arity: usize,
+    /// Erased fn pointer. Cast to the right `unsafe extern "C" fn`
+    /// signature at call time based on `arity`.
+    pub fn_ptr: *const (),
+}
+
+// SAFETY: a raw function pointer is `Send + Sync` — it's plain data.
+// The LLVM-generated machine code is read-only and re-entrant.
+unsafe impl Send for JittedFn {}
+unsafe impl Sync for JittedFn {}
+
+impl JittedFn {
+    /// Call this JITted fn with i64 args. Returns `Some(result)` when
+    /// arity matches a supported overload, `None` otherwise. Caller is
+    /// responsible for keeping the producing JitContext alive — that's
+    /// the unsafe invariant this method enforces minimally (it's
+    /// "safe" because we trust the pointer, but a use-after-free of
+    /// the JitContext would crash here).
+    pub fn call(&self, args: &[i64]) -> Option<i64> {
+        if args.len() != self.arity {
+            return None;
+        }
+        unsafe {
+            match self.arity {
+                0 => {
+                    let f: unsafe extern "C" fn() -> i64 = std::mem::transmute(self.fn_ptr);
+                    Some(f())
+                }
+                1 => {
+                    let f: unsafe extern "C" fn(i64) -> i64 = std::mem::transmute(self.fn_ptr);
+                    Some(f(args[0]))
+                }
+                2 => {
+                    let f: unsafe extern "C" fn(i64, i64) -> i64 = std::mem::transmute(self.fn_ptr);
+                    Some(f(args[0], args[1]))
+                }
+                3 => {
+                    let f: unsafe extern "C" fn(i64, i64, i64) -> i64 =
+                        std::mem::transmute(self.fn_ptr);
+                    Some(f(args[0], args[1], args[2]))
+                }
+                4 => {
+                    let f: unsafe extern "C" fn(i64, i64, i64, i64) -> i64 =
+                        std::mem::transmute(self.fn_ptr);
+                    Some(f(args[0], args[1], args[2], args[3]))
+                }
+                _ => None,
+            }
+        }
+    }
+}
+
 impl<'ctx> JitContext<'ctx> {
     pub fn new(context: &'ctx Context) -> Result<Self, CodegenError> {
         let module = context.create_module("omc_jit");
@@ -107,6 +172,88 @@ impl<'ctx> JitContext<'ctx> {
         lowerer.lower()
     }
 
+    /// Try to JIT every user function in a bytecode `Module` in dual-band
+    /// mode. Functions whose bodies use ops the codegen layer doesn't
+    /// yet support (strings, dicts, builtins, cross-fn calls, etc.)
+    /// are silently skipped — they stay routed through the tree-walk
+    /// interpreter at runtime.
+    ///
+    /// Returns a map of `fn_name -> JittedFn` for every fn that did
+    /// lower successfully. The native function pointers inside
+    /// `JittedFn` are owned by `self` (the underlying ExecutionEngine);
+    /// callers must not invoke the returned fns after `self` is dropped.
+    ///
+    /// The returned name uses the ORIGINAL (un-suffixed) bytecode-side
+    /// fn name; under the hood the LLVM module sees `<name>_hbit` per
+    /// the dual-band lowerer's naming convention.
+    ///
+    /// Session D scope: every user fn is attempted. Sessions later
+    /// add explicit `@hbit` pragma filtering so non-tagged fns aren't
+    /// JIT'd even if they could be.
+    pub fn jit_module(
+        &self,
+        module: &omnimcode_core::bytecode::Module,
+    ) -> Result<HashMap<String, JittedFn>, CodegenError> {
+        let mut out: HashMap<String, JittedFn> = HashMap::new();
+        for (name, cf) in &module.functions {
+            let suffixed = format!("{}_hbit", name);
+            match self.lower_function_dual_band(cf) {
+                Ok(_) => {
+                    // get_function::<F> triggers JIT finalization and
+                    // returns a JitFunction wrapping the raw pointer.
+                    // We dispatch on arity to pick the right F so we
+                    // can extract the raw fn pointer for storage.
+                    let arity = cf.params.len();
+                    let fn_ptr = unsafe { self.extract_raw_fn_ptr(&suffixed, arity)? };
+                    out.insert(
+                        name.clone(),
+                        JittedFn { arity, fn_ptr },
+                    );
+                }
+                Err(_) => {
+                    // Lowering failed mid-emission. The LLVM module
+                    // now contains a partial / broken function — leaving
+                    // it in place would corrupt JIT finalization for
+                    // every subsequent fn (and crash on first call).
+                    // Erase it so the rest of the module stays valid.
+                    if let Some(broken) = self.module.get_function(&suffixed) {
+                        unsafe { broken.delete() };
+                    }
+                }
+            }
+        }
+        Ok(out)
+    }
+
+    /// Erase a typed JitFunction down to a `*const ()` pointer for
+    /// arity-tagged storage in `JittedFn`. Internal helper for
+    /// `jit_module`; the caller is responsible for not invoking the
+    /// returned pointer after `self` is dropped.
+    unsafe fn extract_raw_fn_ptr(
+        &self,
+        name: &str,
+        arity: usize,
+    ) -> Result<*const (), CodegenError> {
+        macro_rules! by_arity {
+            ($t:ty) => {{
+                let jf: JitFunction<'ctx, $t> = self
+                    .engine
+                    .get_function(name)
+                    .map_err(|e| format!("get_function({}): {:?}", name, e))?;
+                jf.into_raw() as *const ()
+            }};
+        }
+        let ptr = match arity {
+            0 => by_arity!(unsafe extern "C" fn() -> i64),
+            1 => by_arity!(unsafe extern "C" fn(i64) -> i64),
+            2 => by_arity!(unsafe extern "C" fn(i64, i64) -> i64),
+            3 => by_arity!(unsafe extern "C" fn(i64, i64, i64) -> i64),
+            4 => by_arity!(unsafe extern "C" fn(i64, i64, i64, i64) -> i64),
+            _ => return Err(format!("arity {} not supported in Session D jit_module", arity)),
+        };
+        Ok(ptr)
+    }
+
     /// JIT-lookup helper for single-arg i64 functions.
     pub unsafe fn get_i64_i64(
         &self,
@@ -187,10 +334,36 @@ impl<'ctx, 'a> FunctionLowerer<'ctx, 'a> {
 
         self.collect_leaders()?;
         self.collect_cleanup_pops();
+        self.bind_params_into_locals()?;
         self.emit_body()?;
         Ok(self.function)
     }
 
+    /// Bind each fn parameter into a named local-variable slot.
+    /// The OMC bytecode compiler emits `Op::LoadVar("x")` for parameter
+    /// access in the body (treating params as locals already in scope).
+    /// The bytecode VM and tree-walk interpreter both pre-populate
+    /// these bindings before executing the body; we mirror that here
+    /// so LoadVar resolves to the actual parameter value rather than
+    /// reading from an uninitialized alloca.
+    fn bind_params_into_locals(&mut self) -> Result<(), CodegenError> {
+        for (i, pname) in self.f.params.clone().iter().enumerate() {
+            let param = self
+                .function
+                .get_nth_param(i as u32)
+                .ok_or_else(|| format!("bind_params: no param at slot {}", i))?;
+            let iv = match param {
+                BasicValueEnum::IntValue(iv) => iv,
+                _ => return Err(format!("bind_params: non-int param at slot {}", i)),
+            };
+            let slot = self.get_or_create_slot(pname)?;
+            self.builder
+                .build_store(slot, iv)
+                .map_err(|e| format!("bind_params store {}: {}", pname, e))?;
+        }
+        Ok(())
+    }
+
     /// First pass: find op-indices that begin a new basic block. An
     /// op-index is a leader if:
     /// - it's 0 (entry)