Skip to content

Commit 5c631e5

Browse files
ShaharNavehShaharNaveh
authored
Newtype var_num oparg (#7400)
1 parent 7e47270 commit 5c631e5

7 files changed

Lines changed: 124 additions & 77 deletions

File tree

‎crates/codegen/src/compile.rs‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -715,14 +715,14 @@ impl Compiler {
715715
}
716716

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

728728
/// Get the index of a global name.
@@ -1283,7 +1283,12 @@ impl Compiler {
12831283
/// if format > VALUE_WITH_FAKE_GLOBALS (2): raise NotImplementedError
12841284
fn emit_format_validation(&mut self) -> CompileResult<()> {
12851285
// Load format parameter (first local variable, index 0)
1286-
emit!(self, Instruction::LoadFast { var_num: 0 });
1286+
emit!(
1287+
self,
1288+
Instruction::LoadFast {
1289+
var_num: oparg::VarNum::from_u32(0)
1290+
}
1291+
);
12871292

12881293
// Load VALUE_WITH_FAKE_GLOBALS constant (2)
12891294
self.emit_load_const(ConstantData::Integer { value: 2.into() });
@@ -1562,15 +1567,19 @@ impl Compiler {
15621567
fn name(&mut self, name: &str) -> bytecode::NameIdx {
15631568
self._name_inner(name, |i| &mut i.metadata.names)
15641569
}
1565-
fn varname(&mut self, name: &str) -> CompileResult<bytecode::NameIdx> {
1570+
1571+
fn varname(&mut self, name: &str) -> CompileResult<oparg::VarNum> {
15661572
// Note: __debug__ checks are now handled in symboltable phase
1567-
Ok(self._name_inner(name, |i| &mut i.metadata.varnames))
1573+
Ok(oparg::VarNum::from_u32(
1574+
self._name_inner(name, |i| &mut i.metadata.varnames),
1575+
))
15681576
}
1577+
15691578
fn _name_inner(
15701579
&mut self,
15711580
name: &str,
15721581
cache: impl FnOnce(&mut ir::CodeInfo) -> &mut IndexSet<String>,
1573-
) -> bytecode::NameIdx {
1582+
) -> u32 {
15741583
let name = self.mangle(name);
15751584
let cache = cache(self.current_code_info());
15761585
cache
@@ -4129,7 +4138,8 @@ impl Compiler {
41294138

41304139
// Load defaults/kwdefaults with LOAD_FAST
41314140
for i in 0..num_typeparam_args {
4132-
emit!(self, Instruction::LoadFast { var_num: i as u32 });
4141+
let var_num = oparg::VarNum::from(i as u32);
4142+
emit!(self, Instruction::LoadFast { var_num });
41334143
}
41344144
}
41354145

‎crates/codegen/src/ir.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1109,8 +1109,8 @@ impl InstrDisplayContext for CodeInfo {
11091109
self.metadata.names[i].as_ref()
11101110
}
11111111

1112-
fn get_varname(&self, i: usize) -> &str {
1113-
self.metadata.varnames[i].as_ref()
1112+
fn get_varname(&self, var_num: oparg::VarNum) -> &str {
1113+
self.metadata.varnames[var_num.as_usize()].as_ref()
11141114
}
11151115

11161116
fn get_cell_name(&self, i: usize) -> &str {

‎crates/compiler-core/src/bytecode.rs‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use bitflags::bitflags;
1111
use core::{
1212
cell::UnsafeCell,
1313
hash, mem,
14-
ops::{Deref, Index},
14+
ops::{Deref, Index, IndexMut},
1515
sync::atomic::{AtomicU8, AtomicU16, AtomicUsize, Ordering},
1616
};
1717
use itertools::Itertools;
@@ -318,6 +318,22 @@ impl<C: Constant> FromIterator<C> for Constants<C> {
318318
}
319319
}
320320

321+
// TODO: Newtype "CodeObject.varnames". Make sure only `oparg:VarNum` can be used as index
322+
impl<T> Index<oparg::VarNum> for [T] {
323+
type Output = T;
324+
325+
fn index(&self, var_num: oparg::VarNum) -> &Self::Output {
326+
&self[var_num.as_usize()]
327+
}
328+
}
329+
330+
// TODO: Newtype "CodeObject.varnames". Make sure only `oparg:VarNum` can be used as index
331+
impl<T> IndexMut<oparg::VarNum> for [T] {
332+
fn index_mut(&mut self, var_num: oparg::VarNum) -> &mut Self::Output {
333+
&mut self[var_num.as_usize()]
334+
}
335+
}
336+
321337
/// Primary container of a single code object. Each python function has
322338
/// a code object. Also a module has a code object.
323339
#[derive(Clone)]
@@ -1037,8 +1053,7 @@ impl<C: Constant> CodeObject<C> {
10371053
pub fn map_bag<Bag: ConstantBag>(self, bag: Bag) -> CodeObject<Bag::Constant> {
10381054
let map_names = |names: Box<[C::Name]>| {
10391055
names
1040-
.into_vec()
1041-
.into_iter()
1056+
.iter()
10421057
.map(|x| bag.make_name(x.as_ref()))
10431058
.collect::<Box<[_]>>()
10441059
};
@@ -1123,7 +1138,7 @@ pub trait InstrDisplayContext {
11231138

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

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

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

1142-
fn get_varname(&self, i: usize) -> &str {
1143-
self.varnames[i].as_ref()
1157+
fn get_varname(&self, var_num: oparg::VarNum) -> &str {
1158+
self.varnames[var_num].as_ref()
11441159
}
11451160

11461161
fn get_cell_name(&self, i: usize) -> &str {

‎crates/compiler-core/src/bytecode/instruction.rs‎

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ pub enum Instruction {
133133
i: Arg<NameIdx>,
134134
} = 62,
135135
DeleteFast {
136-
var_num: Arg<NameIdx>,
136+
var_num: Arg<oparg::VarNum>,
137137
} = 63,
138138
DeleteGlobal {
139139
namei: Arg<NameIdx>,
@@ -192,19 +192,19 @@ pub enum Instruction {
192192
i: Arg<NameIdx>,
193193
} = 83,
194194
LoadFast {
195-
var_num: Arg<NameIdx>,
195+
var_num: Arg<oparg::VarNum>,
196196
} = 84,
197197
LoadFastAndClear {
198-
var_num: Arg<NameIdx>,
198+
var_num: Arg<oparg::VarNum>,
199199
} = 85,
200200
LoadFastBorrow {
201-
var_num: Arg<NameIdx>,
201+
var_num: Arg<oparg::VarNum>,
202202
} = 86,
203203
LoadFastBorrowLoadFastBorrow {
204204
var_nums: Arg<u32>,
205205
} = 87,
206206
LoadFastCheck {
207-
var_num: Arg<NameIdx>,
207+
var_num: Arg<oparg::VarNum>,
208208
} = 88,
209209
LoadFastLoadFast {
210210
var_nums: Arg<u32>,
@@ -276,7 +276,7 @@ pub enum Instruction {
276276
i: Arg<NameIdx>,
277277
} = 111,
278278
StoreFast {
279-
var_num: Arg<NameIdx>,
279+
var_num: Arg<oparg::VarNum>,
280280
} = 112,
281281
StoreFastLoadFast {
282282
var_nums: Arg<StoreFastLoadFast>,
@@ -1105,7 +1105,13 @@ impl InstructionMetadata for Instruction {
11051105
};
11061106
($variant:ident, $map:ident = $arg_marker:expr) => {{
11071107
let arg = $arg_marker.get(arg);
1108-
write!(f, "{:pad$}({}, {})", stringify!($variant), arg, $map(arg))
1108+
write!(
1109+
f,
1110+
"{:pad$}({}, {})",
1111+
stringify!($variant),
1112+
u32::from(arg),
1113+
$map(arg)
1114+
)
11091115
}};
11101116
($variant:ident, $arg_marker:expr) => {
11111117
write!(f, "{:pad$}({})", stringify!($variant), $arg_marker.get(arg))
@@ -1120,7 +1126,7 @@ impl InstructionMetadata for Instruction {
11201126
};
11211127
}
11221128

1123-
let varname = |i: u32| ctx.get_varname(i as usize);
1129+
let varname = |var_num: oparg::VarNum| ctx.get_varname(var_num);
11241130
let name = |i: u32| ctx.get_name(i as usize);
11251131
let cell_name = |i: u32| ctx.get_cell_name(i as usize);
11261132

@@ -1226,16 +1232,16 @@ impl InstructionMetadata for Instruction {
12261232
let oparg = var_nums.get(arg);
12271233
let idx1 = oparg >> 4;
12281234
let idx2 = oparg & 15;
1229-
let name1 = varname(idx1);
1230-
let name2 = varname(idx2);
1235+
let name1 = varname(idx1.into());
1236+
let name2 = varname(idx2.into());
12311237
write!(f, "{:pad$}({}, {})", "LOAD_FAST_LOAD_FAST", name1, name2)
12321238
}
12331239
Self::LoadFastBorrowLoadFastBorrow { var_nums } => {
12341240
let oparg = var_nums.get(arg);
12351241
let idx1 = oparg >> 4;
12361242
let idx2 = oparg & 15;
1237-
let name1 = varname(idx1);
1238-
let name2 = varname(idx2);
1243+
let name1 = varname(idx1.into());
1244+
let name2 = varname(idx2.into());
12391245
write!(
12401246
f,
12411247
"{:pad$}({}, {})",
@@ -1362,8 +1368,8 @@ impl InstructionMetadata for Instruction {
13621368
f,
13631369
"{:pad$}({}, {})",
13641370
"STORE_FAST_STORE_FAST",
1365-
varname(idx1),
1366-
varname(idx2)
1371+
varname(idx1.into()),
1372+
varname(idx2.into())
13671373
)
13681374
}
13691375
Self::StoreGlobal { namei } => w!(STORE_GLOBAL, name = namei),

‎crates/compiler-core/src/bytecode/oparg.rs‎

Lines changed: 44 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -873,38 +873,55 @@ impl LoadAttrBuilder {
873873
}
874874
}
875875

876-
#[derive(Clone, Copy)]
877-
pub struct ConstIdx(u32);
876+
macro_rules! newtype_oparg {
877+
(
878+
$(#[$oparg_meta:meta])*
879+
$vis:vis struct $name:ident(u32)
880+
) => {
881+
$(#[$oparg_meta])*
882+
$vis struct $name(u32);
878883

879-
impl ConstIdx {
880-
#[must_use]
881-
pub const fn from_u32(value: u32) -> Self {
882-
Self(value)
883-
}
884+
impl $name {
885+
#[must_use]
886+
pub const fn from_u32(value: u32) -> Self {
887+
Self(value)
888+
}
884889

885-
/// Returns the index as a `u32` value.
886-
#[must_use]
887-
pub const fn as_u32(self) -> u32 {
888-
self.0
889-
}
890+
/// Returns the oparg as a `u32` value.
891+
#[must_use]
892+
pub const fn as_u32(self) -> u32 {
893+
self.0
894+
}
890895

891-
/// Returns the index as a `usize` value.
892-
#[must_use]
893-
pub const fn as_usize(self) -> usize {
894-
self.0 as usize
895-
}
896-
}
896+
/// Returns the oparg as a `usize` value.
897+
#[must_use]
898+
pub const fn as_usize(self) -> usize {
899+
self.0 as usize
900+
}
901+
}
897902

898-
impl From<u32> for ConstIdx {
899-
fn from(value: u32) -> Self {
900-
Self::from_u32(value)
901-
}
902-
}
903+
impl From<u32> for $name {
904+
fn from(value: u32) -> Self {
905+
Self::from_u32(value)
906+
}
907+
}
908+
909+
impl From<$name> for u32 {
910+
fn from(value: $name) -> Self {
911+
value.as_u32()
912+
}
913+
}
903914

904-
impl From<ConstIdx> for u32 {
905-
fn from(consti: ConstIdx) -> Self {
906-
consti.as_u32()
915+
impl OpArgType for $name {}
907916
}
908917
}
909918

910-
impl OpArgType for ConstIdx {}
919+
newtype_oparg!(
920+
#[derive(Clone, Copy)]
921+
pub struct ConstIdx(u32)
922+
);
923+
924+
newtype_oparg!(
925+
#[derive(Clone, Copy)]
926+
pub struct VarNum(u32)
927+
);

‎crates/jit/src/instructions.rs‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use cranelift::prelude::*;
66
use num_traits::cast::ToPrimitive;
77
use rustpython_compiler_core::bytecode::{
88
self, BinaryOperator, BorrowedConstant, CodeObject, ComparisonOperator, Instruction,
9-
IntrinsicFunction1, Label, OpArg, OpArgState,
9+
IntrinsicFunction1, Label, OpArg, OpArgState, oparg,
1010
};
1111
use std::collections::HashMap;
1212

@@ -94,7 +94,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
9494
let params = compiler.builder.func.dfg.block_params(entry_block).to_vec();
9595
for (i, (ty, val)) in arg_types.iter().zip(params).enumerate() {
9696
compiler
97-
.store_variable(i as u32, JitValue::from_type_and_value(ty.clone(), val))
97+
.store_variable(
98+
(i as u32).into(),
99+
JitValue::from_type_and_value(ty.clone(), val),
100+
)
98101
.unwrap();
99102
}
100103
compiler
@@ -105,14 +108,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
105108
self.stack.drain(stack_len - count..).collect()
106109
}
107110

108-
fn store_variable(
109-
&mut self,
110-
idx: bytecode::NameIdx,
111-
val: JitValue,
112-
) -> Result<(), JitCompileError> {
111+
fn store_variable(&mut self, idx: oparg::VarNum, val: JitValue) -> Result<(), JitCompileError> {
113112
let builder = &mut self.builder;
114113
let ty = val.to_jit_type().ok_or(JitCompileError::NotSupported)?;
115-
let local = self.variables[idx as usize].get_or_insert_with(|| {
114+
let local = self.variables[idx].get_or_insert_with(|| {
116115
let var = builder.declare_var(ty.to_cranelift());
117116
Local {
118117
var,
@@ -649,7 +648,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
649648
Ok(())
650649
}
651650
Instruction::LoadFast { var_num } | Instruction::LoadFastBorrow { var_num } => {
652-
let local = self.variables[var_num.get(arg) as usize]
651+
let local = self.variables[var_num.get(arg)]
653652
.as_ref()
654653
.ok_or(JitCompileError::BadBytecode)?;
655654
self.stack.push(JitValue::from_type_and_value(

0 commit comments

Comments
 (0)