Skip to content

Commit eb0d23c

Browse files
youknowoneyouknowone
committed
fblock
1 parent c33b8c9 commit eb0d23c

File tree

2 files changed

+158
-19
lines changed

2 files changed

+158
-19
lines changed

‎compiler/codegen/src/compile.rs‎

Lines changed: 136 additions & 18 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;
@@ -323,6 +350,7 @@ impl<'src> Compiler<'src> {
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],
@@ -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,38 @@ 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+
let fblock = code.fblock.pop().expect("fblock stack underflow");
513+
// TODO: Add assertion to check expected type matches
514+
// assert!(matches!(fblock.fb_type, expected_type));
515+
fblock
516+
}
517+
457518
// could take impl Into<Cow<str>>, but everything is borrowed from ast structs; we never
458519
// actually have a `String` to pass
459520
fn name(&mut self, name: &str) -> bytecode::NameIdx {
@@ -1086,26 +1147,62 @@ impl Compiler<'_> {
10861147
self.switch_to_block(after_block);
10871148
}
10881149
}
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 });
1150+
Stmt::Break(_) => {
1151+
// Find the innermost loop in fblock stack
1152+
let found_loop = {
1153+
let code = self.current_code_info();
1154+
let mut result = None;
1155+
for i in (0..code.fblock.len() as usize).rev() {
1156+
match code.fblock[i].fb_type {
1157+
FBlockType::WhileLoop | FBlockType::ForLoop => {
1158+
result = Some(code.fblock[i].fb_exit);
1159+
break;
1160+
}
1161+
_ => continue,
1162+
}
1163+
}
1164+
result
1165+
};
1166+
1167+
match found_loop {
1168+
Some(exit_block) => {
1169+
emit!(self, Instruction::Break { target: exit_block });
1170+
}
1171+
None => {
1172+
return Err(
1173+
self.error_ranged(CodegenErrorType::InvalidBreak, statement.range())
1174+
);
1175+
}
11021176
}
1103-
None => {
1104-
return Err(
1105-
self.error_ranged(CodegenErrorType::InvalidContinue, statement.range())
1106-
);
1177+
}
1178+
Stmt::Continue(_) => {
1179+
// Find the innermost loop in fblock stack
1180+
let found_loop = {
1181+
let code = self.current_code_info();
1182+
let mut result = None;
1183+
for i in (0..code.fblock.len() as usize).rev() {
1184+
match code.fblock[i].fb_type {
1185+
FBlockType::WhileLoop | FBlockType::ForLoop => {
1186+
result = Some(code.fblock[i].fb_block);
1187+
break;
1188+
}
1189+
_ => continue,
1190+
}
1191+
}
1192+
result
1193+
};
1194+
1195+
match found_loop {
1196+
Some(loop_block) => {
1197+
emit!(self, Instruction::Continue { target: loop_block });
1198+
}
1199+
None => {
1200+
return Err(
1201+
self.error_ranged(CodegenErrorType::InvalidContinue, statement.range())
1202+
);
1203+
}
11071204
}
1108-
},
1205+
}
11091206
Stmt::Return(StmtReturn { value, .. }) => {
11101207
if !self.ctx.in_func() {
11111208
return Err(
@@ -2028,6 +2125,9 @@ impl Compiler<'_> {
20282125
emit!(self, Instruction::SetupLoop);
20292126
self.switch_to_block(while_block);
20302127

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

20332133
let was_in_loop = self.ctx.loop_data.replace((while_block, after_block));
@@ -2040,6 +2140,9 @@ impl Compiler<'_> {
20402140
}
20412141
);
20422142
self.switch_to_block(else_block);
2143+
2144+
// Pop fblock
2145+
self.pop_fblock(FBlockType::WhileLoop);
20432146
emit!(self, Instruction::PopBlock);
20442147
self.compile_statements(orelse)?;
20452148
self.switch_to_block(after_block);
@@ -2150,6 +2253,10 @@ impl Compiler<'_> {
21502253
emit!(self, Instruction::GetAIter);
21512254

21522255
self.switch_to_block(for_block);
2256+
2257+
// Push fblock for async for loop
2258+
self.push_fblock(FBlockType::ForLoop, for_block, after_block)?;
2259+
21532260
emit!(
21542261
self,
21552262
Instruction::SetupExcept {
@@ -2172,6 +2279,10 @@ impl Compiler<'_> {
21722279
emit!(self, Instruction::GetIter);
21732280

21742281
self.switch_to_block(for_block);
2282+
2283+
// Push fblock for for loop
2284+
self.push_fblock(FBlockType::ForLoop, for_block, after_block)?;
2285+
21752286
emit!(self, Instruction::ForIter { target: else_block });
21762287

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

21862297
self.switch_to_block(else_block);
2298+
2299+
// Pop fblock
2300+
self.pop_fblock(FBlockType::ForLoop);
2301+
21872302
if is_async {
21882303
emit!(self, Instruction::EndAsyncFor);
21892304
}
@@ -4162,6 +4277,9 @@ impl Compiler<'_> {
41624277
// Create magnificent function <listcomp>:
41634278
self.push_output(flags, 1, 1, 0, name.to_owned());
41644279

4280+
// Mark that we're in an inlined comprehension
4281+
self.current_code_info().in_inlined_comp = true;
4282+
41654283
// Set qualname for comprehension
41664284
self.set_qualname();
41674285

‎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 fasthidden: IndexMap<String, bool>, // u_fasthidden
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)