jit: collecting bigint residuals, nursery call frames, vable store-back

majit/gc-root-brackets.baseline.json

@@ -14,6 +14,7 @@
     "unbracketed_calls_fns": 826,
     "unmatched_seeds": [
       "majit_gc::standalone_alloc_fast_nursery_collecting_typed_rooted",
+      "majit_gc::standalone_alloc_fast_nursery_collecting_typed_roots",
       "majit_gc::standalone_alloc_nursery_collecting_typed_rooted"
     ]
   },
@@ -32,6 +33,7 @@
     "unbracketed_calls_fns": 836,
     "unmatched_seeds": [
       "majit_gc::standalone_alloc_fast_nursery_collecting_typed_rooted",
+      "majit_gc::standalone_alloc_fast_nursery_collecting_typed_roots",
       "majit_gc::standalone_alloc_nursery_collecting_typed_rooted"
     ]
   }

majit/majit-backend-dynasm/src/aarch64/assembler.rs

@@ -6108,11 +6108,6 @@ impl<'a> AssemblerARM64<'a> {
                 .copied()
                 .expect("call_assembler missing rewritten jitframe arg");
             let vable_loc = arglocs.get(1).copied();
-            // aarch64/regalloc.py:661-664 routes CALL_ASSEMBLER through
-            // `_call(..., gc_level=2)`, which spills all managed registers.
-            // x19 is already saved by the JIT prologue, so use it as scratch
-            // here without an extra call-site stack save.
-            dynasm!(self.mc ; .arch aarch64 ; mov x19, x29);
             self.emit_load_to_rax(frame_loc);
 
             let descr_arc = op.getdescr();
@@ -6143,9 +6138,6 @@ impl<'a> AssemblerARM64<'a> {
 
             if !is_resolved {
                 let force_addr = crate::call_assembler_force_fn_addr() as i64;
-                dynasm!(self.mc ; .arch aarch64
-                    ; mov x29, x19
-                );
                 if force_addr != 0 {
                     if let Some(vloc) = vable_loc {
                         self.emit_load_to_rax(vloc);
@@ -6190,9 +6182,13 @@ impl<'a> AssemblerARM64<'a> {
             } else if let Some(entry_label) = self.self_entry_label {
                 dynasm!(self.mc ; .arch aarch64 ; bl =>entry_label);
             }
-            dynasm!(self.mc ; .arch aarch64
-                ; mov x29, x19
-            );
+            // aarch64/callbuilder.py `CallBuilderARM64.pop_gcmap` calls
+            // `AssemblerARM64._reload_frame_if_necessary` before clearing
+            // the map.  The callee footer restores the pre-call x29 from the
+            // C stack, but that address may be stale after a minor collection
+            // moved the caller jitframe.  Reload x29 from the shadow stack;
+            // keeping the old address in an allocatable callee-saved register
+            // is both non-orthodox and unsound.
             self.pop_pending_call_gcmap_after_collect(pushed_gcmap);
 
             let fast_path = self.mc.new_dynamic_label();

majit/majit-backend-dynasm/src/lib.rs

@@ -219,7 +219,7 @@ use std::sync::OnceLock;
 /// _prepare_resume_from_failure` hands it to the resumed frame so an
 /// exception guard unwinds into its handler instead of resuming the
 /// no-exception continuation.  `0` = no pending exception.
-pub type BlackholeFn = fn(usize, *const i64, usize, *const i64, usize, i64) -> Option<i64>;
+pub type BlackholeFn = fn(usize, *const i64, usize, *const i64, usize, i64, usize) -> Option<i64>;
 
 /// Bridge compilation: raw values, descr identity, and optional GUARD_VALUE
 /// operand → compiled?
@@ -811,13 +811,18 @@ fn handle_fail_resume_guard(
     // re-reading `jf_guard_exc` here would observe the post-`grab_exc_value`
     // null and drop the exception.
     let bh_result = CA_BLACKHOLE_FN.get().and_then(|blackhole| {
+        // compile.py ResumeGuardForcedDescr.handle_fail reads savedata from
+        // this same deadframe after any bridge attempt. The libc jitframe
+        // tracer updates the field if that attempt collected the nursery.
+        let savedata = unsafe { (*frame_ptr).jf_savedata };
         blackhole(
             descr_raw,
             raw_values.as_ptr(),
             raw_values.len(),
             raw_values.as_ptr(),
             raw_values.len(),
             guard_exc_root.0 as i64,
+            savedata,
         )
     });
     if let Some(bh_result) = bh_result {

majit/majit-backend-dynasm/src/regalloc.rs

@@ -2741,8 +2741,11 @@ impl<'a> RegAlloc<'a> {
             | GuardKind::Overflow
             | GuardKind::NotInvalidated
             | GuardKind::FutureCondition
-            | GuardKind::NotForced
             | GuardKind::AlwaysFails => self.consider_guard_no_args_j2(fail_args, i, output),
+            GuardKind::NotForced if op.opcode == OpCode::GuardNotForced2 => {
+                self.consider_guard_not_forced_2_j2(fail_args, i, output)
+            }
+            GuardKind::NotForced => self.consider_guard_no_args_j2(fail_args, i, output),
             GuardKind::Exception => {
                 self.consider_guard_exception_j2(args, fail_args, op, i, output)
             }
@@ -5230,29 +5233,31 @@ impl<'a> RegAlloc<'a> {
         }
     }
 
-    /// llsupport/regalloc.py locs_for_call_assembler parity.
-    /// RPython syncs args to stack, then before_call spills everything.
-    /// We force-sync register args to frame first via _sync_var_to_stack,
-    /// then before_call spills remaining. arglocs after before_call are
-    /// all Frame or Immed — no register-clobber issues during calloc.
+    /// llsupport/regalloc.py `locs_for_call_assembler` parity.
+    ///
+    /// RPython syncs only argument 1 (the virtualizable) to the frame. The
+    /// callee jitframe in argument 0 stays in its current location, captured
+    /// before `before_call`.
     fn consider_call_assembler(&mut self, op: &Op, i: usize, output: &mut Vec<RegAllocOp>) {
-        // llsupport/regalloc.py: self.rm._sync_var_to_stack(op.getarg(k))
-        // Force all register-held args to frame before before_call.
-        for arg in op.getarglist().iter() {
-            if arg.is_constant() {
-                continue;
-            }
-            let arg = arg.to_opref();
-            let tp = self.tp(arg);
+        assert!(matches!(op.num_args(), 1 | 2));
+        if op.num_args() == 2 {
+            let vable = op.arg(1).to_opref();
+            let tp = self.tp(vable);
             if tp == Type::Float {
                 self.xrm
-                    ._sync_var_to_stack(arg, tp, &mut self.longevity, &mut self.fm);
+                    ._sync_var_to_stack(vable, tp, &mut self.longevity, &mut self.fm);
             } else {
                 self.rm
-                    ._sync_var_to_stack(arg, tp, &mut self.longevity, &mut self.fm);
+                    ._sync_var_to_stack(vable, tp, &mut self.longevity, &mut self.fm);
             }
         }
 
+        let mut arglocs = Vec::with_capacity(op.num_args());
+        for arg in op.getarglist().iter() {
+            let arg = arg.to_opref();
+            arglocs.push(self.loc(arg, self.tp(arg)));
+        }
+
         let type_index = OpTypeIndex::from_parts(
             self.inputargs,
             self.operations,
@@ -5276,14 +5281,6 @@ impl<'a> RegAlloc<'a> {
             &type_index,
         );
 
-        // After before_call, all args are in Frame or Const — safe for calloc.
-        let mut arglocs: Vec<Loc> = Vec::new();
-        for arg in op.getarglist().iter() {
-            let arg = arg.to_opref();
-            let tp = self.tp(arg);
-            arglocs.push(self.loc_must_exist(arg, tp));
-        }
-
         let result_tp = op.opcode.result_type();
         let result_loc = if result_tp != Type::Void {
             let r = if result_tp == Type::Float {
@@ -5307,20 +5304,24 @@ impl<'a> RegAlloc<'a> {
         i: usize,
         output: &mut Vec<RegAllocOp>,
     ) {
-        for &arg in args {
-            if arg.is_constant() {
-                continue;
-            }
-            let tp = self.tp(arg);
+        assert!(matches!(args.len(), 1 | 2));
+        if args.len() == 2 {
+            let vable = args[1];
+            let tp = self.tp(vable);
             if tp == Type::Float {
                 self.xrm
-                    ._sync_var_to_stack(arg, tp, &mut self.longevity, &mut self.fm);
+                    ._sync_var_to_stack(vable, tp, &mut self.longevity, &mut self.fm);
             } else {
                 self.rm
-                    ._sync_var_to_stack(arg, tp, &mut self.longevity, &mut self.fm);
+                    ._sync_var_to_stack(vable, tp, &mut self.longevity, &mut self.fm);
             }
         }
 
+        let mut arglocs = Vec::with_capacity(args.len());
+        for &arg in args {
+            arglocs.push(self.loc(arg, self.tp(arg)));
+        }
+
         let type_index = OpTypeIndex::from_parts(
             self.inputargs,
             self.operations,
@@ -5344,12 +5345,6 @@ impl<'a> RegAlloc<'a> {
             &type_index,
         );
 
-        let mut arglocs: Vec<Loc> = Vec::new();
-        for &arg in args {
-            let tp = self.tp(arg);
-            arglocs.push(self.loc_must_exist(arg, tp));
-        }
-
         let result_tp = op.opcode.result_type();
         let result_loc = if result_tp != Type::Void {
             let dst = dst.unwrap_or(op.pos.get());

majit/majit-backend-dynasm/src/runner.rs

@@ -386,6 +386,9 @@ fn register_active_hooks(supports_guard_gc_type: bool) {
     majit_gc::set_active_alloc_nursery_collecting_typed_rooted(Some(
         dynasm_alloc_nursery_collecting_typed_rooted,
     ));
+    majit_gc::set_active_alloc_nursery_collecting_typed_roots(Some(
+        dynasm_alloc_nursery_collecting_typed_roots,
+    ));
     majit_gc::set_active_alloc_oldgen_typed(Some(dynasm_alloc_oldgen_typed));
     majit_gc::set_active_collect_generation(Some(dynasm_collect_generation));
     majit_gc::set_active_collect_step(Some(dynasm_collect_step));
@@ -706,6 +709,35 @@ unsafe fn dynasm_alloc_nursery_collecting_typed_rooted(
     }
 }
 
+unsafe fn dynasm_alloc_nursery_collecting_typed_roots(
+    type_id: u32,
+    size: usize,
+    roots: *mut GcRef,
+    root_count: usize,
+    needs_write_barrier: *mut bool,
+) -> GcRef {
+    if let Some(r) = gc_box::with_mut(|g| unsafe {
+        g.alloc_fast_nursery_collecting_typed_roots(
+            type_id,
+            size,
+            roots,
+            root_count,
+            needs_write_barrier,
+        )
+    }) {
+        return r;
+    }
+    unsafe {
+        majit_gc::standalone_alloc_fast_nursery_collecting_typed_roots(
+            type_id,
+            size,
+            roots,
+            root_count,
+            needs_write_barrier,
+        )
+    }
+}
+
 /// Host-side old-gen allocation trampoline. Used by
 /// pyre-object allocators (`w_int_new`, `w_float_new`) whose
 /// callers cannot register the returned pointer as a GC root before
@@ -3218,6 +3250,7 @@ impl Backend for DynasmBackend {
         // grab_exc_value (llmodel.py): read jf_guard_exc off the deadframe
         // tip before the libc jitframe chain is freed (same as execute_token).
         let exception_value = GcRef(unsafe { (*result_jf).jf_guard_exc });
+        let savedata = GcRef(unsafe { (*result_jf).jf_savedata });
         let guard_value_operand = majit_backend::guard_value_counter_slot(descr_fd)
             .map(|slot| unsafe { crate::llmodel::get_int_value_direct(result_jf, slot) as i64 });
 
@@ -3232,7 +3265,7 @@ impl Backend for DynasmBackend {
             outputs,
             typed_outputs,
             exit_layout,
-            savedata: None,
+            savedata: (!savedata.is_null()).then_some(savedata),
             exception_value,
             fail_index: descr_fd.fail_index_per_trace(),
             trace_id: descr_fd.trace_id(),
@@ -3316,6 +3349,27 @@ impl Backend for DynasmBackend {
         }
     }
 
+    fn set_savedata_ref(&self, frame: &mut DeadFrame, data: GcRef) {
+        match frame {
+            DeadFrame::JitFrame(jf) => {
+                // llmodel.py set_savedata_ref is a GCREF field store.
+                majit_gc::gc_write_barrier(jf.jf_gcref());
+                jf.set_savedata_ref(data);
+            }
+            DeadFrame::LibcJitFrame(jf) => jf.set_savedata_ref(data),
+            DeadFrame::Boxed(_) => panic!("dynasm deadframe is a jitframe"),
+        }
+    }
+
+    fn get_savedata_ref(&self, frame: &DeadFrame) -> Option<GcRef> {
+        let data = match frame {
+            DeadFrame::JitFrame(jf) => jf.get_savedata_ref(),
+            DeadFrame::LibcJitFrame(jf) => jf.get_savedata_ref(),
+            DeadFrame::Boxed(_) => panic!("dynasm deadframe is a jitframe"),
+        };
+        (!data.is_null()).then_some(data)
+    }
+
     fn clear_stored_exception(&self) {
         crate::jit_exc_clear();
     }