Skip to content

Newtype var_num oparg #7400

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 2 commits into from
Mar 12, 2026
+124 −77
Merged

Newtype var_num oparg #7400

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6f81e21
Newtype oparg::VarNum
ShaharNaveh Mar 9, 2026
87eeba9
fix jit
ShaharNaveh Mar 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 17 additions & 7 deletions crates/codegen/src/compile.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -715,14 +715,14 @@ impl Compiler {
}

/// Get the index of a local variable.
fn get_local_var_index(&mut self, name: &str) -> CompileResult<u32> {
fn get_local_var_index(&mut self, name: &str) -> CompileResult<oparg::VarNum> {
let info = self.code_stack.last_mut().unwrap();
let idx = info
.metadata
.varnames
.get_index_of(name)
.unwrap_or_else(|| info.metadata.varnames.insert_full(name.to_owned()).0);
Ok(idx.to_u32())
Ok(idx.to_u32().into())
}

/// Get the index of a global name.
Expand Down Expand Up @@ -1283,7 +1283,12 @@ impl Compiler {
/// if format > VALUE_WITH_FAKE_GLOBALS (2): raise NotImplementedError
fn emit_format_validation(&mut self) -> CompileResult<()> {
// Load format parameter (first local variable, index 0)
emit!(self, Instruction::LoadFast { var_num: 0 });
emit!(
self,
Instruction::LoadFast {
var_num: oparg::VarNum::from_u32(0)
}
);

// Load VALUE_WITH_FAKE_GLOBALS constant (2)
self.emit_load_const(ConstantData::Integer { value: 2.into() });
Expand Down Expand Up @@ -1562,15 +1567,19 @@ impl Compiler {
fn name(&mut self, name: &str) -> bytecode::NameIdx {
self._name_inner(name, |i| &mut i.metadata.names)
}
fn varname(&mut self, name: &str) -> CompileResult<bytecode::NameIdx> {

fn varname(&mut self, name: &str) -> CompileResult<oparg::VarNum> {
// Note: __debug__ checks are now handled in symboltable phase
Ok(self._name_inner(name, |i| &mut i.metadata.varnames))
Ok(oparg::VarNum::from_u32(
self._name_inner(name, |i| &mut i.metadata.varnames),
))
}

fn _name_inner(
&mut self,
name: &str,
cache: impl FnOnce(&mut ir::CodeInfo) -> &mut IndexSet<String>,
) -> bytecode::NameIdx {
) -> u32 {
let name = self.mangle(name);
let cache = cache(self.current_code_info());
cache
Expand Down Expand Up @@ -4129,7 +4138,8 @@ impl Compiler {

// Load defaults/kwdefaults with LOAD_FAST
for i in 0..num_typeparam_args {
emit!(self, Instruction::LoadFast { var_num: i as u32 });
let var_num = oparg::VarNum::from(i as u32);
emit!(self, Instruction::LoadFast { var_num });
}
}

Expand Down
4 changes: 2 additions & 2 deletions crates/codegen/src/ir.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1109,8 +1109,8 @@ impl InstrDisplayContext for CodeInfo {
self.metadata.names[i].as_ref()
}

fn get_varname(&self, i: usize) -> &str {
self.metadata.varnames[i].as_ref()
fn get_varname(&self, var_num: oparg::VarNum) -> &str {
self.metadata.varnames[var_num.as_usize()].as_ref()
Comment on lines +1112 to +1113
Copy link
Contributor

@coderabbitai coderabbitai bot Mar 11, 2026

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major

Resolve VarNum against the full locals-plus table.

Line 1113 still assumes every VarNum lives in metadata.varnames, but Line 1797 creates LOAD_FAST slots as varnames_len + ... when lowering LOAD_CLOSURE. That means displaying valid closure bytecode can walk past varnames and panic. Please fall back into cellvars/freevars here instead of indexing only varnames.

🛠️ Suggested fix 
 fn get_varname(&self, var_num: oparg::VarNum) -> &str {
-    self.metadata.varnames[var_num.as_usize()].as_ref()
+    let idx = var_num.as_usize();
+    if let Some(name) = self.metadata.varnames.get_index(idx) {
+        name.as_ref()
+    } else {
+        let cell_idx = idx - self.metadata.varnames.len();
+        self.metadata
+            .cellvars
+            .get_index(cell_idx)
+            .or_else(|| self.metadata.freevars.get_index(cell_idx - self.metadata.cellvars.len()))
+            .map(String::as_ref)
+            .unwrap_or("<invalid var num>")
+    }
 }
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@crates/codegen/src/ir.rs` around lines 1112 - 1113, get_varname currently
indexes only metadata.varnames and will panic for VarNum values that refer into
the "locals-plus" table (cellvars/freevars created for LOAD_CLOSURE); update
get_varname to resolve VarNum against the combined locals-plus layout: if
var_num.as_usize() < self.metadata.varnames.len() return that entry, otherwise
subtract varnames.len() and if the remaining index is <
self.metadata.cellvars.len() return from metadata.cellvars, else subtract
cellvars.len() and return from metadata.freevars; ensure indexes use as_usize()
and do bounds-safe lookups so it never panics when encountering
LOAD_CLOSURE-created slots.

}

fn get_cell_name(&self, i: usize) -> &str {
Expand Down
27 changes: 21 additions & 6 deletions crates/compiler-core/src/bytecode.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use bitflags::bitflags;
use core::{
cell::UnsafeCell,
hash, mem,
ops::{Deref, Index},
ops::{Deref, Index, IndexMut},
sync::atomic::{AtomicU8, AtomicU16, AtomicUsize, Ordering},
};
use itertools::Itertools;
Expand Down Expand Up @@ -318,6 +318,22 @@ impl<C: Constant> FromIterator<C> for Constants<C> {
}
}

// TODO: Newtype "CodeObject.varnames". Make sure only `oparg:VarNum` can be used as index
impl<T> Index<oparg::VarNum> for [T] {
type Output = T;

fn index(&self, var_num: oparg::VarNum) -> &Self::Output {
&self[var_num.as_usize()]
}
}

// TODO: Newtype "CodeObject.varnames". Make sure only `oparg:VarNum` can be used as index
impl<T> IndexMut<oparg::VarNum> for [T] {
fn index_mut(&mut self, var_num: oparg::VarNum) -> &mut Self::Output {
&mut self[var_num.as_usize()]
}
}

/// Primary container of a single code object. Each python function has
/// a code object. Also a module has a code object.
#[derive(Clone)]
Expand Down Expand Up @@ -1037,8 +1053,7 @@ impl<C: Constant> CodeObject<C> {
pub fn map_bag<Bag: ConstantBag>(self, bag: Bag) -> CodeObject<Bag::Constant> {
let map_names = |names: Box<[C::Name]>| {
names
.into_vec()
.into_iter()
.iter()
.map(|x| bag.make_name(x.as_ref()))
.collect::<Box<[_]>>()
};
Expand Down Expand Up @@ -1123,7 +1138,7 @@ pub trait InstrDisplayContext {

fn get_name(&self, i: usize) -> &str;

fn get_varname(&self, i: usize) -> &str;
fn get_varname(&self, var_num: oparg::VarNum) -> &str;

fn get_cell_name(&self, i: usize) -> &str;
}
Expand All @@ -1139,8 +1154,8 @@ impl<C: Constant> InstrDisplayContext for CodeObject<C> {
self.names[i].as_ref()
}

fn get_varname(&self, i: usize) -> &str {
self.varnames[i].as_ref()
fn get_varname(&self, var_num: oparg::VarNum) -> &str {
self.varnames[var_num].as_ref()
}

fn get_cell_name(&self, i: usize) -> &str {
Expand Down
34 changes: 20 additions & 14 deletions crates/compiler-core/src/bytecode/instruction.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@ pub enum Instruction {
i: Arg<NameIdx>,
} = 62,
DeleteFast {
var_num: Arg<NameIdx>,
var_num: Arg<oparg::VarNum>,
} = 63,
DeleteGlobal {
namei: Arg<NameIdx>,
Expand Down Expand Up @@ -192,19 +192,19 @@ pub enum Instruction {
i: Arg<NameIdx>,
} = 83,
LoadFast {
var_num: Arg<NameIdx>,
var_num: Arg<oparg::VarNum>,
} = 84,
LoadFastAndClear {
var_num: Arg<NameIdx>,
var_num: Arg<oparg::VarNum>,
} = 85,
LoadFastBorrow {
var_num: Arg<NameIdx>,
var_num: Arg<oparg::VarNum>,
} = 86,
LoadFastBorrowLoadFastBorrow {
var_nums: Arg<u32>,
} = 87,
LoadFastCheck {
var_num: Arg<NameIdx>,
var_num: Arg<oparg::VarNum>,
} = 88,
LoadFastLoadFast {
var_nums: Arg<u32>,
Expand Down Expand Up @@ -276,7 +276,7 @@ pub enum Instruction {
i: Arg<NameIdx>,
} = 111,
StoreFast {
var_num: Arg<NameIdx>,
var_num: Arg<oparg::VarNum>,
} = 112,
StoreFastLoadFast {
var_nums: Arg<StoreFastLoadFast>,
Expand Down Expand Up @@ -1105,7 +1105,13 @@ impl InstructionMetadata for Instruction {
};
($variant:ident, $map:ident = $arg_marker:expr) => {{
let arg = $arg_marker.get(arg);
write!(f, "{:pad$}({}, {})", stringify!($variant), arg, $map(arg))
write!(
f,
"{:pad$}({}, {})",
stringify!($variant),
u32::from(arg),
$map(arg)
)
}};
($variant:ident, $arg_marker:expr) => {
write!(f, "{:pad$}({})", stringify!($variant), $arg_marker.get(arg))
Expand All @@ -1120,7 +1126,7 @@ impl InstructionMetadata for Instruction {
};
}

let varname = |i: u32| ctx.get_varname(i as usize);
let varname = |var_num: oparg::VarNum| ctx.get_varname(var_num);
let name = |i: u32| ctx.get_name(i as usize);
let cell_name = |i: u32| ctx.get_cell_name(i as usize);

Expand Down Expand Up @@ -1226,16 +1232,16 @@ impl InstructionMetadata for Instruction {
let oparg = var_nums.get(arg);
let idx1 = oparg >> 4;
let idx2 = oparg & 15;
let name1 = varname(idx1);
let name2 = varname(idx2);
let name1 = varname(idx1.into());
let name2 = varname(idx2.into());
write!(f, "{:pad$}({}, {})", "LOAD_FAST_LOAD_FAST", name1, name2)
}
Self::LoadFastBorrowLoadFastBorrow { var_nums } => {
let oparg = var_nums.get(arg);
let idx1 = oparg >> 4;
let idx2 = oparg & 15;
let name1 = varname(idx1);
let name2 = varname(idx2);
let name1 = varname(idx1.into());
let name2 = varname(idx2.into());
write!(
f,
"{:pad$}({}, {})",
Expand Down Expand Up @@ -1362,8 +1368,8 @@ impl InstructionMetadata for Instruction {
f,
"{:pad$}({}, {})",
"STORE_FAST_STORE_FAST",
varname(idx1),
varname(idx2)
varname(idx1.into()),
varname(idx2.into())
)
}
Self::StoreGlobal { namei } => w!(STORE_GLOBAL, name = namei),
Expand Down
71 changes: 44 additions & 27 deletions crates/compiler-core/src/bytecode/oparg.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -873,38 +873,55 @@ impl LoadAttrBuilder {
}
}

#[derive(Clone, Copy)]
pub struct ConstIdx(u32);
macro_rules! newtype_oparg {
(
$(#[$oparg_meta:meta])*
$vis:vis struct $name:ident(u32)
) => {
$(#[$oparg_meta])*
$vis struct $name(u32);

impl ConstIdx {
#[must_use]
pub const fn from_u32(value: u32) -> Self {
Self(value)
}
impl $name {
#[must_use]
pub const fn from_u32(value: u32) -> Self {
Self(value)
}

/// Returns the index as a `u32` value.
#[must_use]
pub const fn as_u32(self) -> u32 {
self.0
}
/// Returns the oparg as a `u32` value.
#[must_use]
pub const fn as_u32(self) -> u32 {
self.0
}

/// Returns the index as a `usize` value.
#[must_use]
pub const fn as_usize(self) -> usize {
self.0 as usize
}
}
/// Returns the oparg as a `usize` value.
#[must_use]
pub const fn as_usize(self) -> usize {
self.0 as usize
}
}

impl From<u32> for ConstIdx {
fn from(value: u32) -> Self {
Self::from_u32(value)
}
}
impl From<u32> for $name {
fn from(value: u32) -> Self {
Self::from_u32(value)
}
}

impl From<$name> for u32 {
fn from(value: $name) -> Self {
value.as_u32()
}
}

impl From<ConstIdx> for u32 {
fn from(consti: ConstIdx) -> Self {
consti.as_u32()
impl OpArgType for $name {}
}
}

impl OpArgType for ConstIdx {}
newtype_oparg!(
#[derive(Clone, Copy)]
pub struct ConstIdx(u32)
);

newtype_oparg!(
#[derive(Clone, Copy)]
pub struct VarNum(u32)
);
17 changes: 8 additions & 9 deletions crates/jit/src/instructions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use cranelift::prelude::*;
use num_traits::cast::ToPrimitive;
use rustpython_compiler_core::bytecode::{
self, BinaryOperator, BorrowedConstant, CodeObject, ComparisonOperator, Instruction,
IntrinsicFunction1, Label, OpArg, OpArgState,
IntrinsicFunction1, Label, OpArg, OpArgState, oparg,
};
use std::collections::HashMap;

Expand Down Expand Up @@ -94,7 +94,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
let params = compiler.builder.func.dfg.block_params(entry_block).to_vec();
for (i, (ty, val)) in arg_types.iter().zip(params).enumerate() {
compiler
.store_variable(i as u32, JitValue::from_type_and_value(ty.clone(), val))
.store_variable(
(i as u32).into(),
JitValue::from_type_and_value(ty.clone(), val),
)
.unwrap();
}
compiler
Expand All @@ -105,14 +108,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
self.stack.drain(stack_len - count..).collect()
}

fn store_variable(
&mut self,
idx: bytecode::NameIdx,
val: JitValue,
) -> Result<(), JitCompileError> {
fn store_variable(&mut self, idx: oparg::VarNum, val: JitValue) -> Result<(), JitCompileError> {
let builder = &mut self.builder;
let ty = val.to_jit_type().ok_or(JitCompileError::NotSupported)?;
let local = self.variables[idx as usize].get_or_insert_with(|| {
let local = self.variables[idx].get_or_insert_with(|| {
let var = builder.declare_var(ty.to_cranelift());
Comment on lines +111 to 115
Copy link
Contributor

@coderabbitai coderabbitai bot Mar 11, 2026

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major

Use checked local-slot access instead of panicking on bad bytecode.

Line 114 and Line 651 index self.variables directly, but VarNum is only a typed u32, not a proof that the slot exists. A malformed or stale code object will panic the JIT here instead of taking the existing JitCompileError::BadBytecode path.

🛠️ Suggested fix 
 fn store_variable(&mut self, idx: oparg::VarNum, val: JitValue) -> Result<(), JitCompileError> {
     let builder = &mut self.builder;
     let ty = val.to_jit_type().ok_or(JitCompileError::NotSupported)?;
-    let local = self.variables[idx].get_or_insert_with(|| {
+    let slot = self
+        .variables
+        .get_mut(idx.as_usize())
+        .ok_or(JitCompileError::BadBytecode)?;
+    let local = slot.get_or_insert_with(|| {
         let var = builder.declare_var(ty.to_cranelift());
         Local {
             var,
             ty: ty.clone(),
         }
@@
             Instruction::LoadFast { var_num } | Instruction::LoadFastBorrow { var_num } => {
-                let local = self.variables[var_num.get(arg)]
-                    .as_ref()
+                let local = self
+                    .variables
+                    .get(var_num.get(arg).as_usize())
+                    .and_then(Option::as_ref)
                     .ok_or(JitCompileError::BadBytecode)?;
                 self.stack.push(JitValue::from_type_and_value(
                     local.ty.clone(),
                     self.builder.use_var(local.var),
                 ));

Also applies to: 650-652

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@crates/jit/src/instructions.rs` around lines 111 - 115, The code directly
indexes self.variables with a VarNum (typed u32) which can panic on malformed
bytecode; change both occurrences (e.g., in store_variable and the other site
where self.variables[idx] is used) to first convert the VarNum to usize and use
checked access (self.variables.get_mut(idx_usize) or
self.variables.get(idx_usize)), returning Err(JitCompileError::BadBytecode) when
None. Inside the Some(slot) branch, keep the existing logic (declare_var via
builder and set the slot) but perform it via the checked slot reference instead
of indexing to avoid panics.

Local {
var,
Expand Down Expand Up @@ -649,7 +648,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
Ok(())
}
Instruction::LoadFast { var_num } | Instruction::LoadFastBorrow { var_num } => {
let local = self.variables[var_num.get(arg) as usize]
let local = self.variables[var_num.get(arg)]
.as_ref()
.ok_or(JitCompileError::BadBytecode)?;
self.stack.push(JitValue::from_type_and_value(
Expand Down
Loading
Loading