Skip to content

Commit 3ef0cfc

Browse files
youknowoneyouknowone
authored
compiler enter_scope (#5950)
* enter_scope * drop_class_free * push_output based on enter_scope
1 parent 1303ace commit 3ef0cfc

File tree

2 files changed

+219
-50
lines changed

2 files changed

+219
-50
lines changed

‎compiler/codegen/src/compile.rs

Lines changed: 181 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,7 @@ pub fn compile_program(
176176
.map_err(|e| e.into_codegen_error(source_code.path.to_owned()))?;
177177
let mut compiler = Compiler::new(opts, source_code, "<module>".to_owned());
178178
compiler.compile_program(ast, symbol_table)?;
179-
let code = compiler.pop_code_object();
179+
let code = compiler.exit_scope();
180180
trace!("Compilation completed: {code:?}");
181181
Ok(code)
182182
}
@@ -191,7 +191,7 @@ pub fn compile_program_single(
191191
.map_err(|e| e.into_codegen_error(source_code.path.to_owned()))?;
192192
let mut compiler = Compiler::new(opts, source_code, "<module>".to_owned());
193193
compiler.compile_program_single(&ast.body, symbol_table)?;
194-
let code = compiler.pop_code_object();
194+
let code = compiler.exit_scope();
195195
trace!("Compilation completed: {code:?}");
196196
Ok(code)
197197
}
@@ -205,7 +205,7 @@ pub fn compile_block_expression(
205205
.map_err(|e| e.into_codegen_error(source_code.path.to_owned()))?;
206206
let mut compiler = Compiler::new(opts, source_code, "<module>".to_owned());
207207
compiler.compile_block_expr(&ast.body, symbol_table)?;
208-
let code = compiler.pop_code_object();
208+
let code = compiler.exit_scope();
209209
trace!("Compilation completed: {code:?}");
210210
Ok(code)
211211
}
@@ -219,7 +219,7 @@ pub fn compile_expression(
219219
.map_err(|e| e.into_codegen_error(source_code.path.to_owned()))?;
220220
let mut compiler = Compiler::new(opts, source_code, "<module>".to_owned());
221221
compiler.compile_eval(ast, symbol_table)?;
222-
let code = compiler.pop_code_object();
222+
let code = compiler.exit_scope();
223223
Ok(code)
224224
}
225225

@@ -404,55 +404,121 @@ impl Compiler<'_> {
404404
self.symbol_table_stack.pop().expect("compiler bug")
405405
}
406406

407-
fn push_output(
407+
/// Enter a new scope
408+
// = compiler_enter_scope
409+
fn enter_scope(
408410
&mut self,
409-
flags: bytecode::CodeFlags,
410-
posonlyarg_count: u32,
411-
arg_count: u32,
412-
kwonlyarg_count: u32,
413-
obj_name: String,
414-
) {
411+
name: &str,
412+
scope_type: SymbolTableType,
413+
key: usize, // In RustPython, we use the index in symbol_table_stack as key
414+
lineno: u32,
415+
) -> CompileResult<()> {
416+
// Create location
417+
let location = ruff_source_file::SourceLocation {
418+
row: OneIndexed::new(lineno as usize).unwrap_or(OneIndexed::MIN),
419+
column: OneIndexed::new(1).unwrap(),
420+
};
421+
422+
// Allocate a new compiler unit
423+
424+
// In Rust, we'll create the structure directly
415425
let source_path = self.source_code.path.to_owned();
416-
let first_line_number = self.get_source_line_number();
417426

418-
// Get the private name from current scope if exists
419-
let private = self.code_stack.last().and_then(|info| info.private.clone());
427+
// Lookup symbol table entry using key (_PySymtable_Lookup)
428+
let ste = if key < self.symbol_table_stack.len() {
429+
&self.symbol_table_stack[key]
430+
} else {
431+
return Err(self.error(CodegenErrorType::SyntaxError(
432+
"unknown symbol table entry".to_owned(),
433+
)));
434+
};
420435

421-
let table = self.push_symbol_table();
436+
// Use varnames from symbol table (already collected in definition order)
437+
let varname_cache: IndexSet<String> = ste.varnames.iter().cloned().collect();
422438

423-
let cellvar_cache = table
439+
// Build cellvars using dictbytype (CELL scope, sorted)
440+
let mut cellvar_cache = IndexSet::default();
441+
let mut cell_names: Vec<_> = ste
424442
.symbols
425443
.iter()
426444
.filter(|(_, s)| s.scope == SymbolScope::Cell)
427-
.map(|(var, _)| var.clone())
445+
.map(|(name, _)| name.clone())
428446
.collect();
429-
let freevar_cache = table
447+
cell_names.sort();
448+
for name in cell_names {
449+
cellvar_cache.insert(name);
450+
}
451+
452+
// Handle implicit __class__ cell if needed
453+
if ste.needs_class_closure {
454+
// Cook up an implicit __class__ cell
455+
debug_assert_eq!(scope_type, SymbolTableType::Class);
456+
cellvar_cache.insert("__class__".to_string());
457+
}
458+
459+
// Handle implicit __classdict__ cell if needed
460+
if ste.needs_classdict {
461+
// Cook up an implicit __classdict__ cell
462+
debug_assert_eq!(scope_type, SymbolTableType::Class);
463+
cellvar_cache.insert("__classdict__".to_string());
464+
}
465+
466+
// Build freevars using dictbytype (FREE scope, offset by cellvars size)
467+
let mut freevar_cache = IndexSet::default();
468+
let mut free_names: Vec<_> = ste
430469
.symbols
431470
.iter()
432471
.filter(|(_, s)| {
433472
s.scope == SymbolScope::Free || s.flags.contains(SymbolFlags::FREE_CLASS)
434473
})
435-
.map(|(var, _)| var.clone())
474+
.map(|(name, _)| name.clone())
436475
.collect();
476+
free_names.sort();
477+
for name in free_names {
478+
freevar_cache.insert(name);
479+
}
480+
481+
// Initialize u_metadata fields
482+
let (flags, posonlyarg_count, arg_count, kwonlyarg_count) = match scope_type {
483+
SymbolTableType::Module => (bytecode::CodeFlags::empty(), 0, 0, 0),
484+
SymbolTableType::Class => (bytecode::CodeFlags::empty(), 0, 0, 0),
485+
SymbolTableType::Function | SymbolTableType::Lambda => (
486+
bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
487+
0, // Will be set later in enter_function
488+
0, // Will be set later in enter_function
489+
0, // Will be set later in enter_function
490+
),
491+
SymbolTableType::Comprehension => (
492+
bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
493+
0,
494+
1, // comprehensions take one argument (.0)
495+
0,
496+
),
497+
SymbolTableType::TypeParams => (
498+
bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
499+
0,
500+
0,
501+
0,
502+
),
503+
};
437504

438-
// Initialize varname_cache from SymbolTable::varnames
439-
let varname_cache: IndexSet<String> = table.varnames.iter().cloned().collect();
440-
441-
// Qualname will be set later by set_qualname
442-
let qualname = None;
443-
444-
// Check if this is a class scope
445-
let is_class_scope = table.typ == SymbolTableType::Class;
505+
// Get private name from parent scope
506+
let private = if !self.code_stack.is_empty() {
507+
self.code_stack.last().unwrap().private.clone()
508+
} else {
509+
None
510+
};
446511

447-
let info = ir::CodeInfo {
512+
// Create the new compilation unit
513+
let code_info = ir::CodeInfo {
448514
flags,
449-
source_path,
515+
source_path: source_path.clone(),
450516
private,
451517
blocks: vec![ir::Block::default()],
452-
current_block: ir::BlockIdx(0),
518+
current_block: BlockIdx(0),
453519
metadata: ir::CodeUnitMetadata {
454-
name: obj_name,
455-
qualname,
520+
name: name.to_owned(),
521+
qualname: None, // Will be set below
456522
consts: IndexSet::default(),
457523
names: IndexSet::default(),
458524
varnames: varname_cache,
@@ -462,20 +528,93 @@ impl Compiler<'_> {
462528
argcount: arg_count,
463529
posonlyargcount: posonlyarg_count,
464530
kwonlyargcount: kwonlyarg_count,
465-
firstlineno: first_line_number,
531+
firstlineno: OneIndexed::new(lineno as usize).unwrap_or(OneIndexed::MIN),
466532
},
467-
static_attributes: if is_class_scope {
533+
static_attributes: if scope_type == SymbolTableType::Class {
468534
Some(IndexSet::default())
469535
} else {
470536
None
471537
},
472538
in_inlined_comp: false,
473539
fblock: Vec::with_capacity(MAXBLOCKS),
474540
};
475-
self.code_stack.push(info);
541+
542+
// Push the old compiler unit on the stack (like PyCapsule)
543+
// This happens before setting qualname
544+
self.code_stack.push(code_info);
545+
546+
// Set qualname after pushing (uses compiler_set_qualname logic)
547+
if scope_type != SymbolTableType::Module {
548+
self.set_qualname();
549+
}
550+
551+
// Emit RESUME instruction
552+
let _resume_loc = if scope_type == SymbolTableType::Module {
553+
// Module scope starts with lineno 0
554+
ruff_source_file::SourceLocation {
555+
row: OneIndexed::MIN,
556+
column: OneIndexed::MIN,
557+
}
558+
} else {
559+
location
560+
};
561+
562+
// Set the source range for the RESUME instruction
563+
// For now, just use an empty range at the beginning
564+
self.current_source_range = TextRange::default();
565+
emit!(
566+
self,
567+
Instruction::Resume {
568+
arg: bytecode::ResumeType::AtFuncStart as u32
569+
}
570+
);
571+
572+
if scope_type == SymbolTableType::Module {
573+
// This would be loc.lineno = -1 in CPython
574+
// We handle this differently in RustPython
575+
}
576+
577+
Ok(())
578+
}
579+
580+
fn push_output(
581+
&mut self,
582+
flags: bytecode::CodeFlags,
583+
posonlyarg_count: u32,
584+
arg_count: u32,
585+
kwonlyarg_count: u32,
586+
obj_name: String,
587+
) {
588+
// First push the symbol table
589+
let table = self.push_symbol_table();
590+
let scope_type = table.typ;
591+
592+
// The key is the current position in the symbol table stack
593+
let key = self.symbol_table_stack.len() - 1;
594+
595+
// Get the line number
596+
let lineno = self.get_source_line_number().get();
597+
598+
// Call enter_scope which does most of the work
599+
if let Err(e) = self.enter_scope(&obj_name, scope_type, key, lineno.to_u32()) {
600+
// In the current implementation, push_output doesn't return an error,
601+
// so we panic here. This maintains the same behavior.
602+
panic!("enter_scope failed: {e:?}");
603+
}
604+
605+
// Override the values that push_output sets explicitly
606+
// enter_scope sets default values based on scope_type, but push_output
607+
// allows callers to specify exact values
608+
if let Some(info) = self.code_stack.last_mut() {
609+
info.flags = flags;
610+
info.metadata.argcount = arg_count;
611+
info.metadata.posonlyargcount = posonlyarg_count;
612+
info.metadata.kwonlyargcount = kwonlyarg_count;
613+
}
476614
}
477615

478-
fn pop_code_object(&mut self) -> CodeObject {
616+
// compiler_exit_scope
617+
fn exit_scope(&mut self) -> CodeObject {
479618
let table = self.pop_symbol_table();
480619
assert!(table.sub_tables.is_empty());
481620
let pop = self.code_stack.pop();
@@ -755,7 +894,7 @@ impl Compiler<'_> {
755894
}
756895

757896
fn mangle<'a>(&self, name: &'a str) -> Cow<'a, str> {
758-
// Use u_private from current code unit for name mangling
897+
// Use private from current code unit for name mangling
759898
let private = self
760899
.code_stack
761900
.last()
@@ -1758,14 +1897,6 @@ impl Compiler<'_> {
17581897
.consts
17591898
.insert_full(ConstantData::None);
17601899

1761-
// Emit RESUME instruction at function start
1762-
emit!(
1763-
self,
1764-
Instruction::Resume {
1765-
arg: bytecode::ResumeType::AtFuncStart as u32
1766-
}
1767-
);
1768-
17691900
self.compile_statements(body)?;
17701901

17711902
// Emit None at end:
@@ -1778,7 +1909,7 @@ impl Compiler<'_> {
17781909
}
17791910
}
17801911

1781-
let code = self.pop_code_object();
1912+
let code = self.exit_scope();
17821913
self.ctx = prev_ctx;
17831914

17841915
// Prepare generic type parameters:
@@ -2030,7 +2161,7 @@ impl Compiler<'_> {
20302161

20312162
self.emit_return_value();
20322163

2033-
let code = self.pop_code_object();
2164+
let code = self.exit_scope();
20342165
self.ctx = prev_ctx;
20352166

20362167
emit!(self, Instruction::LoadBuildClass);
@@ -3820,7 +3951,7 @@ impl Compiler<'_> {
38203951

38213952
self.compile_expression(body)?;
38223953
self.emit_return_value();
3823-
let code = self.pop_code_object();
3954+
let code = self.exit_scope();
38243955
if self.build_closure(&code) {
38253956
func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
38263957
}
@@ -4369,7 +4500,7 @@ impl Compiler<'_> {
43694500
self.emit_return_value();
43704501

43714502
// Fetch code for listcomp function:
4372-
let code = self.pop_code_object();
4503+
let code = self.exit_scope();
43734504

43744505
self.ctx = prev_ctx;
43754506

@@ -5076,7 +5207,7 @@ mod tests {
50765207
.unwrap();
50775208
let mut compiler = Compiler::new(opts, source_code, "<module>".to_owned());
50785209
compiler.compile_program(&ast, symbol_table).unwrap();
5079-
compiler.pop_code_object()
5210+
compiler.exit_scope()
50805211
}
50815212

50825213
macro_rules! assert_dis_snapshot {

0 commit comments

Comments
 (0)