fix(codegen): retry oversized RS4GC functions with shadow roots (#8679)

changelog.d/8679-rs4gc-budget-spill-retry.md

@@ -0,0 +1,4 @@
+Fixed large native functions whose statepoint rewrite crossed the LLVM
+instruction budget but fell below the root-spill estimate threshold. Perry now
+re-lowers only those functions with precise shadow-frame roots and retries the
+same optimization pipeline instead of refusing the whole codegen unit.

crates/perry-codegen/src/codegen/helpers.rs

@@ -387,9 +387,10 @@ pub(crate) fn inline_hot_small_max_call_sites() -> u32 {
 /// fan-out it avoids (an ~8M function spilled in 303 s vs 180 s fanned out,
 /// #8620) — and above it fan-out risks not finishing and the shadow frame wins.
 /// The former 4M default fired on ~8M functions that fan out fine in minutes.
-/// The post-RS4GC instruction-budget assertion (#8586, inprocess.rs) backstops
-/// any function this estimate misses: it fails loudly rather than hanging, so
-/// raising the threshold is safe.
+/// The post-RS4GC instruction budget (#8586/#8679, inprocess.rs) backstops any
+/// function this estimate misses: it re-lowers that function onto a precise
+/// shadow frame and retries before LLVM's optimizer can hang, so raising the
+/// estimate threshold is safe.
 ///
 /// `PERRY_ROOT_SPILL_RELOCATIONS=<n>` overrides it; `0` disables spilling
 /// (every function stays on native statepoints, the pre-#8583 behavior).

crates/perry-codegen/src/codegen/mod.rs

@@ -3353,58 +3353,91 @@ pub fn compile_module(hir: &HirModule, opts: CompileOptions) -> Result<Vec<u8>>
             progress.phase(3, "object ready; releasing generated IR");
             return result;
         }
-        let units = llmod.render_codegen_units(n_units);
-        log::debug!(
-            "perry-codegen: split '{}' into {} codegen units",
-            hir.name,
-            units.len()
-        );
-        // #7154: dump the units. The comment above used to claim `PERRY_SAVE_LL`
-        // took the single-text path — it never did; this `return` fires before
-        // the `PERRY_SAVE_LL` write below. So `--trace llvm` silently emitted
-        // NOTHING for any module past `MIN_CALLABLES_TO_SPLIT`, i.e. exactly the
-        // largest modules, which is where a static IR audit
-        // (`scripts/gc_root_dominance_check.py`) most needs to look — a corpus
-        // that quietly omits its biggest members makes a clean verdict
-        // meaningless. One file per unit, not one concatenation: the units are
-        // already materialized here, so this adds no peak.
-        if let Ok(save_dir) = std::env::var("PERRY_SAVE_LL") {
-            for (i, unit) in units.iter().enumerate() {
-                let filename = format!("{}/{}.unit{}.ll", save_dir, module_prefix, i);
-                let _ = std::fs::write(&filename, unit);
+        loop {
+            let units = llmod.render_codegen_units(n_units);
+            log::debug!(
+                "perry-codegen: split '{}' into {} codegen units",
+                hir.name,
+                units.len()
+            );
+            // #7154: dump the units. The comment above used to claim
+            // `PERRY_SAVE_LL` took the single-text path — it never did; this
+            // return fires before the write below. One file per unit, not one
+            // concatenation: the units are already materialized here, so this
+            // adds no peak.
+            if let Ok(save_dir) = std::env::var("PERRY_SAVE_LL") {
+                for (i, unit) in units.iter().enumerate() {
+                    let filename = format!("{}/{}.unit{}.ll", save_dir, module_prefix, i);
+                    let _ = std::fs::write(&filename, unit);
+                }
+            }
+            match crate::linker::compile_units_to_object(&units, opts.target.as_deref()) {
+                Ok(object) => return Ok(object),
+                Err(error) if apply_rs4gc_budget_retry(&mut llmod, &error)? => continue,
+                Err(error) => return Err(error),
             }
         }
-        return crate::linker::compile_units_to_object(&units, opts.target.as_deref());
     }
 
     // exp/llvm-inprocess Phase 2: `PERRY_LLVM_INPROCESS=native` constructs
     // function bodies through the LLVM C API (only the module skeleton is
     // textual); `=diff` builds both arms and diffs them. Unit-split and
     // emit_ir_only paths above stay textual (they fall into the in-process
     // *transport* under these values, so no clang subprocess either way).
-    if let Some(result) = try_native_construction(&llmod, opts.target.as_deref(), &module_prefix) {
+    if let Some(result) =
+        try_native_construction(&mut llmod, opts.target.as_deref(), &module_prefix)
+    {
         return result;
     }
 
-    let ll_text = llmod.to_ir();
-    log::debug!(
-        "perry-codegen: emitted {} bytes of LLVM IR for '{}' ({} interned strings)",
-        ll_text.len(),
-        hir.name,
-        strings.len()
-    );
-    // Save .ll files when PERRY_SAVE_LL=<dir> is set
-    if let Ok(save_dir) = std::env::var("PERRY_SAVE_LL") {
-        let filename = format!("{}/{}.ll", save_dir, module_prefix);
-        let _ = std::fs::write(&filename, &ll_text);
-    }
-    if opts.emit_ir_only {
-        Ok(ll_text.into_bytes())
-    } else {
-        crate::linker::compile_ll_to_object(&ll_text, opts.target.as_deref())
+    loop {
+        let ll_text = llmod.to_ir();
+        log::debug!(
+            "perry-codegen: emitted {} bytes of LLVM IR for '{}' ({} interned strings)",
+            ll_text.len(),
+            hir.name,
+            strings.len()
+        );
+        // Save .ll files when PERRY_SAVE_LL=<dir> is set
+        if let Ok(save_dir) = std::env::var("PERRY_SAVE_LL") {
+            let filename = format!("{}/{}.ll", save_dir, module_prefix);
+            let _ = std::fs::write(&filename, &ll_text);
+        }
+        if opts.emit_ir_only {
+            return Ok(ll_text.into_bytes());
+        }
+        match crate::linker::compile_ll_to_object(&ll_text, opts.target.as_deref()) {
+            Ok(object) => return Ok(object),
+            Err(error) if apply_rs4gc_budget_retry(&mut llmod, &error)? => continue,
+            Err(error) => return Err(error),
+        }
     }
 }
 
+/// Consume the typed post-RS4GC budget signal on text-transport paths. The
+/// native constructors have the same loop closer to their LLVM modules; text
+/// compilation returns through `linker`, so its retry belongs at the last
+/// point where the lowering-owned `LlModule` is still available.
+#[cfg(feature = "llvm-inprocess")]
+fn apply_rs4gc_budget_retry(
+    llmod: &mut crate::module::LlModule,
+    error: &anyhow::Error,
+) -> Result<bool> {
+    let Some(violations) = crate::inprocess::rs4gc_budget_retry(error) else {
+        return Ok(false);
+    };
+    crate::native_emit::apply_budget_spill_retry(llmod.functions_mut(), &violations)?;
+    Ok(true)
+}
+
+#[cfg(not(feature = "llvm-inprocess"))]
+fn apply_rs4gc_budget_retry(
+    _llmod: &mut crate::module::LlModule,
+    _error: &anyhow::Error,
+) -> Result<bool> {
+    Ok(false)
+}
+
 /// exp/llvm-inprocess: unit-split twin of [`try_native_construction`].
 #[cfg(feature = "llvm-inprocess")]
 fn try_native_units(
@@ -3444,7 +3477,7 @@ fn try_native_units(
 /// in-process mode is requested, so the flag can never silently no-op.
 #[cfg(feature = "llvm-inprocess")]
 fn try_native_construction(
-    llmod: &crate::module::LlModule,
+    llmod: &mut crate::module::LlModule,
     target: Option<&str>,
     module_prefix: &str,
 ) -> Option<Result<Vec<u8>>> {
@@ -3471,7 +3504,7 @@ fn try_native_construction(
 
 #[cfg(not(feature = "llvm-inprocess"))]
 fn try_native_construction(
-    _llmod: &crate::module::LlModule,
+    _llmod: &mut crate::module::LlModule,
     _target: Option<&str>,
     _module_prefix: &str,
 ) -> Option<Result<Vec<u8>>> {

crates/perry-codegen/src/function.rs

@@ -302,12 +302,35 @@ impl LlFunction {
     /// into every caller's hot loop. Skip the frame entirely; the
     /// to_ir() rewrite pass keys off `shadow_frame_slot.is_some()`,
     /// so no matching pop is emitted either.
-    /// #8583: route this function's precise roots through the heap shadow
-    /// frame instead of native statepoints. Must be called BEFORE
-    /// `enable_shadow_frame` / `enable_post_init_shadow_frame` so the frame is
-    /// built in shadow form. No effect once a frame has been emitted.
-    pub fn request_shadow_frame_spill(&mut self) {
+    /// #8583/#8679: route this function's precise roots through the heap
+    /// shadow frame instead of native statepoints.
+    ///
+    /// The estimate-driven path calls this before `enable_shadow_frame`, while
+    /// the post-RS4GC budget retry calls it after lowering is complete. In the
+    /// latter case the native-root path deliberately retained the original
+    /// `js_shadow_slot_bind` calls until final rendering, so converting the
+    /// recorded stack-map request back into a shadow-frame push is a complete
+    /// re-lowering: final rendering keeps those binds, adds the matching pops,
+    /// and drops the GC strategy so RS4GC skips the function on retry.
+    ///
+    /// Returns `true` only when this call changed the lowering. A retry driver
+    /// uses that to reject an impossible second retry instead of looping.
+    pub fn request_shadow_frame_spill(&mut self) -> bool {
+        if self.force_shadow_frame {
+            return false;
+        }
         self.force_shadow_frame = true;
+        self.stack_map_requested = false;
+        if self.shadow_frame_requested
+            && self.shadow_frame_slot.is_none()
+            && self.stack_map_slot_count != 0
+        {
+            self.emit_shadow_frame_push(
+                self.stack_map_slot_count,
+                self.shadow_frame_post_init_region,
+            );
+        }
+        true
     }
 
     /// Whether this function spills its roots to the shadow frame (#8583).
@@ -1363,6 +1386,53 @@ mod define_header_tests {
         );
     }
 
+    /// #8679's budget is learned only after RS4GC, so the durable fallback
+    /// necessarily asks an already-lowered function to change root lowering.
+    /// This pins that late request to the same complete shadow-frame shape as
+    /// the estimate-driven early request, including balanced return pops.
+    #[test]
+    fn a_post_lowering_spill_request_rebuilds_the_shadow_frame() {
+        use crate::codegen::helpers::NativeRootsPin;
+        use crate::types::{I64, PTR};
+        const STRATEGY: &str = "gc \"statepoint-example\"";
+
+        let _native = NativeRootsPin::native();
+        let mut function = LlFunction::new("late_spill", crate::types::VOID, vec![]);
+        function.enable_post_init_shadow_frame(0);
+        let idx = function
+            .reserve_shadow_slot()
+            .expect("native lowering reserves a precise-root slot");
+        let root = function.alloca_entry(I64);
+        function.entry_allocas_push_store(I64, "0", &root);
+        function.entry_setup_call_void(
+            "js_shadow_slot_bind",
+            &[(crate::types::I32, &idx.to_string()), (PTR, &root)],
+        );
+        function.mark_entry_init_boundary();
+        let entry = function.create_block("entry");
+        let _ = entry.call(I64, "may_collect", &[]);
+        entry.ret_void();
+
+        let native_ir = function.to_ir();
+        assert!(native_ir.contains(STRATEGY));
+        assert!(!native_ir.contains("@js_shadow_frame_enter"));
+        assert!(!native_ir.contains("@js_shadow_slot_bind"));
+
+        assert!(
+            function.request_shadow_frame_spill(),
+            "the first late request must change the lowering"
+        );
+        assert!(
+            !function.request_shadow_frame_spill(),
+            "a repeated request must report that no retry progress is possible"
+        );
+        let shadow_ir = function.to_ir();
+        assert!(!shadow_ir.contains(STRATEGY), "{shadow_ir}");
+        assert!(shadow_ir.contains("call ptr @js_shadow_frame_enter(i32 1)"));
+        assert!(shadow_ir.contains("call void @js_shadow_slot_bind(i32 0"));
+        assert!(shadow_ir.contains("call void @js_shadow_frame_pop(i64"));
+    }
+
     /// `force_external` drops only the linkage keyword. The codegen-unit path
     /// depends on that and on nothing else changing.
     #[test]