Skip to content

Add 3.12 typing features to the compiler #5302

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 16 commits into from
May 10, 2024
Merged

Add 3.12 typing features to the compiler #5302

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
61a5c9b
Add first draft of typing.rs
dchiquito May 6, 2024
f6a5c12
Split up compiler push_output
dchiquito May 7, 2024
46a7d2c
Add type params to function definitions
dchiquito May 7, 2024
a0f6d5f
Add type params to class definitions
dchiquito May 7, 2024
df8b2c2
Add type params to type alias statements
dchiquito May 7, 2024
bfa18d2
Add TypeVar bytecodes
dchiquito May 7, 2024
42111ae
Add various Types to frame.rs
dchiquito May 7, 2024
d6261f5
Add _typing module to stdlib
dchiquito May 7, 2024
55f916c
Add _typing to default imports
dchiquito May 7, 2024
3d263a0
Remove unused comment
dchiquito May 7, 2024
2d97503
Fix clippy lints
dchiquito May 7, 2024
a230ca3
Use empty_tuple singleton in execute_make_function
dchiquito May 9, 2024
55e37f8
Update bytecode LAST_INSTRUCTION comments
dchiquito May 9, 2024
2f01963
Merge branch 'main' into typing-draft
dchiquito May 9, 2024
c907a4d
Switch emit_constant to emit_load_const
dchiquito May 9, 2024
ed30604
Remove unwrap
youknowone May 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
146 changes: 133 additions & 13 deletions compiler/codegen/src/compile.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,6 +251,26 @@ impl Compiler {
}
}

/// Push the next symbol table on to the stack
fn push_symbol_table(&mut self) -> &SymbolTable {
// Look up the next table contained in the scope of the current table
let table = self
.symbol_table_stack
.last_mut()
.expect("no next symbol table")
.sub_tables
.remove(0);
// Push the next table onto the stack
let last_idx = self.symbol_table_stack.len();
self.symbol_table_stack.push(table);
&self.symbol_table_stack[last_idx]
}

/// Pop the current symbol table off the stack
fn pop_symbol_table(&mut self) -> SymbolTable {
self.symbol_table_stack.pop().expect("compiler bug")
}

fn push_output(
&mut self,
flags: bytecode::CodeFlags,
Expand All @@ -262,12 +282,7 @@ impl Compiler {
let source_path = self.source_path.clone();
let first_line_number = self.get_source_line_number();

let table = self
.symbol_table_stack
.last_mut()
.unwrap()
.sub_tables
.remove(0);
let table = self.push_symbol_table();

let cellvar_cache = table
.symbols
Expand All @@ -284,8 +299,6 @@ impl Compiler {
.map(|(var, _)| var.clone())
.collect();

self.symbol_table_stack.push(table);

let info = ir::CodeInfo {
flags,
posonlyarg_count,
Expand All @@ -307,7 +320,7 @@ impl Compiler {
}

fn pop_code_object(&mut self) -> CodeObject {
let table = self.symbol_table_stack.pop().unwrap();
let table = self.pop_symbol_table();
assert!(table.sub_tables.is_empty());
self.code_stack
.pop()
Expand Down Expand Up @@ -752,6 +765,7 @@ impl Compiler {
body,
decorator_list,
returns,
type_params,
..
}) => self.compile_function_def(
name.as_str(),
Expand All @@ -760,13 +774,15 @@ impl Compiler {
decorator_list,
returns.as_deref(),
false,
type_params,
)?,
Stmt::AsyncFunctionDef(StmtAsyncFunctionDef {
name,
args,
body,
decorator_list,
returns,
type_params,
..
}) => self.compile_function_def(
name.as_str(),
Expand All @@ -775,15 +791,24 @@ impl Compiler {
decorator_list,
returns.as_deref(),
true,
type_params,
)?,
Stmt::ClassDef(StmtClassDef {
name,
body,
bases,
keywords,
decorator_list,
type_params,
..
}) => self.compile_class_def(name.as_str(), body, bases, keywords, decorator_list)?,
}) => self.compile_class_def(
name.as_str(),
body,
bases,
keywords,
decorator_list,
type_params,
)?,
Stmt::Assert(StmtAssert { test, msg, .. }) => {
// if some flag, ignore all assert statements!
if self.opts.optimize == 0 {
Expand Down Expand Up @@ -885,7 +910,27 @@ impl Compiler {
Stmt::Pass(_) => {
// No need to emit any code here :)
}
Stmt::TypeAlias(_) => {}
Stmt::TypeAlias(StmtTypeAlias {
name,
type_params,
value,
..
}) => {
let name_string = name.to_string();
if !type_params.is_empty() {
self.push_symbol_table();
}
self.compile_expression(value)?;
self.compile_type_params(type_params)?;
if !type_params.is_empty() {
self.pop_symbol_table();
}
self.emit_load_const(ConstantData::Str {
value: name_string.clone(),
});
emit!(self, Instruction::TypeAlias);
self.store_name(&name_string)?;
}
}
Ok(())
}
Expand Down Expand Up @@ -1005,6 +1050,47 @@ impl Compiler {
}
}

/// Store each type parameter so it is accessible to the current scope, and leave a tuple of
/// all the type parameters on the stack.
fn compile_type_params(&mut self, type_params: &[located_ast::TypeParam]) -> CompileResult<()> {
for type_param in type_params {
match type_param {
located_ast::TypeParam::TypeVar(located_ast::TypeParamTypeVar {
name,
bound,
..
}) => {
if let Some(expr) = &bound {
self.compile_expression(expr)?;
self.emit_load_const(ConstantData::Str {
value: name.to_string(),
});
emit!(self, Instruction::TypeVarWithBound);
emit!(self, Instruction::Duplicate);
self.store_name(name.as_ref())?;
} else {
// self.store_name(type_name.as_str())?;
self.emit_load_const(ConstantData::Str {
value: name.to_string(),
});
emit!(self, Instruction::TypeVar);
emit!(self, Instruction::Duplicate);
self.store_name(name.as_ref())?;
}
}
located_ast::TypeParam::ParamSpec(_) => todo!(),
located_ast::TypeParam::TypeVarTuple(_) => todo!(),
};
}
emit!(
self,
Instruction::BuildTuple {
size: u32::try_from(type_params.len()).unwrap(),
}
);
Ok(())
}

fn compile_try_statement(
&mut self,
body: &[located_ast::Stmt],
Expand Down Expand Up @@ -1151,6 +1237,7 @@ impl Compiler {
is_forbidden_name(name)
}

#[allow(clippy::too_many_arguments)]
fn compile_function_def(
&mut self,
name: &str,
Expand All @@ -1159,10 +1246,15 @@ impl Compiler {
decorator_list: &[located_ast::Expr],
returns: Option<&located_ast::Expr>, // TODO: use type hint somehow..
is_async: bool,
type_params: &[located_ast::TypeParam],
) -> CompileResult<()> {
// Create bytecode for this function:

self.prepare_decorators(decorator_list)?;

// If there are type params, we need to push a special symbol table just for them
if !type_params.is_empty() {
self.push_symbol_table();
}

let mut func_flags = self.enter_function(name, args)?;
self.current_code_info()
.flags
Expand Down Expand Up @@ -1208,6 +1300,12 @@ impl Compiler {
self.qualified_path.pop();
self.ctx = prev_ctx;

// Prepare generic type parameters:
if !type_params.is_empty() {
self.compile_type_params(type_params)?;
func_flags |= bytecode::MakeFunctionFlags::TYPE_PARAMS;
}

// Prepare type annotations:
let mut num_annotations = 0;

Expand Down Expand Up @@ -1253,6 +1351,11 @@ impl Compiler {
func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
}

// Pop the special type params symbol table
if !type_params.is_empty() {
self.pop_symbol_table();
}

self.emit_load_const(ConstantData::Code {
code: Box::new(code),
});
Expand Down Expand Up @@ -1352,6 +1455,7 @@ impl Compiler {
bases: &[located_ast::Expr],
keywords: &[located_ast::Keyword],
decorator_list: &[located_ast::Expr],
type_params: &[located_ast::TypeParam],
) -> CompileResult<()> {
self.prepare_decorators(decorator_list)?;

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

// If there are type params, we need to push a special symbol table just for them
if !type_params.is_empty() {
self.push_symbol_table();
}

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

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

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

// Prepare generic type parameters:
if !type_params.is_empty() {
self.compile_type_params(type_params)?;
func_flags |= bytecode::MakeFunctionFlags::TYPE_PARAMS;
}

if self.build_closure(&code) {
func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
}

// Pop the special type params symbol table
if !type_params.is_empty() {
self.pop_symbol_table();
}

self.emit_load_const(ConstantData::Code {
code: Box::new(code),
});
Expand Down
Loading