Skip to content

Commit 2a8ee18

Browse files
youknowoneyouknowone
committed
fblock
1 parent bbdee00 commit 2a8ee18

File tree

2 files changed

+159
-21
lines changed

2 files changed

+159
-21
lines changed

‎compiler/codegen/src/compile.rs‎

Lines changed: 137 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,33 @@ use crate::{
1414
symboltable::{self, SymbolFlags, SymbolScope, SymbolTable, SymbolTableType},
1515
unparse::unparse_expr,
1616
};
17+
18+
const MAXBLOCKS: usize = 20;
19+
20+
#[derive(Debug, Clone, Copy)]
21+
pub enum FBlockType {
22+
WhileLoop,
23+
ForLoop,
24+
TryExcept,
25+
FinallyTry,
26+
FinallyEnd,
27+
With,
28+
AsyncWith,
29+
HandlerCleanup,
30+
PopValue,
31+
ExceptionHandler,
32+
ExceptionGroupHandler,
33+
AsyncComprehensionGenerator,
34+
StopIteration,
35+
}
36+
37+
#[derive(Debug, Clone)]
38+
pub struct FBlockInfo {
39+
pub fb_type: FBlockType,
40+
pub fb_block: BlockIdx,
41+
pub fb_exit: BlockIdx,
42+
// fb_datum is not needed in RustPython
43+
}
1744
use itertools::Itertools;
1845
use malachite_bigint::BigInt;
1946
use num_complex::Complex;
@@ -315,14 +342,15 @@ impl<'src> Compiler<'src> {
315342
varnames: IndexSet::default(),
316343
cellvars: IndexSet::default(),
317344
freevars: IndexSet::default(),
318-
fasthidden: IndexMap::default(),
345+
fast_hidden: IndexMap::default(),
319346
argcount: 0,
320347
posonlyargcount: 0,
321348
kwonlyargcount: 0,
322349
firstlineno: OneIndexed::MIN,
323350
},
324351
static_attributes: None,
325352
in_inlined_comp: false,
353+
fblock: Vec::with_capacity(MAXBLOCKS),
326354
};
327355
Compiler {
328356
code_stack: vec![module_code],
@@ -430,7 +458,7 @@ impl Compiler<'_> {
430458
varnames: varname_cache,
431459
cellvars: cellvar_cache,
432460
freevars: freevar_cache,
433-
fasthidden: IndexMap::default(),
461+
fast_hidden: IndexMap::default(),
434462
argcount: arg_count,
435463
posonlyargcount: posonlyarg_count,
436464
kwonlyargcount: kwonlyarg_count,
@@ -442,6 +470,7 @@ impl Compiler<'_> {
442470
None
443471
},
444472
in_inlined_comp: false,
473+
fblock: Vec::with_capacity(MAXBLOCKS),
445474
};
446475
self.code_stack.push(info);
447476
}
@@ -454,6 +483,37 @@ impl Compiler<'_> {
454483
unwrap_internal(self, stack_top.finalize_code(self.opts.optimize))
455484
}
456485

486+
/// Push a new fblock
487+
// = compiler_push_fblock
488+
fn push_fblock(
489+
&mut self,
490+
fb_type: FBlockType,
491+
fb_block: BlockIdx,
492+
fb_exit: BlockIdx,
493+
) -> CompileResult<()> {
494+
let code = self.current_code_info();
495+
if code.fblock.len() >= MAXBLOCKS {
496+
return Err(self.error(CodegenErrorType::SyntaxError(
497+
"too many statically nested blocks".to_owned(),
498+
)));
499+
}
500+
code.fblock.push(FBlockInfo {
501+
fb_type,
502+
fb_block,
503+
fb_exit,
504+
});
505+
Ok(())
506+
}
507+
508+
/// Pop an fblock
509+
// = compiler_pop_fblock
510+
fn pop_fblock(&mut self, _expected_type: FBlockType) -> FBlockInfo {
511+
let code = self.current_code_info();
512+
// TODO: Add assertion to check expected type matches
513+
// assert!(matches!(fblock.fb_type, expected_type));
514+
code.fblock.pop().expect("fblock stack underflow")
515+
}
516+
457517
// could take impl Into<Cow<str>>, but everything is borrowed from ast structs; we never
458518
// actually have a `String` to pass
459519
fn name(&mut self, name: &str) -> bytecode::NameIdx {
@@ -1086,26 +1146,62 @@ impl Compiler<'_> {
10861146
self.switch_to_block(after_block);
10871147
}
10881148
}
1089-
Stmt::Break(_) => match self.ctx.loop_data {
1090-
Some((_, end)) => {
1091-
emit!(self, Instruction::Break { target: end });
1092-
}
1093-
None => {
1094-
return Err(
1095-
self.error_ranged(CodegenErrorType::InvalidBreak, statement.range())
1096-
);
1097-
}
1098-
},
1099-
Stmt::Continue(_) => match self.ctx.loop_data {
1100-
Some((start, _)) => {
1101-
emit!(self, Instruction::Continue { target: start });
1149+
Stmt::Break(_) => {
1150+
// Find the innermost loop in fblock stack
1151+
let found_loop = {
1152+
let code = self.current_code_info();
1153+
let mut result = None;
1154+
for i in (0..code.fblock.len()).rev() {
1155+
match code.fblock[i].fb_type {
1156+
FBlockType::WhileLoop | FBlockType::ForLoop => {
1157+
result = Some(code.fblock[i].fb_exit);
1158+
break;
1159+
}
1160+
_ => continue,
1161+
}
1162+
}
1163+
result
1164+
};
1165+
1166+
match found_loop {
1167+
Some(exit_block) => {
1168+
emit!(self, Instruction::Break { target: exit_block });
1169+
}
1170+
None => {
1171+
return Err(
1172+
self.error_ranged(CodegenErrorType::InvalidBreak, statement.range())
1173+
);
1174+
}
11021175
}
1103-
None => {
1104-
return Err(
1105-
self.error_ranged(CodegenErrorType::InvalidContinue, statement.range())
1106-
);
1176+
}
1177+
Stmt::Continue(_) => {
1178+
// Find the innermost loop in fblock stack
1179+
let found_loop = {
1180+
let code = self.current_code_info();
1181+
let mut result = None;
1182+
for i in (0..code.fblock.len()).rev() {
1183+
match code.fblock[i].fb_type {
1184+
FBlockType::WhileLoop | FBlockType::ForLoop => {
1185+
result = Some(code.fblock[i].fb_block);
1186+
break;
1187+
}
1188+
_ => continue,
1189+
}
1190+
}
1191+
result
1192+
};
1193+
1194+
match found_loop {
1195+
Some(loop_block) => {
1196+
emit!(self, Instruction::Continue { target: loop_block });
1197+
}
1198+
None => {
1199+
return Err(
1200+
self.error_ranged(CodegenErrorType::InvalidContinue, statement.range())
1201+
);
1202+
}
11071203
}
1108-
},
1204+
}
11091205
Stmt::Return(StmtReturn { value, .. }) => {
11101206
if !self.ctx.in_func() {
11111207
return Err(
@@ -2028,6 +2124,9 @@ impl Compiler<'_> {
20282124
emit!(self, Instruction::SetupLoop);
20292125
self.switch_to_block(while_block);
20302126

2127+
// Push fblock for while loop
2128+
self.push_fblock(FBlockType::WhileLoop, while_block, after_block)?;
2129+
20312130
self.compile_jump_if(test, false, else_block)?;
20322131

20332132
let was_in_loop = self.ctx.loop_data.replace((while_block, after_block));
@@ -2040,6 +2139,9 @@ impl Compiler<'_> {
20402139
}
20412140
);
20422141
self.switch_to_block(else_block);
2142+
2143+
// Pop fblock
2144+
self.pop_fblock(FBlockType::WhileLoop);
20432145
emit!(self, Instruction::PopBlock);
20442146
self.compile_statements(orelse)?;
20452147
self.switch_to_block(after_block);
@@ -2150,6 +2252,10 @@ impl Compiler<'_> {
21502252
emit!(self, Instruction::GetAIter);
21512253

21522254
self.switch_to_block(for_block);
2255+
2256+
// Push fblock for async for loop
2257+
self.push_fblock(FBlockType::ForLoop, for_block, after_block)?;
2258+
21532259
emit!(
21542260
self,
21552261
Instruction::SetupExcept {
@@ -2172,6 +2278,10 @@ impl Compiler<'_> {
21722278
emit!(self, Instruction::GetIter);
21732279

21742280
self.switch_to_block(for_block);
2281+
2282+
// Push fblock for for loop
2283+
self.push_fblock(FBlockType::ForLoop, for_block, after_block)?;
2284+
21752285
emit!(self, Instruction::ForIter { target: else_block });
21762286

21772287
// Start of loop iteration, set targets:
@@ -2184,6 +2294,10 @@ impl Compiler<'_> {
21842294
emit!(self, Instruction::Jump { target: for_block });
21852295

21862296
self.switch_to_block(else_block);
2297+
2298+
// Pop fblock
2299+
self.pop_fblock(FBlockType::ForLoop);
2300+
21872301
if is_async {
21882302
emit!(self, Instruction::EndAsyncFor);
21892303
}
@@ -4162,6 +4276,9 @@ impl Compiler<'_> {
41624276
// Create magnificent function <listcomp>:
41634277
self.push_output(flags, 1, 1, 0, name.to_owned());
41644278

4279+
// Mark that we're in an inlined comprehension
4280+
self.current_code_info().in_inlined_comp = true;
4281+
41654282
// Set qualname for comprehension
41664283
self.set_qualname();
41674284

‎compiler/codegen/src/ir.rs‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,24 @@ use ruff_source_file::{OneIndexed, SourceLocation};
66
use rustpython_compiler_core::bytecode::{
77
CodeFlags, CodeObject, CodeUnit, ConstantData, InstrDisplayContext, Instruction, Label, OpArg,
88
};
9+
10+
/// Metadata for a code unit
11+
// = _PyCompile_CodeUnitMetadata
12+
#[derive(Clone, Debug)]
13+
pub struct CodeUnitMetadata {
14+
pub name: String, // u_name (obj_name)
15+
pub qualname: Option<String>, // u_qualname
16+
pub consts: IndexSet<ConstantData>, // u_consts
17+
pub names: IndexSet<String>, // u_names
18+
pub varnames: IndexSet<String>, // u_varnames
19+
pub cellvars: IndexSet<String>, // u_cellvars
20+
pub freevars: IndexSet<String>, // u_freevars
21+
pub fast_hidden: IndexMap<String, bool>, // u_fast_hidden
22+
pub argcount: u32, // u_argcount
23+
pub posonlyargcount: u32, // u_posonlyargcount
24+
pub kwonlyargcount: u32, // u_kwonlyargcount
25+
pub firstlineno: OneIndexed, // u_firstlineno
26+
}
927
// use rustpython_parser_core::source_code::{LineNumber, SourceLocation};
1028

1129
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
@@ -73,14 +91,16 @@ pub struct CodeInfo {
7391
pub blocks: Vec<Block>,
7492
pub current_block: BlockIdx,
7593

76-
// Grouped metadata fields (like CPython's u_metadata)
7794
pub metadata: CodeUnitMetadata,
7895

7996
// For class scopes: attributes accessed via self.X
8097
pub static_attributes: Option<IndexSet<String>>,
8198

8299
// True if compiling an inlined comprehension
83100
pub in_inlined_comp: bool,
101+
102+
// Block stack for tracking nested control structures
103+
pub fblock: Vec<crate::compile::FBlockInfo>,
84104
}
85105
impl CodeInfo {
86106
pub fn finalize_code(mut self, optimize: u8) -> crate::InternalResult<CodeObject> {
@@ -101,6 +121,7 @@ impl CodeInfo {
101121
metadata,
102122
static_attributes: _,
103123
in_inlined_comp: _,
124+
fblock: _,
104125
} = self;
105126

106127
let CodeUnitMetadata {

0 commit comments

Comments
 (0)