diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 0b5733a6d46..7c842286e87 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -715,14 +715,14 @@ impl Compiler { } /// Get the index of a local variable. - fn get_local_var_index(&mut self, name: &str) -> CompileResult { + fn get_local_var_index(&mut self, name: &str) -> CompileResult { 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. @@ -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() }); @@ -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 { + + fn varname(&mut self, name: &str) -> CompileResult { // 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, - ) -> bytecode::NameIdx { + ) -> u32 { let name = self.mangle(name); let cache = cache(self.current_code_info()); cache @@ -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 }); } } diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index a5936ef722b..8d5fbdb8bde 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -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() } fn get_cell_name(&self, i: usize) -> &str { diff --git a/crates/compiler-core/src/bytecode.rs b/crates/compiler-core/src/bytecode.rs index d7ed55930f6..80cf01bc027 100644 --- a/crates/compiler-core/src/bytecode.rs +++ b/crates/compiler-core/src/bytecode.rs @@ -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; @@ -318,6 +318,22 @@ impl FromIterator for Constants { } } +// TODO: Newtype "CodeObject.varnames". Make sure only `oparg:VarNum` can be used as index +impl Index 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 IndexMut 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)] @@ -1037,8 +1053,7 @@ impl CodeObject { pub fn map_bag(self, bag: Bag) -> CodeObject { let map_names = |names: Box<[C::Name]>| { names - .into_vec() - .into_iter() + .iter() .map(|x| bag.make_name(x.as_ref())) .collect::>() }; @@ -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; } @@ -1139,8 +1154,8 @@ impl InstrDisplayContext for CodeObject { 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 { diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs index e1b2ffd65c7..7b3f2c816b6 100644 --- a/crates/compiler-core/src/bytecode/instruction.rs +++ b/crates/compiler-core/src/bytecode/instruction.rs @@ -133,7 +133,7 @@ pub enum Instruction { i: Arg, } = 62, DeleteFast { - var_num: Arg, + var_num: Arg, } = 63, DeleteGlobal { namei: Arg, @@ -192,19 +192,19 @@ pub enum Instruction { i: Arg, } = 83, LoadFast { - var_num: Arg, + var_num: Arg, } = 84, LoadFastAndClear { - var_num: Arg, + var_num: Arg, } = 85, LoadFastBorrow { - var_num: Arg, + var_num: Arg, } = 86, LoadFastBorrowLoadFastBorrow { var_nums: Arg, } = 87, LoadFastCheck { - var_num: Arg, + var_num: Arg, } = 88, LoadFastLoadFast { var_nums: Arg, @@ -276,7 +276,7 @@ pub enum Instruction { i: Arg, } = 111, StoreFast { - var_num: Arg, + var_num: Arg, } = 112, StoreFastLoadFast { var_nums: Arg, @@ -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)) @@ -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); @@ -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$}({}, {})", @@ -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), diff --git a/crates/compiler-core/src/bytecode/oparg.rs b/crates/compiler-core/src/bytecode/oparg.rs index f57a2b6fdab..2cb79213125 100644 --- a/crates/compiler-core/src/bytecode/oparg.rs +++ b/crates/compiler-core/src/bytecode/oparg.rs @@ -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 for ConstIdx { - fn from(value: u32) -> Self { - Self::from_u32(value) - } -} + impl From 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 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) +); diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 01f13e9d289..a215710da3b 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -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; @@ -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 @@ -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()); Local { var, @@ -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( diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index a70236a223a..3cfd3eef105 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -2162,9 +2162,9 @@ impl ExecutingFrame<'_> { self.cell_ref(i.get(arg) as usize).set(None); Ok(None) } - Instruction::DeleteFast { var_num: idx } => { + Instruction::DeleteFast { var_num } => { let fastlocals = self.localsplus.fastlocals_mut(); - let idx = idx.get(arg) as usize; + let idx = var_num.get(arg); if fastlocals[idx].is_none() { return Err(vm.new_exception_msg( vm.ctx.exceptions.unbound_local_error.to_owned(), @@ -2680,7 +2680,7 @@ impl ExecutingFrame<'_> { self.push_value(x); Ok(None) } - Instruction::LoadFast { var_num: idx } => { + Instruction::LoadFast { var_num } => { #[cold] fn reference_error( varname: &'static PyStrInterned, @@ -2691,27 +2691,27 @@ impl ExecutingFrame<'_> { format!("local variable '{varname}' referenced before assignment").into(), ) } - let idx = idx.get(arg) as usize; + let idx = var_num.get(arg); let x = self.localsplus.fastlocals()[idx] .clone() .ok_or_else(|| reference_error(self.code.varnames[idx], vm))?; self.push_value(x); Ok(None) } - Instruction::LoadFastAndClear { var_num: idx } => { + Instruction::LoadFastAndClear { var_num } => { // Load value and clear the slot (for inlined comprehensions) // If slot is empty, push None (not an error - variable may not exist yet) - let idx = idx.get(arg) as usize; + let idx = var_num.get(arg); let x = self.localsplus.fastlocals_mut()[idx] .take() .unwrap_or_else(|| vm.ctx.none()); self.push_value(x); Ok(None) } - Instruction::LoadFastCheck { var_num: idx } => { + Instruction::LoadFastCheck { var_num } => { // Same as LoadFast but explicitly checks for unbound locals // (LoadFast in RustPython already does this check) - let idx = idx.get(arg) as usize; + let idx = var_num.get(arg); let x = self.localsplus.fastlocals()[idx].clone().ok_or_else(|| { vm.new_exception_msg( vm.ctx.exceptions.unbound_local_error.to_owned(), @@ -2759,8 +2759,8 @@ impl ExecutingFrame<'_> { // Borrow optimization not yet active; falls back to clone. // push_borrowed() is available but disabled until stack // lifetime issues at yield/exception points are resolved. - Instruction::LoadFastBorrow { var_num: idx } => { - let idx = idx.get(arg) as usize; + Instruction::LoadFastBorrow { var_num } => { + let idx = var_num.get(arg); let x = self.localsplus.fastlocals()[idx].clone().ok_or_else(|| { vm.new_exception_msg( vm.ctx.exceptions.unbound_local_error.to_owned(), @@ -3299,10 +3299,10 @@ impl ExecutingFrame<'_> { self.cell_ref(i.get(arg) as usize).set(Some(value)); Ok(None) } - Instruction::StoreFast { var_num: idx } => { + Instruction::StoreFast { var_num } => { let value = self.pop_value(); let fastlocals = self.localsplus.fastlocals_mut(); - fastlocals[idx.get(arg) as usize] = Some(value); + fastlocals[var_num.get(arg)] = Some(value); Ok(None) } Instruction::StoreFastLoadFast { var_nums } => {