Skip to content

Commit ac08f44

Browse files
dchiquitodchiquito
authored
Add 3.12 typing features to the compiler (#5302)
1 parent d243c90 commit ac08f44

File tree

7 files changed

+440
-19
lines changed

7 files changed

+440
-19
lines changed

‎compiler/codegen/src/compile.rs

Lines changed: 133 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,26 @@ impl Compiler {
251251
}
252252
}
253253

254+
/// Push the next symbol table on to the stack
255+
fn push_symbol_table(&mut self) -> &SymbolTable {
256+
// Look up the next table contained in the scope of the current table
257+
let table = self
258+
.symbol_table_stack
259+
.last_mut()
260+
.expect("no next symbol table")
261+
.sub_tables
262+
.remove(0);
263+
// Push the next table onto the stack
264+
let last_idx = self.symbol_table_stack.len();
265+
self.symbol_table_stack.push(table);
266+
&self.symbol_table_stack[last_idx]
267+
}
268+
269+
/// Pop the current symbol table off the stack
270+
fn pop_symbol_table(&mut self) -> SymbolTable {
271+
self.symbol_table_stack.pop().expect("compiler bug")
272+
}
273+
254274
fn push_output(
255275
&mut self,
256276
flags: bytecode::CodeFlags,
@@ -262,12 +282,7 @@ impl Compiler {
262282
let source_path = self.source_path.clone();
263283
let first_line_number = self.get_source_line_number();
264284

265-
let table = self
266-
.symbol_table_stack
267-
.last_mut()
268-
.unwrap()
269-
.sub_tables
270-
.remove(0);
285+
let table = self.push_symbol_table();
271286

272287
let cellvar_cache = table
273288
.symbols
@@ -284,8 +299,6 @@ impl Compiler {
284299
.map(|(var, _)| var.clone())
285300
.collect();
286301

287-
self.symbol_table_stack.push(table);
288-
289302
let info = ir::CodeInfo {
290303
flags,
291304
posonlyarg_count,
@@ -307,7 +320,7 @@ impl Compiler {
307320
}
308321

309322
fn pop_code_object(&mut self) -> CodeObject {
310-
let table = self.symbol_table_stack.pop().unwrap();
323+
let table = self.pop_symbol_table();
311324
assert!(table.sub_tables.is_empty());
312325
self.code_stack
313326
.pop()
@@ -752,6 +765,7 @@ impl Compiler {
752765
body,
753766
decorator_list,
754767
returns,
768+
type_params,
755769
..
756770
}) => self.compile_function_def(
757771
name.as_str(),
@@ -760,13 +774,15 @@ impl Compiler {
760774
decorator_list,
761775
returns.as_deref(),
762776
false,
777+
type_params,
763778
)?,
764779
Stmt::AsyncFunctionDef(StmtAsyncFunctionDef {
765780
name,
766781
args,
767782
body,
768783
decorator_list,
769784
returns,
785+
type_params,
770786
..
771787
}) => self.compile_function_def(
772788
name.as_str(),
@@ -775,15 +791,24 @@ impl Compiler {
775791
decorator_list,
776792
returns.as_deref(),
777793
true,
794+
type_params,
778795
)?,
779796
Stmt::ClassDef(StmtClassDef {
780797
name,
781798
body,
782799
bases,
783800
keywords,
784801
decorator_list,
802+
type_params,
785803
..
786-
}) => self.compile_class_def(name.as_str(), body, bases, keywords, decorator_list)?,
804+
}) => self.compile_class_def(
805+
name.as_str(),
806+
body,
807+
bases,
808+
keywords,
809+
decorator_list,
810+
type_params,
811+
)?,
787812
Stmt::Assert(StmtAssert { test, msg, .. }) => {
788813
// if some flag, ignore all assert statements!
789814
if self.opts.optimize == 0 {
@@ -885,7 +910,27 @@ impl Compiler {
885910
Stmt::Pass(_) => {
886911
// No need to emit any code here :)
887912
}
888-
Stmt::TypeAlias(_) => {}
913+
Stmt::TypeAlias(StmtTypeAlias {
914+
name,
915+
type_params,
916+
value,
917+
..
918+
}) => {
919+
let name_string = name.to_string();
920+
if !type_params.is_empty() {
921+
self.push_symbol_table();
922+
}
923+
self.compile_expression(value)?;
924+
self.compile_type_params(type_params)?;
925+
if !type_params.is_empty() {
926+
self.pop_symbol_table();
927+
}
928+
self.emit_load_const(ConstantData::Str {
929+
value: name_string.clone(),
930+
});
931+
emit!(self, Instruction::TypeAlias);
932+
self.store_name(&name_string)?;
933+
}
889934
}
890935
Ok(())
891936
}
@@ -1005,6 +1050,47 @@ impl Compiler {
10051050
}
10061051
}
10071052

1053+
/// Store each type parameter so it is accessible to the current scope, and leave a tuple of
1054+
/// all the type parameters on the stack.
1055+
fn compile_type_params(&mut self, type_params: &[located_ast::TypeParam]) -> CompileResult<()> {
1056+
for type_param in type_params {
1057+
match type_param {
1058+
located_ast::TypeParam::TypeVar(located_ast::TypeParamTypeVar {
1059+
name,
1060+
bound,
1061+
..
1062+
}) => {
1063+
if let Some(expr) = &bound {
1064+
self.compile_expression(expr)?;
1065+
self.emit_load_const(ConstantData::Str {
1066+
value: name.to_string(),
1067+
});
1068+
emit!(self, Instruction::TypeVarWithBound);
1069+
emit!(self, Instruction::Duplicate);
1070+
self.store_name(name.as_ref())?;
1071+
} else {
1072+
// self.store_name(type_name.as_str())?;
1073+
self.emit_load_const(ConstantData::Str {
1074+
value: name.to_string(),
1075+
});
1076+
emit!(self, Instruction::TypeVar);
1077+
emit!(self, Instruction::Duplicate);
1078+
self.store_name(name.as_ref())?;
1079+
}
1080+
}
1081+
located_ast::TypeParam::ParamSpec(_) => todo!(),
1082+
located_ast::TypeParam::TypeVarTuple(_) => todo!(),
1083+
};
1084+
}
1085+
emit!(
1086+
self,
1087+
Instruction::BuildTuple {
1088+
size: u32::try_from(type_params.len()).unwrap(),
1089+
}
1090+
);
1091+
Ok(())
1092+
}
1093+
10081094
fn compile_try_statement(
10091095
&mut self,
10101096
body: &[located_ast::Stmt],
@@ -1151,6 +1237,7 @@ impl Compiler {
11511237
is_forbidden_name(name)
11521238
}
11531239

1240+
#[allow(clippy::too_many_arguments)]
11541241
fn compile_function_def(
11551242
&mut self,
11561243
name: &str,
@@ -1159,10 +1246,15 @@ impl Compiler {
11591246
decorator_list: &[located_ast::Expr],
11601247
returns: Option<&located_ast::Expr>, // TODO: use type hint somehow..
11611248
is_async: bool,
1249+
type_params: &[located_ast::TypeParam],
11621250
) -> CompileResult<()> {
1163-
// Create bytecode for this function:
1164-
11651251
self.prepare_decorators(decorator_list)?;
1252+
1253+
// If there are type params, we need to push a special symbol table just for them
1254+
if !type_params.is_empty() {
1255+
self.push_symbol_table();
1256+
}
1257+
11661258
let mut func_flags = self.enter_function(name, args)?;
11671259
self.current_code_info()
11681260
.flags
@@ -1208,6 +1300,12 @@ impl Compiler {
12081300
self.qualified_path.pop();
12091301
self.ctx = prev_ctx;
12101302

1303+
// Prepare generic type parameters:
1304+
if !type_params.is_empty() {
1305+
self.compile_type_params(type_params)?;
1306+
func_flags |= bytecode::MakeFunctionFlags::TYPE_PARAMS;
1307+
}
1308+
12111309
// Prepare type annotations:
12121310
let mut num_annotations = 0;
12131311

@@ -1253,6 +1351,11 @@ impl Compiler {
12531351
func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
12541352
}
12551353

1354+
// Pop the special type params symbol table
1355+
if !type_params.is_empty() {
1356+
self.pop_symbol_table();
1357+
}
1358+
12561359
self.emit_load_const(ConstantData::Code {
12571360
code: Box::new(code),
12581361
});
@@ -1352,6 +1455,7 @@ impl Compiler {
13521455
bases: &[located_ast::Expr],
13531456
keywords: &[located_ast::Keyword],
13541457
decorator_list: &[located_ast::Expr],
1458+
type_params: &[located_ast::TypeParam],
13551459
) -> CompileResult<()> {
13561460
self.prepare_decorators(decorator_list)?;
13571461

@@ -1378,6 +1482,11 @@ impl Compiler {
13781482
self.push_qualified_path(name);
13791483
let qualified_name = self.qualified_path.join(".");
13801484

1485+
// If there are type params, we need to push a special symbol table just for them
1486+
if !type_params.is_empty() {
1487+
self.push_symbol_table();
1488+
}
1489+
13811490
self.push_output(bytecode::CodeFlags::empty(), 0, 0, 0, name.to_owned());
13821491

13831492
let (doc_str, body) = split_doc(body, &self.opts);
@@ -1428,10 +1537,21 @@ impl Compiler {
14281537

14291538
let mut func_flags = bytecode::MakeFunctionFlags::empty();
14301539

1540+
// Prepare generic type parameters:
1541+
if !type_params.is_empty() {
1542+
self.compile_type_params(type_params)?;
1543+
func_flags |= bytecode::MakeFunctionFlags::TYPE_PARAMS;
1544+
}
1545+
14311546
if self.build_closure(&code) {
14321547
func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
14331548
}
14341549

1550+
// Pop the special type params symbol table
1551+
if !type_params.is_empty() {
1552+
self.pop_symbol_table();
1553+
}
1554+
14351555
self.emit_load_const(ConstantData::Code {
14361556
code: Box::new(code),
14371557
});

0 commit comments

Comments
 (0)