Tuple Support and WIP Object Support in JIT by LuigiTheJokesterAlvarez · Pull Request #8188 · RustPython/RustPython · GitHub
Skip to content

Tuple Support and WIP Object Support in JIT#8188

Open
LuigiTheJokesterAlvarez wants to merge 12 commits into
RustPython:mainfrom
LuigiTheJokesterAlvarez:abi-fix
Open

Tuple Support and WIP Object Support in JIT#8188
LuigiTheJokesterAlvarez wants to merge 12 commits into
RustPython:mainfrom
LuigiTheJokesterAlvarez:abi-fix

Conversation

@LuigiTheJokesterAlvarez

@LuigiTheJokesterAlvarez LuigiTheJokesterAlvarez commented Jun 28, 2026

Copy link
Copy Markdown

Summary

This changes the JIT to support tuples and possibly, objects

Summary by CodeRabbit

  • New Features

    • Added heap-backed, shape-aware tuple support across the JIT ABI and runtime.
    • Extended the JIT ABI/type system with a new object/pointer-like value type and tuple constant lowering.
    • JIT type annotations now accept tuple, mapping it to the new representation.
  • Bug Fixes

    • Fixed tuple unpacking to use the heap-backed layout with element-count validation.
    • Corrected local load/store handling for scalar vs object-backed values.
    • Tightened boolean conversion and improved return value handling to match the pre-declared return type.
  • Tests

    • Added tuple_identity and expanded JIT snippet coverage with explicit tuple-related type annotations.

@coderabbitai

coderabbitai Bot commented Jun 28, 2026

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

Caution

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

⚠️ Outside diff range comments (1)
crates/jit/src/instructions.rs (1)

55-60: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Preserve tuple shape for object arguments before supporting tuple unpacking there.

JitType::Object arguments are materialized as ObjectKind::Opaque, but UnpackSequence only accepts ObjectKind::Tuple. A JIT function taking a tuple argument and unpacking it will reject its own object ABI value despite the new tuple support.

Also applies to: 878-884

🤖 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 `@crates/jit/src/instructions.rs` around lines 55 - 60, In from_type_and_value,
JitType::Object is currently always wrapped as ObjectKind::Opaque, which breaks
tuple arguments that later flow into UnpackSequence. Update the object-argument
materialization path in from_type_and_value so tuple-shaped values are preserved
as ObjectKind::Tuple when the ABI value represents a tuple, and keep
UnpackSequence aligned to accept that preserved shape instead of rejecting the
function’s own argument value.
🤖 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 `@crates/jit/src/instructions.rs`:
- Around line 155-161: The changed Rust code needs rustfmt cleanup, with several
signatures and calls in functions like local_to_value and other touched Rust
sections needing standard formatting and trailing whitespace removed. Run cargo
fmt on the affected code so the updated blocks follow the default Rust style
consistently across the modified instructions and nearby changed lines.
- Around line 429-435: The return ABI handling in build_function/instructions.rs
is being mutated too late and can diverge across branches, so make the return
signature fixed before any return emission and enforce that all subsequent
returns match the first inferred ret_ty. Update the logic around self.sig.ret
and builder.func.signature.returns so the signature is established once, then
validate or reject mixed return types in the return-building path rather than
appending/changing ABI params later.

In `@crates/jit/src/lib.rs`:
- Around line 41-50: Add a matching deallocation path for the heap tuples
allocated by jit_alloc_tuple, since BuildTuple and tuple constants currently
return raw allocations with no owner or free. Introduce a clear ownership
strategy in crates/jit/src/lib.rs, and ensure the JIT runtime exposes a
corresponding free/release function or otherwise transfers ownership so tuple
allocations are reclaimed after use.

---

Outside diff comments:
In `@crates/jit/src/instructions.rs`:
- Around line 55-60: In from_type_and_value, JitType::Object is currently always
wrapped as ObjectKind::Opaque, which breaks tuple arguments that later flow into
UnpackSequence. Update the object-argument materialization path in
from_type_and_value so tuple-shaped values are preserved as ObjectKind::Tuple
when the ABI value represents a tuple, and keep UnpackSequence aligned to accept
that preserved shape instead of rejecting the function’s own argument value.
🪄 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: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: e9eeb398-7e6d-4522-bc07-630d242533be

📥 Commits

Reviewing files that changed from the base of the PR and between 0009dd6 and bdf64b7.

📒 Files selected for processing (3)
  • crates/jit/src/instructions.rs
  • crates/jit/src/lib.rs
  • crates/vm/src/builtins/function.rs

Comment thread crates/jit/src/instructions.rs Outdated
Comment thread crates/jit/src/instructions.rs Outdated
Comment thread crates/jit/src/lib.rs
Comment on lines +41 to +50
extern "C" fn jit_alloc_tuple(len: i64) -> i64 {
let layout = Layout::array::<i64>(len as usize + 1).unwrap();
let ptr = unsafe { alloc(layout) };
if ptr.is_null() {
handle_alloc_error(layout);
}
unsafe {
*(ptr as *mut i64) = len;
}
ptr as i64

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

Add an ownership/freeing strategy for heap tuples.

jit_alloc_tuple allocates with the global allocator and returns a raw pointer, but this PR adds no matching deallocation path. Since BuildTuple and tuple constants call this allocator at runtime, repeated JIT invocations leak every tuple allocation.

🤖 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 `@crates/jit/src/lib.rs` around lines 41 - 50, Add a matching deallocation path
for the heap tuples allocated by jit_alloc_tuple, since BuildTuple and tuple
constants currently return raw allocations with no owner or free. Introduce a
clear ownership strategy in crates/jit/src/lib.rs, and ensure the JIT runtime
exposes a corresponding free/release function or otherwise transfers ownership
so tuple allocations are reclaimed after use.

@ShaharNaveh ShaharNaveh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi, thanks for contributing to RustPython:)

Can you please check the AI policy? (#8188 (comment))

Comment thread crates/jit/src/instructions.rs Outdated
@@ -1,6 +1,8 @@
// instructions.rs

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the reasoning for this comment?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

when i forget what im doing i comment it out to make myself know again what im doing

yeah it's weird

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ill delete unnecessary comments and if i can right now, document the functions i did, and making the changes coderabbit requested

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i did documentation yesterday lol

Comment thread crates/jit/src/lib.rs Outdated
@@ -1,7 +1,9 @@
// lib.rs

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 (1)
crates/jit/src/instructions.rs (1)

382-406: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

Heap tuple allocations need a reclamation path crates/jit/src/instructions.rs:382-406

build_heap_tuple allocates raw memory via jit_alloc_tuple, but nothing in the JIT frees those buffers. CompiledCode::drop only releases the compiled module, not tuple storage. Since both LoadConst tuples and BuildTuple hit this path, repeated execution will leak one allocation per tuple unless ownership is transferred to the VM or a matching free path is added.

🤖 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 `@crates/jit/src/instructions.rs` around lines 382 - 406, The heap tuple
allocation in build_heap_tuple currently has no matching reclamation path, so
tuple buffers allocated through alloc_func will leak across repeated execution.
Update the JIT tuple lifecycle around build_heap_tuple, TupleShape, and
CompiledCode::drop so ownership of these allocations is clearly transferred to
the VM or a corresponding free is registered and invoked when the compiled
result is released. Make sure both LoadConst tuples and BuildTuple allocations
are covered by the same cleanup path.
🧹 Nitpick comments (2)
crates/vm/src/builtins/function/jit.rs (1)

55-63: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Error message is now stale — tuple is accepted but not mentioned.

The error message still says "Jit requires argument to be either int, float or bool"-style text, but the branch above now also accepts tuples. Update the message so it doesn't mislead users about supported argument types.

📝 Proposed fix
             Err(new_jit_error(
-                "Jit requires argument to be either int, float or bool".to_owned(),
+                "Jit requires argument to be either int, float, bool or tuple".to_owned(),
                 vm,
             ))
🤖 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 `@crates/vm/src/builtins/function/jit.rs` around lines 55 - 63, The fallback
error in JIT argument validation is stale because the `jit` type check now also
accepts tuples via `value.is(vm.ctx.types.tuple_type)`. Update the
`new_jit_error` message in `jit.rs` to include tuple among the supported
argument types, keeping the message aligned with the branches in the JIT
type-detection logic.
crates/jit/src/instructions.rs (1)

616-622: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Deduplicate the tuple JitValue::Object construction. The JitValue::Object(ptr, Rc::new(ObjectKind::Tuple(shape))) wrapping is identical here and in prepare_const (Lines 374-376). Consider having build_heap_tuple return the finished JitValue (or a small helper) so the wrapping lives in one place. As per coding guidelines: "When branches differ only in a value but share common logic, extract the differing value first, then call the common logic once to avoid duplicate code."

🤖 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 `@crates/jit/src/instructions.rs` around lines 616 - 622, The tuple
JitValue::Object construction is duplicated in Instruction::BuildTuple and
prepare_const, so move the wrapping into one shared place. Update
build_heap_tuple to return the finished JitValue, or add a small helper used by
both build_heap_tuple and prepare_const, and have Instruction::BuildTuple and
prepare_const call that shared logic instead of rebuilding JitValue::Object(ptr,
Rc::new(ObjectKind::Tuple(shape))) in multiple places.

Source: Coding guidelines

🤖 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 `@crates/vm/src/builtins/function/jit.rs`:
- Around line 57-58: The conditional in the JIT function implementation is
formatted incorrectly, with the closing brace and else split across separate
lines. Update the `jit.rs` block around the `Function`/`jit` logic so it follows
rustfmt style and uses a single-line `} else {` form consistent with `cargo
fmt`.

---

Outside diff comments:
In `@crates/jit/src/instructions.rs`:
- Around line 382-406: The heap tuple allocation in build_heap_tuple currently
has no matching reclamation path, so tuple buffers allocated through alloc_func
will leak across repeated execution. Update the JIT tuple lifecycle around
build_heap_tuple, TupleShape, and CompiledCode::drop so ownership of these
allocations is clearly transferred to the VM or a corresponding free is
registered and invoked when the compiled result is released. Make sure both
LoadConst tuples and BuildTuple allocations are covered by the same cleanup
path.

---

Nitpick comments:
In `@crates/jit/src/instructions.rs`:
- Around line 616-622: The tuple JitValue::Object construction is duplicated in
Instruction::BuildTuple and prepare_const, so move the wrapping into one shared
place. Update build_heap_tuple to return the finished JitValue, or add a small
helper used by both build_heap_tuple and prepare_const, and have
Instruction::BuildTuple and prepare_const call that shared logic instead of
rebuilding JitValue::Object(ptr, Rc::new(ObjectKind::Tuple(shape))) in multiple
places.

In `@crates/vm/src/builtins/function/jit.rs`:
- Around line 55-63: The fallback error in JIT argument validation is stale
because the `jit` type check now also accepts tuples via
`value.is(vm.ctx.types.tuple_type)`. Update the `new_jit_error` message in
`jit.rs` to include tuple among the supported argument types, keeping the
message aligned with the branches in the JIT type-detection logic.
🪄 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: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: e4d2b7ab-8b30-431a-b98b-975a72432d63

📥 Commits

Reviewing files that changed from the base of the PR and between f32149d and 661b659.

📒 Files selected for processing (3)
  • crates/jit/src/instructions.rs
  • crates/jit/src/lib.rs
  • crates/vm/src/builtins/function/jit.rs
🚧 Files skipped from review as they are similar to previous changes (1)
  • crates/jit/src/lib.rs

Comment thread crates/vm/src/builtins/function/jit.rs Outdated
@LuigiTheJokesterAlvarez

Copy link
Copy Markdown
Author

I will fix the leak with tuple alloc later

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
extra_tests/snippets/jit.py (1)

24-24: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Variable name violates PEP 8 (snake_case).

tupleId should be tuple_id per PEP 8 naming conventions for local variables.

As per coding guidelines, **/*.py: "Follow PEP 8 style for custom Python code".

🎨 Proposed fix
-    tupleId = (1, 2)
-    assert tuple_identity(tupleId) == tupleId
+    tuple_id = (1, 2)
+    assert tuple_identity(tuple_id) == tuple_id
🤖 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 `@extra_tests/snippets/jit.py` at line 24, Rename the local variable tupleId to
tuple_id in jit.py to comply with PEP 8 snake_case naming; update any nearby
references in the same snippet or function so the new identifier is used
consistently.

Source: Coding guidelines

🤖 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.

Nitpick comments:
In `@extra_tests/snippets/jit.py`:
- Line 24: Rename the local variable tupleId to tuple_id in jit.py to comply
with PEP 8 snake_case naming; update any nearby references in the same snippet
or function so the new identifier is used consistently.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: 01ba0ae0-098c-406f-9653-2d60610b3f24

📥 Commits

Reviewing files that changed from the base of the PR and between 661b659 and c12d905.

📒 Files selected for processing (2)
  • crates/vm/src/builtins/function/jit.rs
  • extra_tests/snippets/jit.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • crates/vm/src/builtins/function/jit.rs

@LuigiTheJokesterAlvarez

Copy link
Copy Markdown
Author

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants