Skip to content

Commit 82e9b5d

Browse files
ShaharNavehShaharNaveh
authored
Newtype ConstIdx, Constants (#7358)
* Newtype ConstIdx, Constants * Set generic
1 parent fa4f84c commit 82e9b5d

File tree

10 files changed

+116
-52
lines changed

10 files changed

+116
-52
lines changed

‎crates/codegen/src/compile.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ use rustpython_compiler_core::{
2929
self, AnyInstruction, Arg as OpArgMarker, BinaryOperator, BuildSliceArgCount, CodeObject,
3030
ComparisonOperator, ConstantData, ConvertValueOparg, Instruction, IntrinsicFunction1,
3131
Invert, LoadAttr, LoadSuperAttr, OpArg, OpArgType, PseudoInstruction, SpecialMethod,
32-
UnpackExArgs,
32+
UnpackExArgs, oparg,
3333
},
3434
};
3535
use rustpython_wtf8::Wtf8Buf;
@@ -8075,9 +8075,9 @@ impl Compiler {
80758075

80768076
// fn block_done()
80778077

8078-
fn arg_constant(&mut self, constant: ConstantData) -> u32 {
8078+
fn arg_constant(&mut self, constant: ConstantData) -> oparg::ConstIdx {
80798079
let info = self.current_code_info();
8080-
info.metadata.consts.insert_full(constant).0.to_u32()
8080+
info.metadata.consts.insert_full(constant).0.to_u32().into()
80818081
}
80828082

80838083
fn emit_load_const(&mut self, constant: ConstantData) {

‎crates/codegen/src/ir.rs‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use rustpython_compiler_core::{
1010
bytecode::{
1111
AnyInstruction, Arg, CodeFlags, CodeObject, CodeUnit, CodeUnits, ConstantData,
1212
ExceptionTableEntry, InstrDisplayContext, Instruction, InstructionMetadata, Label, OpArg,
13-
PseudoInstruction, PyCodeLocationInfoKind, encode_exception_table,
13+
PseudoInstruction, PyCodeLocationInfoKind, encode_exception_table, oparg,
1414
},
1515
varint::{write_signed_varint, write_varint},
1616
};
@@ -695,15 +695,15 @@ impl CodeInfo {
695695
}
696696
(Instruction::LoadConst { consti }, Instruction::ToBool) => {
697697
let consti = consti.get(curr.arg);
698-
let constant = &self.metadata.consts[consti as usize];
698+
let constant = &self.metadata.consts[consti.as_usize()];
699699
if let ConstantData::Boolean { .. } = constant {
700-
Some((curr_instr, OpArg::from(consti)))
700+
Some((curr_instr, OpArg::from(consti.as_u32())))
701701
} else {
702702
None
703703
}
704704
}
705705
(Instruction::LoadConst { consti }, Instruction::UnaryNot) => {
706-
let constant = &self.metadata.consts[consti.get(curr.arg) as usize];
706+
let constant = &self.metadata.consts[consti.get(curr.arg).as_usize()];
707707
match constant {
708708
ConstantData::Boolean { value } => {
709709
let (const_idx, _) = self
@@ -1100,15 +1100,19 @@ impl CodeInfo {
11001100

11011101
impl InstrDisplayContext for CodeInfo {
11021102
type Constant = ConstantData;
1103-
fn get_constant(&self, i: usize) -> &ConstantData {
1104-
&self.metadata.consts[i]
1103+
1104+
fn get_constant(&self, consti: oparg::ConstIdx) -> &ConstantData {
1105+
&self.metadata.consts[consti.as_usize()]
11051106
}
1107+
11061108
fn get_name(&self, i: usize) -> &str {
11071109
self.metadata.names[i].as_ref()
11081110
}
1111+
11091112
fn get_varname(&self, i: usize) -> &str {
11101113
self.metadata.varnames[i].as_ref()
11111114
}
1115+
11121116
fn get_cell_name(&self, i: usize) -> &str {
11131117
self.metadata
11141118
.cellvars

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

Lines changed: 32 additions & 8 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,
14+
ops::{Deref, Index},
1515
sync::atomic::{AtomicU8, AtomicU16, AtomicUsize, Ordering},
1616
};
1717
use itertools::Itertools;
@@ -32,7 +32,7 @@ pub use crate::bytecode::{
3232
};
3333

3434
mod instruction;
35-
mod oparg;
35+
pub mod oparg;
3636

3737
/// Exception table entry for zero-cost exception handling
3838
/// Format: (start, size, target, depth<<1|lasti)
@@ -293,6 +293,31 @@ impl ConstantBag for BasicBag {
293293
}
294294
}
295295

296+
#[derive(Clone)]
297+
pub struct Constants<C: Constant>(Box<[C]>);
298+
299+
impl<C: Constant> Deref for Constants<C> {
300+
type Target = [C];
301+
302+
fn deref(&self) -> &Self::Target {
303+
&self.0
304+
}
305+
}
306+
307+
impl<C: Constant> Index<oparg::ConstIdx> for Constants<C> {
308+
type Output = C;
309+
310+
fn index(&self, consti: oparg::ConstIdx) -> &Self::Output {
311+
&self.0[consti.as_usize()]
312+
}
313+
}
314+
315+
impl<C: Constant> FromIterator<C> for Constants<C> {
316+
fn from_iter<T: IntoIterator<Item = C>>(iter: T) -> Self {
317+
Self(iter.into_iter().collect())
318+
}
319+
}
320+
296321
/// Primary container of a single code object. Each python function has
297322
/// a code object. Also a module has a code object.
298323
#[derive(Clone)]
@@ -312,7 +337,7 @@ pub struct CodeObject<C: Constant = ConstantData> {
312337
/// Qualified name of the object (like CPython's co_qualname)
313338
pub qualname: C::Name,
314339
pub cell2arg: Option<Box<[i32]>>,
315-
pub constants: Box<[C]>,
340+
pub constants: Constants<C>,
316341
pub names: Box<[C::Name]>,
317342
pub varnames: Box<[C::Name]>,
318343
pub cellvars: Box<[C::Name]>,
@@ -1020,8 +1045,7 @@ impl<C: Constant> CodeObject<C> {
10201045
CodeObject {
10211046
constants: self
10221047
.constants
1023-
.into_vec()
1024-
.into_iter()
1048+
.iter()
10251049
.map(|x| bag.make_constant(x.borrow_constant()))
10261050
.collect(),
10271051
names: map_names(self.names),
@@ -1095,7 +1119,7 @@ impl<C: Constant> fmt::Display for CodeObject<C> {
10951119
pub trait InstrDisplayContext {
10961120
type Constant: Constant;
10971121

1098-
fn get_constant(&self, i: usize) -> &Self::Constant;
1122+
fn get_constant(&self, consti: oparg::ConstIdx) -> &Self::Constant;
10991123

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

@@ -1107,8 +1131,8 @@ pub trait InstrDisplayContext {
11071131
impl<C: Constant> InstrDisplayContext for CodeObject<C> {
11081132
type Constant = C;
11091133

1110-
fn get_constant(&self, i: usize) -> &C {
1111-
&self.constants[i]
1134+
fn get_constant(&self, consti: oparg::ConstIdx) -> &C {
1135+
&self.constants[consti]
11121136
}
11131137

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

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

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use crate::{
44
bytecode::{
55
BorrowedConstant, Constant, InstrDisplayContext,
66
oparg::{
7-
BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator,
7+
self, BinaryOperator, BuildSliceArgCount, CommonConstant, ComparisonOperator,
88
ConvertValueOparg, IntrinsicFunction1, IntrinsicFunction2, Invert, Label, LoadAttr,
99
LoadSuperAttr, MakeFunctionFlags, NameIdx, OpArg, OpArgByte, OpArgType, RaiseKind,
1010
SpecialMethod, StoreFastLoadFast, UnpackExArgs,
@@ -186,7 +186,7 @@ pub enum Instruction {
186186
idx: Arg<CommonConstant>,
187187
} = 81,
188188
LoadConst {
189-
consti: Arg<u32>,
189+
consti: Arg<oparg::ConstIdx>,
190190
} = 82,
191191
LoadDeref {
192192
i: Arg<NameIdx>,
@@ -1124,22 +1124,25 @@ impl InstructionMetadata for Instruction {
11241124
let name = |i: u32| ctx.get_name(i as usize);
11251125
let cell_name = |i: u32| ctx.get_cell_name(i as usize);
11261126

1127-
let fmt_const =
1128-
|op: &str, arg: OpArg, f: &mut fmt::Formatter<'_>, idx: &Arg<u32>| -> fmt::Result {
1129-
let value = ctx.get_constant(idx.get(arg) as usize);
1130-
match value.borrow_constant() {
1131-
BorrowedConstant::Code { code } if expand_code_objects => {
1132-
write!(f, "{op:pad$}({code:?}):")?;
1133-
code.display_inner(f, true, level + 1)?;
1134-
Ok(())
1135-
}
1136-
c => {
1137-
write!(f, "{op:pad$}(")?;
1138-
c.fmt_display(f)?;
1139-
write!(f, ")")
1140-
}
1127+
let fmt_const = |op: &str,
1128+
arg: OpArg,
1129+
f: &mut fmt::Formatter<'_>,
1130+
consti: &Arg<oparg::ConstIdx>|
1131+
-> fmt::Result {
1132+
let value = ctx.get_constant(consti.get(arg));
1133+
match value.borrow_constant() {
1134+
BorrowedConstant::Code { code } if expand_code_objects => {
1135+
write!(f, "{op:pad$}({code:?}):")?;
1136+
code.display_inner(f, true, level + 1)?;
1137+
Ok(())
11411138
}
1142-
};
1139+
c => {
1140+
write!(f, "{op:pad$}(")?;
1141+
c.fmt_display(f)?;
1142+
write!(f, ")")
1143+
}
1144+
}
1145+
};
11431146

11441147
match self {
11451148
Self::BinarySlice => w!(BINARY_SLICE),

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -872,3 +872,39 @@ impl LoadAttrBuilder {
872872
self
873873
}
874874
}
875+
876+
#[derive(Clone, Copy)]
877+
pub struct ConstIdx(u32);
878+
879+
impl ConstIdx {
880+
#[must_use]
881+
pub const fn from_u32(value: u32) -> Self {
882+
Self(value)
883+
}
884+
885+
/// Returns the index as a `u32` value.
886+
#[must_use]
887+
pub const fn as_u32(self) -> u32 {
888+
self.0
889+
}
890+
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+
}
897+
898+
impl From<u32> for ConstIdx {
899+
fn from(value: u32) -> Self {
900+
Self::from_u32(value)
901+
}
902+
}
903+
904+
impl From<ConstIdx> for u32 {
905+
fn from(consti: ConstIdx) -> Self {
906+
consti.as_u32()
907+
}
908+
}
909+
910+
impl OpArgType for ConstIdx {}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -240,7 +240,7 @@ pub fn deserialize_code<R: Read, Bag: ConstantBag>(
240240
let len = rdr.read_u32()?;
241241
let constants = (0..len)
242242
.map(|_| deserialize_value(rdr, bag))
243-
.collect::<Result<Box<[_]>>>()?;
243+
.collect::<Result<_>>()?;
244244

245245
let mut read_names = || {
246246
let len = rdr.read_u32()?;

‎crates/jit/src/instructions.rs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -637,9 +637,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
637637
Ok(())
638638
}
639639
Instruction::LoadConst { consti } => {
640-
let val = self.prepare_const(
641-
bytecode.constants[consti.get(arg) as usize].borrow_constant(),
642-
)?;
640+
let val =
641+
self.prepare_const(bytecode.constants[consti.get(arg)].borrow_constant())?;
643642
self.stack.push(val);
644643
Ok(())
645644
}

‎crates/jit/tests/common.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use core::ops::ControlFlow;
22
use rustpython_compiler_core::bytecode::{
3-
CodeObject, ConstantData, Instruction, OpArg, OpArgState,
3+
CodeObject, ConstantData, Constants, Instruction, OpArg, OpArgState,
44
};
55
use rustpython_jit::{CompiledCode, JitType};
66
use rustpython_wtf8::{Wtf8, Wtf8Buf};
@@ -77,7 +77,7 @@ fn extract_annotations_from_annotate_code(code: &CodeObject) -> HashMap<Wtf8Buf,
7777

7878
match instruction {
7979
Instruction::LoadConst { consti } => {
80-
stack.push((true, consti.get(arg) as usize));
80+
stack.push((true, consti.get(arg).as_usize()));
8181
}
8282
Instruction::LoadName { namei } => {
8383
stack.push((false, namei.get(arg) as usize));
@@ -99,7 +99,8 @@ fn extract_annotations_from_annotate_code(code: &CodeObject) -> HashMap<Wtf8Buf,
9999

100100
// Key should be a const string (parameter name)
101101
if key_is_const
102-
&& let ConstantData::Str { value } = &code.constants[key_idx]
102+
&& let ConstantData::Str { value } =
103+
&code.constants[(key_idx as u32).into()]
103104
{
104105
let param_name = value;
105106
// Value can be a name (type ref) or a const string (forward ref)
@@ -185,16 +186,15 @@ impl StackMachine {
185186
&mut self,
186187
instruction: Instruction,
187188
arg: OpArg,
188-
constants: &[ConstantData],
189+
constants: &Constants<ConstantData>,
189190
names: &[String],
190191
) -> ControlFlow<()> {
191192
match instruction {
192193
Instruction::Resume { .. } | Instruction::Cache | Instruction::NotTaken => {
193194
// No-op for JIT tests
194195
}
195196
Instruction::LoadConst { consti } => {
196-
let idx = consti.get(arg);
197-
self.stack.push(constants[idx as usize].clone().into())
197+
self.stack.push(constants[consti.get(arg)].clone().into())
198198
}
199199
Instruction::LoadName { namei } => self
200200
.stack

‎crates/vm/src/builtins/code.rs‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -538,16 +538,14 @@ impl Constructor for PyCode {
538538
.map_err(|e| vm.new_value_error(format!("invalid bytecode: {}", e)))?;
539539

540540
// Convert constants
541-
let constants: Box<[Literal]> = args
541+
let constants = args
542542
.consts
543543
.iter()
544544
.map(|obj| {
545-
// Convert PyObject to Literal constant
546-
// For now, just wrap it
545+
// Convert PyObject to Literal constant. For now, just wrap it
547546
Literal(obj.clone())
548547
})
549-
.collect::<Vec<_>>()
550-
.into_boxed_slice();
548+
.collect();
551549

552550
// Create locations (start and end pairs)
553551
let row = if args.firstlineno > 0 {

‎crates/vm/src/frame.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2375,8 +2375,8 @@ impl ExecutingFrame<'_> {
23752375
});
23762376
Ok(None)
23772377
}
2378-
Instruction::LoadConst { consti: idx } => {
2379-
self.push_value(self.code.constants[idx.get(arg) as usize].clone().into());
2378+
Instruction::LoadConst { consti } => {
2379+
self.push_value(self.code.constants[consti.get(arg)].clone().into());
23802380
// Mirror CPython's LOAD_CONST family transition. RustPython does
23812381
// not currently distinguish immortal constants at runtime.
23822382
let instr_idx = self.lasti() as usize - 1;
@@ -2388,7 +2388,7 @@ impl ExecutingFrame<'_> {
23882388
Ok(None)
23892389
}
23902390
Instruction::LoadConstMortal | Instruction::LoadConstImmortal => {
2391-
self.push_value(self.code.constants[u32::from(arg) as usize].clone().into());
2391+
self.push_value(self.code.constants[u32::from(arg).into()].clone().into());
23922392
Ok(None)
23932393
}
23942394
Instruction::LoadCommonConstant { idx } => {

0 commit comments

Comments
 (0)