feat(reg-vm): OSR×J2 inline closures with Float captures (Pending #2 broadening)

A captured f64 is bit-reinterpreted from its i64 capture slot (not integer-
converted), reusing the native f64<->i64 bitcast convention, so a closure
capturing Floats inlines inside an OSR loop. Int/Bool captures unchanged;
non-scalar captures still bail. OSR differential byte-identical on/off
(float parity exact); 10.3x faster on osr_float_closure_loop. Positive +
negative tests.

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

Author: olwang

benchmarks/vm-jit/kernels/osr_float_closure_loop.rss

@@ -0,0 +1,54 @@
+// OSR × J2 FLOAT-CAPTURE kernel (Pending #2 broadening). A higher-order function
+// `apply_loop` takes a closure PARAMETER `g: read Fn(Float) -> Float` and calls it
+// in a hot loop. The closure CAPTURES a FLOAT (`scale`), so the inline must
+// materialize that f64 capture into the loop body by BIT-REINTERPRETING its i64
+// capture slot (the `closure_capture` helper returns `f64::to_bits` as i64; the
+// inlined body's Float-class capture register bit-casts it back to f64 — NOT an
+// integer→float conversion). `apply_loop` is also I/O-tangled (`Log.write` around
+// the loop in the same, once-called function), so the whole function is native-
+// INELIGIBLE — only OSR can run the loop natively, and only by inlining the
+// capturing closure into it (otherwise the per-iteration `CallClosure` keeps the
+// loop off the native subset).
+//
+// Run OSR off (`--mode jit-native`) vs OSR on (`--mode jit-native-osr`, or the
+// default auto-trigger); the OSR median should beat the off median. The closure is
+// monomorphic at the call site (one callee, one scalar Float capture), the
+// precondition for the OSR × J2 inline.
+
+fn bench_size(default: Int) -> Int {
+    let raw = Args.get_or_default(index: 0, default: read String.from_int(value: default))
+    match String.parse_int(value: read raw) {
+        Some(value) => {
+            return value
+        }
+        None => {
+            return default
+        }
+    }
+}
+
+fn apply_loop(g: read Fn(Float) -> Float, limit: Int, seed: Float) -> Float {
+    Log.write(message: read "osr-begin")
+    let mut i = 0
+    let mut total = seed
+    let mut x = 0.0
+    while i < limit {
+        total = total + g(read x)
+        x = x + 1.0
+        i = i + 1
+    }
+    Log.write(message: read String.from_float(value: total))
+    return total
+}
+
+fn main() -> Unit {
+    let limit = bench_size(default: 2000000)
+    let scale = 1.5
+    let g = fn(value) captures(read scale) effects(pure) {
+        return value * scale + scale
+    }
+    let result = apply_loop(g: read g, limit: limit, seed: 0.0)
+    Log.write(message: read String.from_float(value: result))
+    return Unit
+}
+features: local

crates/rsscript/src/reg_vm/mod.rs

@@ -3576,13 +3576,16 @@ fn translate_to_native_jit(
                         && native_set_ty(ty, *dst, NativeTy::Int, c)
                 }
                 // Capturing-closure inline (OSR × J2): the closure operand is a
-                // native handle; the materialized capture `dst` is an `Int`-class
-                // register (the `closure_capture` helper returns scalar bits). A
-                // body that uses the capture as a `Float` conflicts here ⇒ bail, so
-                // only Int/Bool-class captures inline natively (Float deferred).
-                RegInstr::NativeClosureCapture { dst, closure, .. } => {
+                // native handle; the materialized capture `dst` carries the
+                // capture's scalar bits (the `closure_capture` helper returns the
+                // raw i64 bit pattern: an `Int` directly, a `Bool` as 0/1, a
+                // `Float` reinterpreted via `f64::to_bits`). Leave `dst`'s class to
+                // flow from its uses — exactly like a `GetFieldSlot` read — so an
+                // Int/Bool capture stays Int-class and a Float capture becomes
+                // Float-class. Lowering admits only a provably Int/Bool/Float `dst`
+                // (and the Float arm bit-reinterprets the i64 slot to f64).
+                RegInstr::NativeClosureCapture { dst: _, closure, .. } => {
                     native_set_ty(ty, *closure, NativeTy::Handle, c)
-                        && native_set_ty(ty, *dst, NativeTy::Int, c)
                 }
                 _ => true,
             };
@@ -3996,10 +3999,16 @@ fn translate_to_native_jit(
                 index,
             } => {
                 // The closure handle is a native-readable handle; materialize
-                // capture `index`'s scalar bits into the `Int`-class `dst` (the
-                // inlined body's capture register). A non-scalar capture bails
-                // out-of-band in the host helper.
-                require(handle_reg(*closure) && int_or_free(*dst))?;
+                // capture `index`'s scalar bits into `dst` (the inlined body's
+                // capture register). `dst` may be Int/Bool (i64 used directly) or
+                // Float (the i64 slot is `f64::to_bits`, bit-reinterpreted to f64 in
+                // codegen); an unconstrained `dst` defaults to Int. A non-scalar
+                // (Handle/flat) `dst` bails. A non-scalar capture VALUE additionally
+                // bails out-of-band in the host helper at runtime.
+                require(
+                    handle_reg(*closure)
+                        && (int_or_free(*dst) || bool_ty(*dst) || float(*dst)),
+                )?;
                 let index = i64::try_from(*index).ok()?;
                 let index = u32::try_from(index).ok()?;
                 JitInstr::ClosureCapture {
@@ -4355,18 +4364,20 @@ fn translate_osr_loop(
                 }
                 // Synthetic closure-inline ops (only present when an inlined
                 // capturing/monomorphic closure body landed in the OSR region):
-                // the closure operand is a native Handle param; an id read or a
-                // materialized capture is `Int`-class.
+                // the closure operand is a native Handle param; an id read is
+                // `Int`-class. A materialized capture's `dst` is left to flow from
+                // its uses (Int/Bool/Float) — the helper returns the raw scalar bit
+                // pattern (a `Float` via `f64::to_bits`), and lowering admits only a
+                // provably Int/Bool/Float `dst`, bit-reinterpreting a Float slot.
                 RegInstr::NativeGuardClosureId { closure, .. } => {
                     native_set_ty(ty, *closure, NativeTy::Handle, c)
                 }
                 RegInstr::NativeClosureId { dst, closure } => {
                     native_set_ty(ty, *closure, NativeTy::Handle, c)
                         && native_set_ty(ty, *dst, NativeTy::Int, c)
                 }
-                RegInstr::NativeClosureCapture { dst, closure, .. } => {
+                RegInstr::NativeClosureCapture { dst: _, closure, .. } => {
                     native_set_ty(ty, *closure, NativeTy::Handle, c)
-                        && native_set_ty(ty, *dst, NativeTy::Int, c)
                 }
                 _ => true,
             };
@@ -4552,7 +4563,14 @@ fn translate_osr_loop(
                 JitInstr::ClosureId { dst: r(*dst), base: r(*closure) }
             }
             RegInstr::NativeClosureCapture { dst, closure, index } => {
-                require(handle_reg(*closure) && int_or_free(*dst))?;
+                // The capture's `dst` may be Int/Bool (i64 slot used directly) or
+                // Float (the i64 slot is `f64::to_bits`, bit-reinterpreted to f64
+                // in codegen). An unconstrained `dst` defaults to Int. A non-scalar
+                // `dst` (Handle/flat array) cannot hold a scalar capture ⇒ bail.
+                require(
+                    handle_reg(*closure)
+                        && (int_or_free(*dst) || bool_ty(*dst) || float(*dst)),
+                )?;
                 let index = u32::try_from(*index).ok()?;
                 JitInstr::ClosureCapture { dst: r(*dst), base: r(*closure), index }
             }

crates/rsscript/tests/jit_acceptance.rs

@@ -1392,6 +1392,119 @@ fn main() -> Unit {
     );
 }
 
+/// OSR × J2 POSITIVE test (Pending #2 Float-capture broadening) — a CAPTURING
+/// monomorphic param-handle closure whose capture is a FLOAT, called in a hot loop
+/// inside an I/O-tangled (native-INELIGIBLE) function, MUST inline (bit-reinterpreting
+/// the f64 capture from its i64 slot) and OSR to an allocation-free native loop.
+/// `apply_loop` takes `g: read Fn(Float) -> Float` and calls it every iteration;
+/// `g = |value| value * scale + scale` captures the Float `scale = 1.5`. The capture
+/// is materialized via the `closure_capture` helper, which returns `f64::to_bits` as
+/// i64; the inlined body's Float-class capture register bit-reinterprets that i64 to
+/// f64 (NOT an integer→float conversion). Forcing OSR on, stdout (a Float, formatted
+/// via `String.from_float`) MUST be byte-identical to the pure interpreter AND
+/// `osr_entries > 0` proves the loop inlined the capturing closure + materialized the
+/// FLOAT capture + OSR'd. If the bits were int-converted, the Float would be wrong and
+/// stdout would diverge.
+#[cfg(feature = "native-jit")]
+#[test]
+fn native_osr_j2_float_capture_closure_loop_matches_interpreter() {
+    let source = "\
+fn apply_loop(g: read Fn(Float) -> Float, limit: Int, seed: Float) -> Float {
+    Log.write(message: read \"begin\")
+    let mut i = 0
+    let mut total = seed
+    let mut x = 0.0
+    while i < limit {
+        total = total + g(read x)
+        x = x + 1.0
+        i = i + 1
+    }
+    Log.write(message: read String.from_float(value: total))
+    return total
+}
+
+fn main() -> Unit {
+    let scale = 1.5
+    let g = fn(value) captures(read scale) effects(pure) {
+        return value * scale + scale
+    }
+    Log.write(message: read String.from_float(value: apply_loop(g: read g, limit: read 4000, seed: read 0.0)))
+    return Unit
+}
+";
+    let file = "jit-osr-j2-float-capture-closure.rss";
+    let interp = common::run_vm_source(file, source, &[]).expect("interp run");
+    let executable = rsscript::reg_vm_compile_source(file, source).expect("source compiles");
+    let (osr, stats) = executable
+        .eval_main_with_args_native_osr_with_stats(std::iter::empty::<String>())
+        .expect("osr native run");
+    assert_eq!(
+        interp.stdout, osr.stdout,
+        "OSR × J2 Float-capture closure loop must be byte-identical to the interpreter \
+         (the f64 capture must be bit-reinterpreted, not int-converted)"
+    );
+    assert!(
+        stats.osr_entries > 0,
+        "a capturing monomorphic param-handle closure with a FLOAT capture called in an \
+         I/O-tangled hot loop must inline (bit-reinterpreting the f64 capture) and OSR \
+         natively: {stats:?}",
+    );
+}
+
+/// OSR × J2 NEGATIVE test (Pending #2) — a closure capturing a NON-scalar value that
+/// HAPPENS TO BE float-typed (`scales: List<Float>`, a heap value) MUST NOT OSR. The
+/// Float broadening admits only FLAT scalar captures (Int/Bool/Float); a heap capture
+/// — even one whose elements are Floats — cannot be materialized as a scalar via the
+/// `closure_capture` helper, so the inline gate's `captures_all_scalar` profile bit
+/// goes false on the first observation, the site never inlines, the per-iteration
+/// `CallClosure` keeps the loop off the native subset, and `osr_entries` MUST stay 0.
+/// Output must still be interpreter-identical. This guards the Float broadening from
+/// over-reaching: widening flat-scalar captures to include `Float` must NOT start
+/// admitting a heap aggregate just because it carries Floats (which would read garbage
+/// bits from a heap handle's slot).
+#[cfg(feature = "native-jit")]
+#[test]
+fn native_osr_j2_float_heap_capture_closure_does_not_osr() {
+    let source = "\
+fn apply_loop(g: read Fn(Int) -> Float, limit: Int, seed: Float) -> Float {
+    Log.write(message: read \"begin\")
+    let mut i = 0
+    let mut total = seed
+    while i < limit {
+        total = total + g(i)
+        i = i + 1
+    }
+    Log.write(message: read String.from_float(value: total))
+    return total
+}
+
+fn main() -> Unit {
+    let scales: List<Float> = [1.5, 2.25, 3.75]
+    let g = fn(value) captures(read scales) effects(pure) {
+        let n = List.len<Float>(list: read scales)
+        return List.get<Float>(list: read scales, index: value - (value / n) * n)
+    }
+    Log.write(message: read String.from_float(value: apply_loop(g: read g, limit: read 4000, seed: read 0.0)))
+    return Unit
+}
+";
+    let file = "jit-osr-j2-float-heap-capture-closure.rss";
+    let interp = common::run_vm_source(file, source, &[]).expect("interp run");
+    let executable = rsscript::reg_vm_compile_source(file, source).expect("source compiles");
+    let (osr, stats) = executable
+        .eval_main_with_args_native_osr_with_stats(std::iter::empty::<String>())
+        .expect("osr native run");
+    assert_eq!(
+        interp.stdout, osr.stdout,
+        "a closure capturing a List<Float> (heap) must still be interpreter-identical under OSR"
+    );
+    assert_eq!(
+        stats.osr_entries, 0,
+        "a closure capturing a heap value (List<Float>) must NOT OSR even though its \
+         elements are Floats — only FLAT scalar captures inline: {stats:?}",
+    );
+}
+
 /// OSR × J2 NEGATIVE test — a closure with a NON-scalar (heap) capture MUST NOT OSR.
 /// `g` captures `tag: List<Int>` (a heap value) and reads its length each call, so
 /// the capture cannot be materialized as a scalar via the `closure_capture` helper:

crates/vm-jit/src/lib.rs

@@ -1334,7 +1334,15 @@ fn validate(program: &JitFunction) -> Result<(), JitError> {
             }
             JitInstr::ClosureCapture { dst, base, .. } => {
                 require_class(*base, JitValueType::Handle, "ClosureCapture base")?;
-                require_class(*dst, JitValueType::Int, "ClosureCapture result")?;
+                // The capture result `dst` is Int-class (an Int/Bool capture, used
+                // as i64 directly) or Float-class (a Float capture, whose i64 slot
+                // is `f64::to_bits` and is bit-reinterpreted to f64 in codegen).
+                if !matches!(class(*dst), JitValueType::Int | JitValueType::Float) {
+                    return Err(JitError(format!(
+                        "ClosureCapture result: register {dst} is {:?}, expected Int or Float",
+                        class(*dst)
+                    )));
+                }
             }
             JitInstr::FieldHandle { dst, base, .. } => {
                 require_class(*base, JitValueType::Handle, "FieldHandle base")?;
@@ -2634,7 +2642,17 @@ fn build_function(
                     deopt!(i),
                 );
                 bcx.switch_to_block(cont);
-                bcx.def_var(reg(*dst), result);
+                // The helper returns the capture's raw scalar bits as i64. For an
+                // Int/Bool-class `dst` that IS the value; for a Float-class `dst`
+                // the i64 is `f64::to_bits`, so bit-reinterpret it to f64 (NOT an
+                // integer-to-float conversion) — the same convention by which a
+                // float register's arg slot is loaded as f64 on entry.
+                let stored = if program.is_float(*dst) {
+                    bcx.ins().bitcast(types::F64, MemFlags::new(), result)
+                } else {
+                    result
+                };
+                bcx.def_var(reg(*dst), stored);
             }
             JitInstr::FieldHandle { dst, base, slot } => {
                 // Fetch a heap field (e.g. a stored closure) as a fresh handle.