Audit RPython parity in OptVirtualize and fix deviations

[lifthrasiir]

This pull request removes the unused imported_virtual_heads field from the OptContext struct and updates related logic to align with RPython's handling of virtual heads during optimization. It also clarifies how virtualizables are treated in the optimizer, ensuring consistency with RPython's semantics.

Key changes:

Refactoring and Cleanup:

• Removed the imported_virtual_heads field from OptContext and all related initialization code, as imported virtual heads are now managed differently and no longer require a side table.
• Updated the logic in Optimizer::import_state to avoid populating imported_virtual_heads, with explanatory comments about how replaying short preambles populates the necessary cache instead.

Behavioral Clarification:

• Added documentation and an early return in Optimizer::force_box to clarify that "virtualizable" objects should not be treated as true virtuals, preventing destructive state changes and maintaining RPython parity.

Summary by CodeRabbit

• Optimization & Performance

  • Enhanced virtual object allocation tracking and optimization pipeline
  • Improved cache handling for field access patterns

• Bug Fixes

  • Fixed virtual allocation forcing logic to prevent state tracking corruption

• Chores

  • Refactored virtualizable frame initialization and tracking for improved maintainability
  • Updated internal virtual reference modeling and field tracking

[coderabbitai]

Warning

Rate limit exceeded

@lifthrasiir has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 15 minutes and 47 seconds before requesting another review.

To continue reviewing without waiting, purchase usage credits in the billing tab.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 20095468-492b-460d-ab23-54ad28fe4a4e

📥 Commits

Reviewing files that changed from the base of the PR and between 1f2e752 and 5cfc206.

📒 Files selected for processing (2)

• majit/majit-metainterp/src/optimizeopt/rewrite.rs
• majit/majit-metainterp/src/optimizeopt/virtualize.rs

Walkthrough

This PR refactors virtualizable frame tracking in the optimization pass. It introduces VirtualizableTracker to manage state initialization, restructures JitVirtualRef from three tracked fields to a two-field model (virtual_token, forced), and reworks array/field optimizations to directly consume and produce virtual state instead of relying on side-table lookups. The OptContext constructors are updated to reflect struct layout changes.

Changes

Virtualizable Frame Tracking Refactor

Layer / File(s)	Summary
Virtual Reference Field Layout majit/majit-metainterp/src/optimizeopt/virtualize.rs	JitVirtualRef tracked field indices change from three fields (including type tag) to two fields: virtual_token (index 0) and forced (index 1); vtable identity now comes from descriptor.
VirtualizableTracker Infrastructure majit/majit-metainterp/src/optimizeopt/virtualize.rs	Introduces VirtualizableTracker struct that lazily seeds PtrInfo::Virtualizable during optimization by mapping configured byte offsets to virtual field indices and input typed slots; adds helper methods to determine standard virtualizable ref, map array indices, and mirror Setarrayitem* writes into tracked state.
OptVirtualize Struct Refactor majit/majit-metainterp/src/optimizeopt/virtualize.rs	Consolidates virtualizable tracking into vable: Option<VirtualizableTracker>; removes public is_phase2 field and prior vable_config/vable_initialized/needs_vable_setup fields; makes all fields private.
Virtualizable Setup Integration majit/majit-metainterp/src/optimizeopt/virtualize.rs	propagate_forward entry setup calls VirtualizableTracker::ensure_setup() to initialize state on demand.
Array and Field Operation Optimizations majit/majit-metainterp/src/optimizeopt/virtualize.rs	optimize_setarrayitem_gc removes ops that update existing PtrInfo::VirtualArray slots and mirrors writes into PtrInfo::Virtualizable via tracker; optimize_getarrayitem_gc returns InvalidLoop for out-of-range/uninitialized indices; optimize_getfield_gc, optimize_getinteriorfield_gc, and optimize_setinteriorfield_gc resolve and forward PtrInfo::VirtualArray/VirtualArrayStruct values; ArraylenGc is registered as pure operation.
Virtual Reference Creation and Finalization majit/majit-metainterp/src/optimizeopt/virtualize.rs	optimize_virtual_ref now creates struct with two-field layout using VRefSizeDescr for vtable; optimize_virtual_ref_finish updates forced and clears virtual_token to null on normal finish; emits SetfieldGc for escaped values.
Optimizer Integration and Context Updates majit/majit-metainterp/src/optimizeopt/optimizer.rs, majit/majit-metainterp/src/optimizeopt/mod.rs	force_box short-circuits for virtualizable_via_box to preserve tracked state; install_imported_virtuals no longer pushes to side table, relying on short preamble replay; OptContext::new and OptContext::with_num_inputs_and_start_pos initializers updated to reflect struct layout changes.

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

🐰 A rabbit hops through virtual frames,
Tracking tokens, forced and free,
Arrays mirror, fields align,
Two-field structs dance gracefully,
Optimization flows like spring! 🌱

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately reflects the main change: a comprehensive audit of RPython parity in OptVirtualize and correction of identified deviations throughout the optimization pipeline.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (3)

majit/majit-metainterp/src/optimizeopt/virtualize.rs (3)

169-201: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick win

Length-mismatch between parallel array config vectors silently truncates.

init zips array_field_offsets with array_lengths, so any mismatch in their lengths silently drops entries; array_field_descrs.get(array_idx) then falls back to a synthetic make_field_index_descr, which loses the parent_descr backreference that static_field_descrs/the comment at lines 30–47 emphasizes is required for the force path. Consider asserting array_field_offsets.len() == array_lengths.len() (and ideally == array_item_types.len() == array_field_descrs.len()) at config construction or at the top of init to catch misconfigured VirtualizableConfig builders before the optimizer silently swaps in placeholder descrs.

🛡️ Proposed assertion

     fn init(&mut self, ctx: &mut OptContext) {
         self.initialized = true;
         if ctx.num_inputs() <= 1 {
             return;
         }
+        debug_assert_eq!(
+            self.config.array_field_offsets.len(),
+            self.config.array_lengths.len(),
+            "VirtualizableConfig: array_field_offsets/array_lengths must be parallel",
+        );
+        debug_assert!(
+            self.config.array_field_descrs.is_empty()
+                || self.config.array_field_descrs.len()
+                    == self.config.array_field_offsets.len(),
+            "VirtualizableConfig: array_field_descrs must be empty or parallel to array_field_offsets",
+        );

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@majit/majit-metainterp/src/optimizeopt/virtualize.rs` around lines 169 - 201,
The loop in init that zips self.config.array_field_offsets with
self.config.array_lengths can silently drop entries and cause
array_field_descrs.get(array_idx) to fall back to make_field_index_descr (losing
required parent_descr backrefs); fix by adding an assertion (or explicit
length-check and early error) that self.config.array_field_offsets.len() ==
self.config.array_lengths.len() (and ideally also ==
self.config.array_item_types.len() and == self.config.array_field_descrs.len())
either when constructing VirtualizableConfig or at the top of init in
virtualize.rs so misconfigured VirtualizableConfig builders fail fast instead of
generating placeholder descrs.

---

664-694: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Negative constant index falls through to mirror_setarrayitem and over-allocates.

When index is a constant negative integer, idx = index as usize becomes ~usize::MAX. The idx < vinfo.items.len() check on the virtual-array branch correctly fails, but execution then drops into the mirror_setarrayitem(array_ref, index, ...) call with the negative i64. Inside, elem_idx = index as usize is huge, and set_array_element performs elems.resize(elem_idx + 1, OpRef::NONE) / vec![OpRef::NONE; elem_idx + 1] (lines 1756 and 1760), which either panics on elem_idx + 1 overflow in debug, or attempts to allocate ~usize::MAX elements in release.

Note that the symmetric optimize_getarrayitem_gc does guard this case (line 707, returns InvalidLoop); the setter path should at minimum reject negatives before mirroring.

🛡️ Proposed fix to skip negative indices on the mirror path

         if let Some(index) = ctx.get_constant_int(index_ref) {
             let idx = index as usize;
             let did_virtual_write = ctx
                 .with_ptr_info_mut(array_ref, |info| {
                     if let PtrInfo::VirtualArray(vinfo) = info {
                         if idx < vinfo.items.len() {
                             vinfo.items[idx] = value_ref;
                             return true;
                         }
                     }
                     false
                 })
                 .unwrap_or(false);
             if did_virtual_write {
                 return OptimizationResult::Remove;
             }
-            if let Some(ref vt) = self.vable {
+            if index >= 0 && let Some(ref vt) = self.vable {
                 vt.mirror_setarrayitem(array_ref, index, value_ref, ctx);
             }
         }

(Adjust split-if form if MSRV does not include the let-chain stabilization; equivalent guard via nested if.)

---

1042-1063: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Type mismatch breaks optimize_jit_force_virtual — token field written as Ref but read as Int.

VRefFieldDescr::field_type declares VREF_VIRTUAL_TOKEN_FIELD_INDEX as Type::Int (line 1652), and optimize_jit_force_virtual reads it via ctx.get_constant_int(token_ref).is_some_and(|v| v == 0) (line 1121). However, both finish paths write using ctx.emit_constant_ref(majit_ir::GcRef(0)) (lines 1042 and 1061), which creates a Value::Ref-typed constant. Since get_constant_int() only returns Some for Value::Int and returns None for any other type, the read will always fail, and the optimization silently never fires on the normal vref-finish path — the case it is designed to handle.

The comment at lines 1119–1120 explicitly claims the value uses emit_constant_int(0), which contradicts the actual code.

Change the writers to use ctx.emit_constant_int(0) and update the comment, or change the reader to extract the Ref-typed value and adjust the field type accordingly.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@majit/majit-metainterp/src/optimizeopt/virtualize.rs` around lines 1042 -
1063, The token field VREF_VIRTUAL_TOKEN_FIELD_INDEX is declared as Type::Int
but the code writes a Ref constant (ctx.emit_constant_ref(majit_ir::GcRef(0)))
so get_constant_int in optimize_jit_force_virtual can never see a 0; change the
writers to emit an integer constant using ctx.emit_constant_int(0) for both
places where null_ref is created/set (the earlier early-return path and the
set_token creation), and update the adjacent comment that currently claims
emit_constant_int(0) to reflect the corrected writer; ensure references to
VREF_VIRTUAL_TOKEN_FIELD_INDEX, optimize_jit_force_virtual,
ctx.emit_constant_ref and ctx.emit_constant_int are adjusted accordingly.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@majit/majit-metainterp/src/optimizeopt/virtualize.rs`:
- Around line 91-134: The struct field `initialized` on VirtualizableTracker is
unused; remove the `initialized` field and all assignments to it (in
VirtualizableTracker::new, VirtualizableTracker::setup, and
VirtualizableTracker::init) and keep behavior of ensure_setup as-is, or
alternatively if idempotency was intended, change ensure_setup to call init only
when `!self.initialized` and set `self.initialized = true` inside `init`; pick
one approach and update VirtualizableTracker, its new(), setup(), init(), and
ensure_setup() accordingly to eliminate the dead state.

---

Outside diff comments:
In `@majit/majit-metainterp/src/optimizeopt/virtualize.rs`:
- Around line 169-201: The loop in init that zips
self.config.array_field_offsets with self.config.array_lengths can silently drop
entries and cause array_field_descrs.get(array_idx) to fall back to
make_field_index_descr (losing required parent_descr backrefs); fix by adding an
assertion (or explicit length-check and early error) that
self.config.array_field_offsets.len() == self.config.array_lengths.len() (and
ideally also == self.config.array_item_types.len() and ==
self.config.array_field_descrs.len()) either when constructing
VirtualizableConfig or at the top of init in virtualize.rs so misconfigured
VirtualizableConfig builders fail fast instead of generating placeholder descrs.
- Around line 1042-1063: The token field VREF_VIRTUAL_TOKEN_FIELD_INDEX is
declared as Type::Int but the code writes a Ref constant
(ctx.emit_constant_ref(majit_ir::GcRef(0))) so get_constant_int in
optimize_jit_force_virtual can never see a 0; change the writers to emit an
integer constant using ctx.emit_constant_int(0) for both places where null_ref
is created/set (the earlier early-return path and the set_token creation), and
update the adjacent comment that currently claims emit_constant_int(0) to
reflect the corrected writer; ensure references to
VREF_VIRTUAL_TOKEN_FIELD_INDEX, optimize_jit_force_virtual,
ctx.emit_constant_ref and ctx.emit_constant_int are adjusted accordingly.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: b5a35005-1206-42e7-bc82-cf07b1c9d740

📥 Commits

Reviewing files that changed from the base of the PR and between 64aa654 and 1f2e752.

📒 Files selected for processing (3)

• majit/majit-metainterp/src/optimizeopt/mod.rs
• majit/majit-metainterp/src/optimizeopt/optimizer.rs
• majit/majit-metainterp/src/optimizeopt/virtualize.rs

[youknowone]

generated review:

 - RECORD_EXACT_VALUE_I/R optimization was removed.
    Current majit/majit-metainterp/src/optimizeopt/virtualize.rs:1430 says
    RECORD_EXACT_CLASS / RECORD_EXACT_VALUE_* are handled by OptRewrite, but
    actual majit/majit-metainterp/src/optimizeopt/rewrite.rs:2878 only handles
    RecordExactClass; it has no RecordExactValueI/R handling. RPython handles
    this in rpython/jit/metainterp/optimizeopt/rewrite.py:388 via
    optimize_record_exact_value, making the box constant to the expected const.
    main still had behavior for RecordExactValueI/R, even though it lived in
    OptVirtualize rather than the RPython-matching pass, so this patch is a
    parity regression relative to main.

[lifthrasiir]

generated review:

 - RECORD_EXACT_VALUE_I/R optimization was removed.
    Current majit/majit-metainterp/src/optimizeopt/virtualize.rs:1430 says
    RECORD_EXACT_CLASS / RECORD_EXACT_VALUE_* are handled by OptRewrite, but
    actual majit/majit-metainterp/src/optimizeopt/rewrite.rs:2878 only handles
    RecordExactClass; it has no RecordExactValueI/R handling. RPython handles
    this in rpython/jit/metainterp/optimizeopt/rewrite.py:388 via
    optimize_record_exact_value, making the box constant to the expected const.
    main still had behavior for RecordExactValueI/R, even though it lived in
    OptVirtualize rather than the RPython-matching pass, so this patch is a
    parity regression relative to main.

Fixed in 5cfc206.