Skip to content

Commit 3c6e2cd

Browse files
youknowoneyouknowone
committed
rework type_params scope
1 parent e146de5 commit 3c6e2cd

File tree

1 file changed

+108
-70
lines changed

1 file changed

+108
-70
lines changed

‎compiler/codegen/src/compile.rs

Lines changed: 108 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -2173,103 +2173,97 @@ impl Compiler<'_> {
21732173
) -> CompileResult<()> {
21742174
self.prepare_decorators(decorator_list)?;
21752175

2176-
let prev_ctx = self.ctx;
2177-
self.ctx = CompileContext {
2178-
func: FunctionContext::NoFunction,
2179-
in_class: true,
2180-
loop_data: None,
2181-
};
2176+
let has_type_params = type_params.is_some();
2177+
let firstlineno = self.get_source_line_number().get().to_u32();
2178+
2179+
if has_type_params {
2180+
// For PEP 695 classes, we need to:
2181+
// 1. Enter type params scope
2182+
// 2. Compile type params
2183+
// 3. Call compile_class_body (creates nested scope)
2184+
// 4. Generate class creation code
2185+
// 5. Exit type params scope and wrap everything in a function
2186+
2187+
let type_params_name = format!("<generic parameters of {}>", name);
2188+
self.push_output(
2189+
bytecode::CodeFlags::IS_OPTIMIZED | bytecode::CodeFlags::NEW_LOCALS,
2190+
0,
2191+
0,
2192+
0,
2193+
type_params_name.clone(),
2194+
);
2195+
2196+
// Set private name for name mangling
2197+
self.code_stack.last_mut().unwrap().private = Some(name.to_owned());
21822198

2183-
// If there are type params, we need to push a special symbol table just for them
2184-
if let Some(type_params) = type_params {
2185-
self.push_symbol_table();
2186-
// Save current private name to restore later
2187-
let saved_private = self.code_stack.last().and_then(|info| info.private.clone());
21882199
// Compile type parameters and store as .type_params
2189-
self.compile_type_params(type_params)?;
2190-
// Restore private name after type param scope
2191-
if let Some(private) = saved_private {
2192-
self.code_stack.last_mut().unwrap().private = Some(private);
2193-
}
2200+
self.compile_type_params(type_params.unwrap())?;
21942201
let dot_type_params = self.name(".type_params");
21952202
emit!(self, Instruction::StoreLocal(dot_type_params));
21962203

2197-
// Create .generic_base in the type params scope (like CPython)
2198-
// Load .type_params
2199-
emit!(self, Instruction::LoadNameAny(dot_type_params));
2204+
// Compile the class body (creates its own scope)
2205+
let prev_ctx = self.ctx;
2206+
self.ctx = CompileContext {
2207+
func: FunctionContext::NoFunction,
2208+
in_class: true,
2209+
loop_data: None,
2210+
};
2211+
let class_code = self.compile_class_body(name, body, type_params, firstlineno)?;
2212+
self.ctx = prev_ctx;
22002213

2201-
// Call INTRINSIC_SUBSCRIPT_GENERIC to create Generic[*type_params]
2214+
// Back in type params scope, create .generic_base
2215+
emit!(self, Instruction::LoadNameAny(dot_type_params));
22022216
emit!(
22032217
self,
22042218
Instruction::CallIntrinsic1 {
22052219
func: bytecode::IntrinsicFunction1::SubscriptGeneric
22062220
}
22072221
);
2208-
2209-
// Store as .generic_base in the type params scope
22102222
let dot_generic_base = self.name(".generic_base");
22112223
emit!(self, Instruction::StoreLocal(dot_generic_base));
2212-
}
2213-
2214-
// Compile the class body into a code object
2215-
let firstlineno = self.get_source_line_number().get().to_u32();
2216-
let code = self.compile_class_body(name, body, type_params, firstlineno)?;
2217-
self.ctx = prev_ctx;
22182224

2219-
emit!(self, Instruction::LoadBuildClass);
2225+
// Generate the class creation code (still in type params scope)
2226+
emit!(self, Instruction::LoadBuildClass);
22202227

2221-
let mut func_flags = bytecode::MakeFunctionFlags::empty();
2228+
// Set up the class function
2229+
let mut func_flags = bytecode::MakeFunctionFlags::empty();
22222230

2223-
// Prepare generic type parameters:
2224-
if type_params.is_some() {
2225-
// Load .type_params from the type params scope
2226-
let dot_type_params = self.name(".type_params");
2231+
// Load .type_params for the class function
22272232
emit!(self, Instruction::LoadNameAny(dot_type_params));
22282233
func_flags |= bytecode::MakeFunctionFlags::TYPE_PARAMS;
2229-
}
22302234

2231-
if self.build_closure(&code) {
2232-
func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
2233-
}
2234-
2235-
self.emit_load_const(ConstantData::Code {
2236-
code: Box::new(code),
2237-
});
2238-
self.emit_load_const(ConstantData::Str { value: name.into() });
2235+
if self.build_closure(&class_code) {
2236+
func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
2237+
}
22392238

2240-
// Turn code object into function object:
2241-
emit!(self, Instruction::MakeFunction(func_flags));
2239+
self.emit_load_const(ConstantData::Code {
2240+
code: Box::new(class_code),
2241+
});
2242+
self.emit_load_const(ConstantData::Str { value: name.into() });
22422243

2243-
self.emit_load_const(ConstantData::Str { value: name.into() });
2244+
emit!(self, Instruction::MakeFunction(func_flags));
22442245

2245-
// For PEP 695 classes: handle Generic base creation
2246-
if type_params.is_some() {
2247-
// Stack currently: [function, class_name]
2246+
self.emit_load_const(ConstantData::Str { value: name.into() });
22482247

2249-
// Compile the original bases (if any)
2250-
if let Some(arguments) = arguments {
2251-
// Compile all the positional arguments (bases)
2248+
// Compile bases with .generic_base appended
2249+
let base_count = if let Some(arguments) = arguments {
22522250
for arg in &arguments.args {
22532251
self.compile_expression(arg)?;
22542252
}
2255-
}
2253+
arguments.args.len()
2254+
} else {
2255+
0
2256+
};
22562257

2257-
// Load .generic_base from the type params scope
2258-
let dot_generic_base = self.name(".generic_base");
2258+
// Load .generic_base as the last base
22592259
emit!(self, Instruction::LoadNameAny(dot_generic_base));
22602260

2261-
// Calculate total number of positional args
2262-
let nargs = if let Some(arguments) = arguments {
2263-
2 + arguments.args.len() as u32 + 1 // function, name, bases..., generic_base
2264-
} else {
2265-
3 // function, name, generic_base
2266-
};
2261+
let nargs = 2 + base_count as u32 + 1; // function, name, bases..., generic_base
22672262

2268-
// Handle keyword arguments if any
2263+
// Handle keyword arguments
22692264
if let Some(arguments) = arguments
22702265
&& !arguments.keywords.is_empty()
22712266
{
2272-
// Compile keyword arguments
22732267
for keyword in &arguments.keywords {
22742268
if let Some(name) = &keyword.arg {
22752269
self.emit_load_const(ConstantData::Str {
@@ -2288,8 +2282,58 @@ impl Compiler<'_> {
22882282
} else {
22892283
emit!(self, Instruction::CallFunctionPositional { nargs });
22902284
}
2285+
2286+
// Return the created class
2287+
self.emit_return_value();
2288+
2289+
// Exit type params scope and get the code object
2290+
let type_params_code = self.exit_scope();
2291+
2292+
// Now execute the type params function
2293+
if self.build_closure(&type_params_code) {
2294+
// Should not need closure
2295+
}
2296+
self.emit_load_const(ConstantData::Code {
2297+
code: Box::new(type_params_code),
2298+
});
2299+
self.emit_load_const(ConstantData::Str {
2300+
value: type_params_name.into(),
2301+
});
2302+
emit!(
2303+
self,
2304+
Instruction::MakeFunction(bytecode::MakeFunctionFlags::empty())
2305+
);
2306+
2307+
emit!(self, Instruction::CallFunctionPositional { nargs: 0 });
22912308
} else {
2292-
// No type params, normal compilation
2309+
// Traditional class without type params
2310+
let prev_ctx = self.ctx;
2311+
self.ctx = CompileContext {
2312+
func: FunctionContext::NoFunction,
2313+
in_class: true,
2314+
loop_data: None,
2315+
};
2316+
2317+
let code = self.compile_class_body(name, body, type_params, firstlineno)?;
2318+
self.ctx = prev_ctx;
2319+
2320+
emit!(self, Instruction::LoadBuildClass);
2321+
2322+
let mut func_flags = bytecode::MakeFunctionFlags::empty();
2323+
2324+
if self.build_closure(&code) {
2325+
func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
2326+
}
2327+
2328+
self.emit_load_const(ConstantData::Code {
2329+
code: Box::new(code),
2330+
});
2331+
self.emit_load_const(ConstantData::Str { value: name.into() });
2332+
2333+
emit!(self, Instruction::MakeFunction(func_flags));
2334+
2335+
self.emit_load_const(ConstantData::Str { value: name.into() });
2336+
22932337
let call = if let Some(arguments) = arguments {
22942338
self.compile_call_inner(2, arguments)?
22952339
} else {
@@ -2298,13 +2342,7 @@ impl Compiler<'_> {
22982342
self.compile_normal_call(call);
22992343
}
23002344

2301-
// Pop the special type params symbol table
2302-
if type_params.is_some() {
2303-
self.pop_symbol_table();
2304-
}
2305-
23062345
self.apply_decorators(decorator_list);
2307-
23082346
self.store_name(name)
23092347
}
23102348

0 commit comments

Comments
 (0)