jit: trace the cranelift force token as a GC slot, plus _pickle and JitCodeBuilder follow-ups

majit/majit-backend-cranelift/src/compiler.rs

@@ -4698,21 +4698,13 @@ fn build_known_values_set(inputargs: &[InputArg], ops: &[Op]) -> IndexSet<u32> {
     known
 }
 
-fn build_force_token_set(inputargs: &[InputArg], ops: &[Op]) -> IndexSet<u32> {
-    // FORCE_TOKEN (resoperation.py:1090 'FORCE_TOKEN/0/r') yields the raw
-    // jitframe handle. Its Ref result is the frame pointer, not a heap GCREF,
-    // so it is excluded from the ref-root slots the GC traces and relocates.
-    let mut force_tokens = IndexSet::new();
-    for (op_idx, op) in ops.iter().enumerate() {
-        if op.pos.get().is_none() {
-            continue;
-        }
-        if op.opcode == OpCode::ForceToken {
-            let result_var = op_var_index(op, op_idx, inputargs.len()) as u32;
-            force_tokens.insert(result_var);
-        }
-    }
-    force_tokens
+fn build_force_token_set(_inputargs: &[InputArg], _ops: &[Op]) -> IndexSet<u32> {
+    // FORCE_TOKEN is a GCREF to the active JITFRAME
+    // (`virtualizable.py:315-318`, `resoperation.py:1090`). Keep its in-frame
+    // copies in the ordinary Ref root set so moving collectors update them.
+    // The empty compatibility set leaves the existing exit-layout plumbing in
+    // place while giving FORCE_TOKEN the same treatment as every other Ref.
+    IndexSet::new()
 }
 
 /// Auxiliary LABEL type overrides on top of `OpTypeIndex`.
@@ -8753,9 +8745,10 @@ impl CraneliftBackend {
         //    overwrite by index.  That disjoint home is why non-refs (incl.
         //    inputargs) are safe to demote: the ref-inputarg staleness argument
         //    does not apply because the home is re-seeded on every LABEL entry
-        //    edge (preamble fall-through and loader re-entry).  Floats, SIMD
-        //    lanes, and force tokens are excluded — an 8-byte home cannot hold a
-        //    vector and float typing must round-trip through resume.
+        //    edge (preamble fall-through and loader re-entry). Floats and SIMD
+        //    lanes are excluded — an 8-byte home cannot hold a vector and float
+        //    typing must round-trip through resume. FORCE_TOKEN is a Ref and uses
+        //    the forwarded ref-root home.
         if !loop_phi_keep_by_label.is_empty() {
             let input_idxs: indexmap::IndexSet<u32> = inputargs.iter().map(|ia| ia.index).collect();
             loop_phi_keep_by_label.retain(|&label_idx, keep| {

majit/majit-metainterp/src/jitcode/assembler.rs

@@ -624,6 +624,28 @@ impl JitCodeBuilder {
         headerless: bool,
         fields: &[(usize, bool, &str)],
     ) {
+        // Every emit site re-registers the layout it accesses, so the same few
+        // type_ids arrive hundreds of times.  When the spec already lists an
+        // offset for each incoming field the merge below pushes nothing, and
+        // the re-sort/re-index that follow it are no-ops on an already sorted,
+        // already indexed vector — the branch never touches `size`,
+        // `is_gc_managed` or `headerless`, so returning here is exactly
+        // equivalent and skips building `new_fields` (one owned String per
+        // field) only to discard it.
+        if self
+            .struct_size_specs
+            .get(&type_id)
+            .is_some_and(|existing| {
+                fields.iter().all(|&(offset, _, _)| {
+                    existing
+                        .all_fielddescrs
+                        .iter()
+                        .any(|ef| ef.offset == offset)
+                })
+            })
+        {
+            return;
+        }
         let new_fields = Self::field_specs_from_layout(fields);
         // Merge into existing spec if present — each getfield/setfield
         // site registers only the field it accesses, so the complete
@@ -760,8 +782,7 @@ impl JitCodeBuilder {
         type_id: u64,
         field_name: &str,
     ) -> u16 {
-        let parent = self.struct_size_specs.get(&type_id).cloned();
-        let Some(parent_spec) = parent.as_ref() else {
+        let Some(parent_spec) = self.struct_size_specs.get(&type_id) else {
             return self.add_scalar_field_descr(offset, field_type);
         };
         let (field_flag, is_field_signed) = match field_type {
@@ -793,6 +814,22 @@ impl JitCodeBuilder {
             field_slot_in(&parent_spec.all_fielddescrs, field_name, offset)
                 .map(|idx| (idx, parent_spec.all_fielddescrs[idx].name.clone()))
                 .unwrap_or((0, String::new()));
+        // Carry the scalars only.  `patch_field_descr_parents`, called
+        // unconditionally from `try_finish` after the decline early-return,
+        // replaces this snapshot with `struct_size_specs`' final merged spec
+        // for the same `type_id` — and entries there are only inserted or
+        // merged, never removed, so a type_id present now is present then.
+        // `type_id` is therefore the only load-bearing part; deep-copying
+        // `all_fielddescrs` (a String per field) once per field-descr mint
+        // would build a table that is overwritten before anything reads it.
+        let parent = Some(BhSizeSpec {
+            size: parent_spec.size,
+            type_id: parent_spec.type_id,
+            vtable: parent_spec.vtable,
+            is_gc_managed: parent_spec.is_gc_managed,
+            headerless: parent_spec.headerless,
+            all_fielddescrs: Vec::new(),
+        });
         self.add_bh_descr(CanonicalBhDescr::Field {
             offset,
             field_size: scalar_size(field_type),

majit/majit-metainterp/src/jitdriver.rs

@@ -2984,6 +2984,13 @@ impl<S: JitState> JitDriver<S> {
                                 // false now and true at the next visit of this header.
                                 if let Some(ctx) = self.meta.trace_ctx() {
                                     if attempted {
+                                        // The two sibling close sites report a declined
+                                        // attempt under the same tally; this one is the
+                                        // third.  Bump it on the same condition the latch
+                                        // uses, so the census counts closes an optimizer
+                                        // pass actually rejected and not headers the gate
+                                        // above skipped.
+                                        crate::mc_diag_bump(50); // bridge_declined_close
                                         ctx.note_cross_loop_close_declined(target_key);
                                     }
                                     ctx.close_greens = None;

majit/majit-metainterp/src/jitprof.rs

@@ -21,9 +21,9 @@ use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 #[cfg(not(target_arch = "wasm32"))]
-use std::time::Instant;
+pub use std::time::Instant;
 #[cfg(target_arch = "wasm32")]
-use wasm_clock::Instant;
+pub use wasm_clock::Instant;
 
 /// Monotonic substitute for `std::time::Instant` on wasm32-unknown-unknown,
 /// which has no clock (`Instant::now()` there panics). The profiler only needs
@@ -33,7 +33,7 @@ use wasm_clock::Instant;
 /// `Instant::now()` / `saturating_duration_since` surface the timer uses, so
 /// the timing code below is platform-agnostic.
 #[cfg(target_arch = "wasm32")]
-mod wasm_clock {
+pub mod wasm_clock {
     use core::sync::atomic::{AtomicU64, Ordering};
     use std::time::Duration;
 

majit/majit-metainterp/src/pyjitpl.rs

@@ -21,6 +21,7 @@ use std::sync::{
     atomic::{AtomicU64, Ordering},
 };
 
+use crate::jitprof::Instant;
 use crate::optimizeopt::optimizer::{Optimizer, PendingBridgeRd};
 use majit_backend::{Backend, ExitRecoveryLayout, JitCellToken};
 #[cfg(all(feature = "cranelift", not(target_arch = "wasm32")))]
@@ -6061,6 +6062,7 @@ impl<M: Clone> MetaInterp<M> {
             Vec<majit_ir::OpRc>,
             crate::optimizeopt::unroll::ExportedState,
         )> = None;
+        let optimize_start = Instant::now();
         let optimize_result = if no_unroll {
             if crate::majit_log_enabled() {
                 eprintln!(
@@ -6263,6 +6265,7 @@ impl<M: Clone> MetaInterp<M> {
                 optimized_ops
             }
         };
+        let opt_time = Instant::now().saturating_duration_since(optimize_start);
         let num_ops_after = optimized_ops.len();
         if crate::majit_log_enabled() {
             eprintln!(
@@ -6625,12 +6628,14 @@ impl<M: Clone> MetaInterp<M> {
                 compiled_ops.len()
             );
         }
+        let compile_start = Instant::now();
         let compile_result = {
             let _backend_scope = self.staticdata.profiler.enter_backend();
             std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                 self.backend.compile_loop(&inputargs, &compiled_ops, &token)
             }))
         };
+        let compile_time = Instant::now().saturating_duration_since(compile_start);
         let compile_result = match compile_result {
             Ok(r) => r,
             Err(e) => {
@@ -6808,8 +6813,8 @@ impl<M: Clone> MetaInterp<M> {
                     green_key,
                     num_ops_before,
                     num_ops_after,
-                    std::time::Duration::ZERO,
-                    std::time::Duration::ZERO,
+                    opt_time,
+                    compile_time,
                 );
                 // warmstate.py:339-348 attach the same compiled token object.
                 self.attach_procedure_with_redirect(green_key, Arc::clone(&token));
@@ -7614,6 +7619,7 @@ impl<M: Clone> MetaInterp<M> {
         // the pre-peel arg set for slots the loop label rebinds.
         unroll_opt.emit_start_label = false;
 
+        let optimize_start = Instant::now();
         let optimize_result = unroll_opt.optimize_trace_with_constants_and_inputs_vable(
             &trace_ops,
             &mut constants,
@@ -7634,6 +7640,7 @@ impl<M: Clone> MetaInterp<M> {
                 return false;
             }
         };
+        let opt_time = Instant::now().saturating_duration_since(optimize_start);
         // compile.py:384-390: merge loop_info deps first, then deps carried
         // from the exported start_state.
         let mut quasi_immutable_deps = std::mem::take(&mut unroll_opt.quasi_immutable_deps);
@@ -7809,6 +7816,7 @@ impl<M: Clone> MetaInterp<M> {
         // compile.py:532-546 `debug_start("jit-backend") +
         // profiler.start_backend() ... try: do_compile_loop ... finally:
         // ... profiler.end_backend() + debug_stop("jit-backend")`.
+        let compile_start = Instant::now();
         let compile_result = {
             let _backend_scope = self.staticdata.profiler.enter_backend();
             std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
@@ -7820,6 +7828,7 @@ impl<M: Clone> MetaInterp<M> {
                 )
             }))
         };
+        let compile_time = Instant::now().saturating_duration_since(compile_start);
         let compile_result = match compile_result {
             Ok(r) => r,
             Err(payload) => {
@@ -7974,8 +7983,8 @@ impl<M: Clone> MetaInterp<M> {
                     green_key,
                     num_ops_before,
                     num_combined_ops,
-                    std::time::Duration::ZERO,
-                    std::time::Duration::ZERO,
+                    opt_time,
+                    compile_time,
                 );
                 self.attach_procedure_with_redirect(green_key, Arc::clone(&token));
                 self.stats.loops_compiled += 1;
@@ -8569,6 +8578,7 @@ impl<M: Clone> MetaInterp<M> {
 
         // InvalidLoop during optimization should abort the trace, not crash
         // the process. Matches compile_loop.
+        let optimize_start = Instant::now();
         let optimize_result = optimizer.optimize_with_constants_and_inputs_oprc(
             // `trace.ops` are the canonical `Rc<Op>`, so `input_ops`
             // seeds identity directly from them.
@@ -8599,6 +8609,7 @@ impl<M: Clone> MetaInterp<M> {
                 return Err(SwitchToBlackhole::giveup());
             }
         };
+        let opt_time = Instant::now().saturating_duration_since(optimize_start);
         // RPython optimizer.py:552-556 (flush=True): Finish/Jump is sent
         // through passes inside propagate_all_forward and ends up in
         // new_operations naturally — no restoration needed.
@@ -8731,11 +8742,13 @@ impl<M: Clone> MetaInterp<M> {
         // compile.py:532-546 `debug_start("jit-backend") +
         // profiler.start_backend() ... try: do_compile_loop ... finally:
         // ... profiler.end_backend() + debug_stop("jit-backend")`.
+        let compile_start = Instant::now();
         let compile_loop_result = {
             let _backend_guard = self.staticdata.profiler.enter_backend();
             self.backend
                 .compile_loop(&inputargs, &optimized_ops, &token)
         };
+        let compile_time = Instant::now().saturating_duration_since(compile_start);
         match compile_loop_result {
             Ok(_) => {
                 self.last_compiled_artifact_invalidation_flag = Some(token.invalidation_flag());
@@ -8857,6 +8870,13 @@ impl<M: Clone> MetaInterp<M> {
                         },
                     );
                 }
+                self.warm_state.log_compile(
+                    green_key,
+                    num_ops_before,
+                    num_ops_after,
+                    opt_time,
+                    compile_time,
+                );
                 self.attach_procedure_with_redirect(green_key, Arc::clone(&token));
                 self.stats.loops_compiled += 1;
                 // `cpu.tracker.total_compiled_loops` is bumped inside
@@ -8990,6 +9010,7 @@ impl<M: Clone> MetaInterp<M> {
         optimizer.snapshot_vref_boxes = snapshot_vref_map;
         optimizer.snapshot_frame_pcs = snapshot_pc_map;
 
+        let optimize_start = Instant::now();
         let optimize_result = optimizer.optimize_with_constants_and_inputs_oprc(
             // Canonical `Rc<Op>`; `input_ops` seeds identity from them.
             &trace.ops,
@@ -9013,6 +9034,7 @@ impl<M: Clone> MetaInterp<M> {
                 return None;
             }
         };
+        let opt_time = Instant::now().saturating_duration_since(optimize_start);
 
         // optimizer.py:557 self.resumedata_memo.update_counters(profiler)
         optimizer.update_counters(&self.staticdata.profiler);
@@ -9106,6 +9128,7 @@ impl<M: Clone> MetaInterp<M> {
         // compile.py:532-546 `debug_start("jit-backend") +
         // profiler.start_backend() ... try: do_compile_loop ... finally:
         // ... profiler.end_backend() + debug_stop("jit-backend")`.
+        let compile_start = Instant::now();
         let compile_loop_result = {
             let _backend_guard = self.staticdata.profiler.enter_backend();
             self.backend.compile_loop(
@@ -9115,6 +9138,7 @@ impl<M: Clone> MetaInterp<M> {
                     .expect("JitCellToken must stay uniquely owned until backend compile"),
             )
         };
+        let compile_time = Instant::now().saturating_duration_since(compile_start);
         match compile_loop_result {
             Ok(_) => {
                 self.assign_guard_hashes(token.as_ref());
@@ -9198,6 +9222,13 @@ impl<M: Clone> MetaInterp<M> {
                         next_global_opref,
                     },
                 );
+                self.warm_state.log_compile(
+                    green_key,
+                    num_ops_before,
+                    num_ops_after,
+                    opt_time,
+                    compile_time,
+                );
                 self.stats.loops_compiled += 1;
                 // `cpu.tracker.total_compiled_loops` is bumped inside
                 // `CompiledLoopToken::new` (model.py:297 parity).
@@ -11381,6 +11412,7 @@ impl<M: Clone> MetaInterp<M> {
         // constant pool merge. Const objects flow via rd_consts + fresh
         // decode (resume.py:1245-1282).
         let retrace_limit = self.warm_state.retrace_limit();
+        let optimize_start = Instant::now();
         let bridge_optimize_result = {
             let compiled = self.compiled_loops.get_mut(&green_key).unwrap();
             optimizer.optimize_bridge(
@@ -11419,6 +11451,7 @@ impl<M: Clone> MetaInterp<M> {
                 return false;
             }
         };
+        let opt_time = Instant::now().saturating_duration_since(optimize_start);
         // optimizer.py:557 self.resumedata_memo.update_counters(profiler)
         optimizer.update_counters(&self.staticdata.profiler);
         // RPython-orthodox: unroll.py replay uses Const args directly;
@@ -11504,13 +11537,15 @@ impl<M: Clone> MetaInterp<M> {
         // compile.py:532-546 `debug_start("jit-backend") +
         // profiler.start_backend() ... try: do_compile_loop ... finally:
         // ... profiler.end_backend() + debug_stop("jit-backend")`.
+        let compile_start = Instant::now();
         let compile_result = {
             let _backend_scope = self.staticdata.profiler.enter_backend();
             std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                 self.backend
                     .compile_loop(bridge_inputargs, &optimized_ops, &token)
             }))
         };
+        let compile_time = Instant::now().saturating_duration_since(compile_start);
         let compile_result = match compile_result {
             Ok(r) => r,
             Err(payload) => {
@@ -11636,6 +11671,13 @@ impl<M: Clone> MetaInterp<M> {
                         next_global_opref,
                     },
                 );
+                self.warm_state.log_compile(
+                    original_green_key,
+                    bridge_ops.len(),
+                    num_optimized_ops,
+                    opt_time,
+                    compile_time,
+                );
                 self.attach_procedure_with_redirect(original_green_key, Arc::clone(&token));
                 self.stats.loops_compiled += 1;
                 // `cpu.tracker.total_compiled_loops` is bumped inside