Skip to content

Commit aa9fc7f

Browse files
youknowonegithub-actions[bot]
youknowone
and
github-actions[bot]
authored
Pseudo ops (#6678)
* better expect * pseudo instruction * Fix CallKw narg * Auto-format: cargo fmt --all --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent 2d3ef86 commit aa9fc7f

File tree

6 files changed

+172
-61
lines changed

6 files changed

+172
-61
lines changed

‎crates/codegen/src/compile.rs‎

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6212,23 +6212,33 @@ impl Compiler {
62126212
}
62136213
emit!(self, Instruction::CallFunctionEx { has_kwargs });
62146214
} else if !arguments.keywords.is_empty() {
6215-
let mut kwarg_names = vec![];
6215+
// No **kwargs in this branch (has_double_star is false),
6216+
// so all keywords have arg.is_some()
6217+
let mut kwarg_names = Vec::with_capacity(arguments.keywords.len());
62166218
for keyword in &arguments.keywords {
6217-
if let Some(name) = &keyword.arg {
6218-
kwarg_names.push(ConstantData::Str {
6219-
value: name.as_str().into(),
6220-
});
6221-
} else {
6222-
// This means **kwargs!
6223-
panic!("name must be set");
6224-
}
6219+
let name = keyword
6220+
.arg
6221+
.as_ref()
6222+
.expect("has_double_star is false, so arg must be Some");
6223+
kwarg_names.push(ConstantData::Str {
6224+
value: name.as_str().into(),
6225+
});
62256226
self.compile_expression(&keyword.value)?;
62266227
}
62276228

62286229
self.emit_load_const(ConstantData::Tuple {
62296230
elements: kwarg_names,
62306231
});
6231-
emit!(self, Instruction::CallKw { nargs: count });
6232+
// nargs = positional args + keyword args
6233+
let positional = additional_positional
6234+
.checked_add(u32::try_from(arguments.args.len()).expect("too many positional args"))
6235+
.expect("too many positional args");
6236+
let keyword_count =
6237+
u32::try_from(arguments.keywords.len()).expect("too many keyword args");
6238+
let nargs = positional
6239+
.checked_add(keyword_count)
6240+
.expect("too many arguments");
6241+
emit!(self, Instruction::CallKw { nargs });
62326242
} else {
62336243
emit!(self, Instruction::Call { nargs: count });
62346244
}

‎crates/codegen/src/ir.rs‎

Lines changed: 56 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@ use crate::{IndexMap, IndexSet, error::InternalError};
44
use rustpython_compiler_core::{
55
OneIndexed, SourceLocation,
66
bytecode::{
7-
CodeFlags, CodeObject, CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry,
7+
Arg, CodeFlags, CodeObject, CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry,
88
InstrDisplayContext, Instruction, Label, OpArg, PyCodeLocationInfoKind,
9-
encode_exception_table,
9+
encode_exception_table, encode_load_attr_arg,
1010
},
1111
varint::{write_signed_varint, write_varint},
1212
};
@@ -189,6 +189,34 @@ impl CodeInfo {
189189
let mut instructions = Vec::new();
190190
let mut locations = Vec::new();
191191

192+
// convert_pseudo_ops: instructions before the main loop
193+
for block in blocks
194+
.iter_mut()
195+
.filter(|b| b.next != BlockIdx::NULL || !b.instructions.is_empty())
196+
{
197+
for info in &mut block.instructions {
198+
match info.instr {
199+
// LOAD_ATTR_METHOD pseudo → LOAD_ATTR (with method flag=1)
200+
Instruction::LoadAttrMethod { idx } => {
201+
let encoded = encode_load_attr_arg(idx.get(info.arg), true);
202+
info.arg = OpArg(encoded);
203+
info.instr = Instruction::LoadAttr { idx: Arg::marker() };
204+
}
205+
// LOAD_ATTR → encode with method flag=0
206+
Instruction::LoadAttr { idx } => {
207+
let encoded = encode_load_attr_arg(idx.get(info.arg), false);
208+
info.arg = OpArg(encoded);
209+
info.instr = Instruction::LoadAttr { idx: Arg::marker() };
210+
}
211+
// POP_BLOCK pseudo → NOP
212+
Instruction::PopBlock => {
213+
info.instr = Instruction::Nop;
214+
}
215+
_ => {}
216+
}
217+
}
218+
}
219+
192220
let mut block_to_offset = vec![Label(0); blocks.len()];
193221
// block_to_index: maps block idx to instruction index (for exception table)
194222
// This is the index into the final instructions array, including EXTENDED_ARG
@@ -213,23 +241,44 @@ impl CodeInfo {
213241
let mut next_block = BlockIdx(0);
214242
while next_block != BlockIdx::NULL {
215243
let block = &mut blocks[next_block];
244+
// Track current instruction offset for jump direction resolution
245+
let mut current_offset = block_to_offset[next_block.idx()].0;
216246
for info in &mut block.instructions {
217-
let (op, arg, target) = (info.instr, &mut info.arg, info.target);
247+
let target = info.target;
218248
if target != BlockIdx::NULL {
219249
let new_arg = OpArg(block_to_offset[target.idx()].0);
220-
recompile_extended_arg |= new_arg.instr_size() != arg.instr_size();
221-
*arg = new_arg;
250+
recompile_extended_arg |= new_arg.instr_size() != info.arg.instr_size();
251+
info.arg = new_arg;
222252
}
223-
let (extras, lo_arg) = arg.split();
253+
254+
// Convert JUMP pseudo to real instructions (direction depends on offset)
255+
let op = match info.instr {
256+
Instruction::Jump { .. } if target != BlockIdx::NULL => {
257+
let target_offset = block_to_offset[target.idx()].0;
258+
if target_offset > current_offset {
259+
Instruction::JumpForward {
260+
target: Arg::marker(),
261+
}
262+
} else {
263+
Instruction::JumpBackward {
264+
target: Arg::marker(),
265+
}
266+
}
267+
}
268+
other => other,
269+
};
270+
271+
let (extras, lo_arg) = info.arg.split();
224272
locations.extend(core::iter::repeat_n(
225273
(info.location, info.end_location),
226-
arg.instr_size(),
274+
info.arg.instr_size(),
227275
));
228276
instructions.extend(
229277
extras
230278
.map(|byte| CodeUnit::new(Instruction::ExtendedArg, byte))
231279
.chain([CodeUnit { op, arg: lo_arg }]),
232280
);
281+
current_offset += info.arg.instr_size() as u32;
233282
}
234283
next_block = block.next;
235284
}

‎crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap‎

Lines changed: 14 additions & 14 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,20 @@ pub fn find_exception_handler(table: &[u8], offset: u32) -> Option<ExceptionTabl
8484
None
8585
}
8686

87+
/// Encode LOAD_ATTR oparg: bit 0 = method flag, bits 1+ = name index.
88+
#[inline]
89+
pub const fn encode_load_attr_arg(name_idx: u32, is_method: bool) -> u32 {
90+
(name_idx << 1) | (is_method as u32)
91+
}
92+
93+
/// Decode LOAD_ATTR oparg: returns (name_idx, is_method).
94+
#[inline]
95+
pub const fn decode_load_attr_arg(oparg: u32) -> (u32, bool) {
96+
let is_method = (oparg & 1) == 1;
97+
let name_idx = oparg >> 1;
98+
(name_idx, is_method)
99+
}
100+
87101
/// Oparg values for [`Instruction::ConvertValue`].
88102
///
89103
/// ## See also
@@ -1734,6 +1748,9 @@ impl Instruction {
17341748
pub const fn label_arg(&self) -> Option<Arg<Label>> {
17351749
match self {
17361750
Jump { target: l }
1751+
| JumpBackward { target: l }
1752+
| JumpBackwardNoInterrupt { target: l }
1753+
| JumpForward { target: l }
17371754
| JumpIfNotExcMatch(l)
17381755
| PopJumpIfTrue { target: l }
17391756
| PopJumpIfFalse { target: l }
@@ -1760,6 +1777,9 @@ impl Instruction {
17601777
matches!(
17611778
self,
17621779
Jump { .. }
1780+
| JumpForward { .. }
1781+
| JumpBackward { .. }
1782+
| JumpBackwardNoInterrupt { .. }
17631783
| Continue { .. }
17641784
| Break { .. }
17651785
| ReturnValue
@@ -2053,11 +2073,27 @@ impl Instruction {
20532073
ImportName { idx } => w!(IMPORT_NAME, name = idx),
20542074
IsOp(inv) => w!(IS_OP, ?inv),
20552075
Jump { target } => w!(JUMP, target),
2076+
JumpBackward { target } => w!(JUMP_BACKWARD, target),
2077+
JumpBackwardNoInterrupt { target } => w!(JUMP_BACKWARD_NO_INTERRUPT, target),
2078+
JumpForward { target } => w!(JUMP_FORWARD, target),
20562079
JumpIfFalseOrPop { target } => w!(JUMP_IF_FALSE_OR_POP, target),
20572080
JumpIfNotExcMatch(target) => w!(JUMP_IF_NOT_EXC_MATCH, target),
20582081
JumpIfTrueOrPop { target } => w!(JUMP_IF_TRUE_OR_POP, target),
20592082
ListAppend { i } => w!(LIST_APPEND, i),
2060-
LoadAttr { idx } => w!(LOAD_ATTR, name = idx),
2083+
LoadAttr { idx } => {
2084+
let encoded = idx.get(arg);
2085+
let (name_idx, is_method) = decode_load_attr_arg(encoded);
2086+
let attr_name = name(name_idx);
2087+
if is_method {
2088+
write!(
2089+
f,
2090+
"{:pad$}({}, {}, method=true)",
2091+
"LOAD_ATTR", encoded, attr_name
2092+
)
2093+
} else {
2094+
write!(f, "{:pad$}({}, {})", "LOAD_ATTR", encoded, attr_name)
2095+
}
2096+
}
20612097
LoadAttrMethod { idx } => w!(LOAD_ATTR_METHOD, name = idx),
20622098
LoadBuildClass => w!(LOAD_BUILD_CLASS),
20632099
LoadClassDeref(idx) => w!(LOAD_CLASSDEREF, cell_name = idx),

‎crates/jit/src/instructions.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
212212
match instruction {
213213
Instruction::ReturnValue
214214
| Instruction::ReturnConst { .. }
215-
| Instruction::Jump { .. } => {
215+
| Instruction::Jump { .. }
216+
| Instruction::JumpBackward { .. }
217+
| Instruction::JumpBackwardNoInterrupt { .. }
218+
| Instruction::JumpForward { .. } => {
216219
in_unreachable_code = true;
217220
}
218221
_ => {}
@@ -558,7 +561,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
558561
}
559562
Instruction::ExtendedArg => Ok(()),
560563

561-
Instruction::Jump { target } => {
564+
Instruction::Jump { target }
565+
| Instruction::JumpBackward { target }
566+
| Instruction::JumpBackwardNoInterrupt { target }
567+
| Instruction::JumpForward { target } => {
562568
let target_block = self.get_or_create_block(target.get(arg));
563569
self.builder.ins().jump(target_block, &[]);
564570
Ok(())

0 commit comments

Comments
 (0)