Skip to content

Commit 4894ab4

Browse files
youknowoneyouknowone
committed
Add SymbolUsage::TypeParams
1 parent 18d7c1b commit 4894ab4

File tree

2 files changed

+69
-24
lines changed

2 files changed

+69
-24
lines changed

‎compiler/codegen/src/compile.rs‎

Lines changed: 52 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use crate::{
1111
IndexSet, ToPythonName,
1212
error::{CodegenError, CodegenErrorType, PatternUnreachableReason},
1313
ir::{self, BlockIdx},
14-
symboltable::{self, SymbolFlags, SymbolScope, SymbolTable},
14+
symboltable::{self, SymbolFlags, SymbolScope, SymbolTable, SymbolTableType},
1515
unparse::unparse_expr,
1616
};
1717
use itertools::Itertools;
@@ -385,6 +385,30 @@ impl Compiler<'_> {
385385
let source_path = self.source_code.path.to_owned();
386386
let first_line_number = self.get_source_line_number();
387387

388+
// Check if this is a class declared as global in parent scope before pushing symbol table
389+
let is_global_class =
390+
if !self.symbol_table_stack.is_empty() && !self.qualified_path.is_empty() {
391+
// Look at the next table to be pushed (not yet removed from parent)
392+
let parent_table = self.symbol_table_stack.last().unwrap();
393+
let next_table = parent_table.sub_tables.first();
394+
395+
if let Some(next_table) = next_table {
396+
if next_table.typ == SymbolTableType::Class {
397+
// Check if the class is global in parent scope
398+
parent_table
399+
.lookup(&obj_name)
400+
.map(|symbol| symbol.is_global())
401+
.unwrap_or(false)
402+
} else {
403+
false
404+
}
405+
} else {
406+
false
407+
}
408+
} else {
409+
false
410+
};
411+
388412
let table = self.push_symbol_table();
389413

390414
let cellvar_cache = table
@@ -405,6 +429,9 @@ impl Compiler<'_> {
405429
// Calculate qualname based on the current qualified path
406430
let qualname = if self.qualified_path.is_empty() {
407431
Some(obj_name.clone())
432+
} else if is_global_class {
433+
// If the class is global, use just the class name as qualname
434+
Some(obj_name.clone())
408435
} else {
409436
Some(self.qualified_path.join("."))
410437
};
@@ -638,7 +665,11 @@ impl Compiler<'_> {
638665
cache = &mut info.cellvar_cache;
639666
NameOpType::Deref
640667
} // TODO: is this right?
641-
// SymbolScope::Unknown => NameOpType::Global,
668+
SymbolScope::TypeParams => {
669+
// Type parameters are always cell variables
670+
cache = &mut info.cellvar_cache;
671+
NameOpType::Deref
672+
} // SymbolScope::Unknown => NameOpType::Global,
642673
};
643674

644675
if NameUsage::Load == usage && name == "__debug__" {
@@ -1630,6 +1661,7 @@ impl Compiler<'_> {
16301661
let vars = match symbol.scope {
16311662
SymbolScope::Free => &parent_code.freevar_cache,
16321663
SymbolScope::Cell => &parent_code.cellvar_cache,
1664+
SymbolScope::TypeParams => &parent_code.cellvar_cache,
16331665
_ if symbol.flags.contains(SymbolFlags::FREE_CLASS) => &parent_code.freevar_cache,
16341666
x => unreachable!(
16351667
"var {} in a {:?} should be free or cell but it's {:?}",
@@ -1706,20 +1738,7 @@ impl Compiler<'_> {
17061738

17071739
let prev_class_name = self.class_name.replace(name.to_owned());
17081740

1709-
// Check if the class is declared global
1710-
let symbol_table = self.symbol_table_stack.last().unwrap();
1711-
let symbol = unwrap_internal(
1712-
self,
1713-
symbol_table
1714-
.lookup(name.as_ref())
1715-
.ok_or_else(|| InternalError::MissingSymbol(name.to_owned())),
1716-
);
1717-
let mut global_path_prefix = Vec::new();
1718-
if symbol.scope == SymbolScope::GlobalExplicit {
1719-
global_path_prefix.append(&mut self.qualified_path);
1720-
}
17211741
self.push_qualified_path(name);
1722-
let qualified_name = self.qualified_path.join(".");
17231742

17241743
// If there are type params, we need to push a special symbol table just for them
17251744
if let Some(type_params) = type_params {
@@ -1732,17 +1751,23 @@ impl Compiler<'_> {
17321751

17331752
self.push_output(bytecode::CodeFlags::empty(), 0, 0, 0, name.to_owned());
17341753

1735-
// Update the qualname in the current code info
1736-
self.code_stack.last_mut().unwrap().qualname = Some(qualified_name.clone());
1737-
17381754
let (doc_str, body) = split_doc(body, &self.opts);
17391755

17401756
let dunder_name = self.name("__name__");
17411757
emit!(self, Instruction::LoadGlobal(dunder_name));
17421758
let dunder_module = self.name("__module__");
17431759
emit!(self, Instruction::StoreLocal(dunder_module));
1760+
1761+
// Get the qualname from the current code info (set by push_output)
1762+
let current_qualname = self
1763+
.code_stack
1764+
.last()
1765+
.unwrap()
1766+
.qualname
1767+
.clone()
1768+
.unwrap_or_else(|| name.to_owned());
17441769
self.emit_load_const(ConstantData::Str {
1745-
value: qualified_name.into(),
1770+
value: current_qualname.into(),
17461771
});
17471772
let qualname = self.name("__qualname__");
17481773
emit!(self, Instruction::StoreLocal(qualname));
@@ -1790,7 +1815,7 @@ impl Compiler<'_> {
17901815

17911816
self.class_name = prev_class_name;
17921817
self.qualified_path.pop();
1793-
self.qualified_path.append(global_path_prefix.as_mut());
1818+
// self.qualified_path.append(global_path_prefix.as_mut());
17941819
self.ctx = prev_ctx;
17951820

17961821
emit!(self, Instruction::LoadBuildClass);
@@ -3529,10 +3554,16 @@ impl Compiler<'_> {
35293554
if self.build_closure(&code) {
35303555
func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
35313556
}
3557+
3558+
// Get the qualname from the code object
3559+
let qualname = code.qualname.clone();
3560+
35323561
self.emit_load_const(ConstantData::Code {
35333562
code: Box::new(code),
35343563
});
3535-
self.emit_load_const(ConstantData::Str { value: name.into() });
3564+
self.emit_load_const(ConstantData::Str {
3565+
value: qualname.into(),
3566+
});
35363567
// Turn code object into function object:
35373568
emit!(self, Instruction::MakeFunction(func_flags));
35383569

‎compiler/codegen/src/symboltable.rs‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ pub enum SymbolScope {
113113
GlobalImplicit,
114114
Free,
115115
Cell,
116+
TypeParams,
116117
}
117118

118119
bitflags! {
@@ -359,6 +360,10 @@ impl SymbolTableAnalyzer {
359360
SymbolScope::Local | SymbolScope::Cell => {
360361
// all is well
361362
}
363+
SymbolScope::TypeParams => {
364+
// Type parameters are always cell variables in their scope
365+
symbol.scope = SymbolScope::Cell;
366+
}
362367
SymbolScope::Unknown => {
363368
// Try hard to figure out what the scope of this symbol is.
364369
let scope = if symbol.is_bound() {
@@ -557,6 +562,7 @@ enum SymbolUsage {
557562
AnnotationParameter,
558563
AssignedNamedExprInComprehension,
559564
Iter,
565+
TypeParam,
560566
}
561567

562568
struct SymbolTableBuilder<'src> {
@@ -1267,6 +1273,9 @@ impl SymbolTableBuilder<'_> {
12671273
}
12681274

12691275
fn scan_type_params(&mut self, type_params: &TypeParams) -> SymbolTableResult {
1276+
// Register .type_params as a type parameter (automatically becomes cell variable)
1277+
self.register_name(".type_params", SymbolUsage::TypeParam, type_params.range)?;
1278+
12701279
// First register all type parameters
12711280
for type_param in &type_params.type_params {
12721281
match type_param {
@@ -1276,7 +1285,7 @@ impl SymbolTableBuilder<'_> {
12761285
range: type_var_range,
12771286
..
12781287
}) => {
1279-
self.register_name(name.as_str(), SymbolUsage::Assigned, *type_var_range)?;
1288+
self.register_name(name.as_str(), SymbolUsage::TypeParam, *type_var_range)?;
12801289
if let Some(binding) = bound {
12811290
self.scan_expression(binding, ExpressionContext::Load)?;
12821291
}
@@ -1286,14 +1295,14 @@ impl SymbolTableBuilder<'_> {
12861295
range: param_spec_range,
12871296
..
12881297
}) => {
1289-
self.register_name(name, SymbolUsage::Assigned, *param_spec_range)?;
1298+
self.register_name(name, SymbolUsage::TypeParam, *param_spec_range)?;
12901299
}
12911300
TypeParam::TypeVarTuple(TypeParamTypeVarTuple {
12921301
name,
12931302
range: type_var_tuple_range,
12941303
..
12951304
}) => {
1296-
self.register_name(name, SymbolUsage::Assigned, *type_var_tuple_range)?;
1305+
self.register_name(name, SymbolUsage::TypeParam, *type_var_tuple_range)?;
12971306
}
12981307
}
12991308
}
@@ -1544,6 +1553,11 @@ impl SymbolTableBuilder<'_> {
15441553
SymbolUsage::Iter => {
15451554
flags.insert(SymbolFlags::ITER);
15461555
}
1556+
SymbolUsage::TypeParam => {
1557+
// Type parameters are always cell variables in their scope
1558+
symbol.scope = SymbolScope::Cell;
1559+
flags.insert(SymbolFlags::ASSIGNED);
1560+
}
15471561
}
15481562

15491563
// and even more checking

0 commit comments

Comments
 (0)