Fast locals part 3

src/compile.rs

@@ -273,14 +273,13 @@ impl Compiler {
         let cellvar_cache = table
             .symbols
             .iter()
-            .filter(|(_, s)| matches!(s.scope, SymbolScope::Cell))
+            .filter(|(_, s)| s.scope == SymbolScope::Cell)
             .map(|(var, _)| var.clone())
             .collect();
         let freevar_cache = table
             .symbols
             .iter()
-            // TODO: check if Free or FREE_CLASS symbol
+            .filter(|(_, s)| s.scope == SymbolScope::Free || s.is_free_class)
-            .filter(|(_, s)| matches!(s.scope, SymbolScope::Free))
             .map(|(var, _)| var.clone())
             .collect();
 
@@ -336,6 +335,11 @@ impl Compiler {
             let doc = self.name("__doc__");
             self.emit(Instruction::StoreGlobal(doc))
         }
+
+        if self.find_ann(statements) {
+            self.emit(Instruction::SetupAnnotation);
+        }
+
         self.compile_statements(statements)?;
 
         assert_eq!(self.code_stack.len(), size_before);
@@ -433,9 +437,12 @@ impl Compiler {
                 cache = &mut info.varname_cache;
                 NameOpType::Fast
             }
-            SymbolScope::GlobalImplicit if self.ctx.in_func() => NameOpType::Global,
-            SymbolScope::Local | SymbolScope::GlobalImplicit => NameOpType::Local,
             SymbolScope::GlobalExplicit => NameOpType::Global,
+            SymbolScope::GlobalImplicit | SymbolScope::Unknown if self.ctx.in_func() => {
+                NameOpType::Global
+            }
+            SymbolScope::GlobalImplicit | SymbolScope::Unknown => NameOpType::Local,
+            SymbolScope::Local => NameOpType::Local,
             SymbolScope::Free => {
                 cache = &mut info.freevar_cache;
                 NameOpType::Deref
@@ -444,8 +451,8 @@ impl Compiler {
                 cache = &mut info.cellvar_cache;
                 NameOpType::Deref
             }
-            // TODO: is this right?
+            // // TODO: is this right?
-            SymbolScope::Unknown => NameOpType::Global,
+            // SymbolScope::Unknown => NameOpType::Global,
         };
         let mut idx = cache
             .get_index_of(name)
@@ -527,6 +534,10 @@ impl Compiler {
                 let module_idx = module.as_ref().map(|s| self.name(s));
 
                 if import_star {
+                    if self.ctx.in_func() {
+                        return Err(self
+                            .error_loc(CompileErrorType::FunctionImportStar, statement.location));
+                    }
                     let star = self.name("*");
                     // from .... import *
                     self.emit(Instruction::Import {
@@ -1069,6 +1080,8 @@ impl Compiler {
         }
 
         let mut code = self.pop_code_object();
+        self.current_qualified_path = old_qualified_path;
+        self.ctx = prev_ctx;
 
         // Prepare type annotations:
         let mut num_annotations = 0;
@@ -1138,9 +1151,6 @@ impl Compiler {
         let doc = self.name("__doc__");
         self.emit(Instruction::StoreAttr { idx: doc });
 
-        self.current_qualified_path = old_qualified_path;
-        self.ctx = prev_ctx;
-
         self.apply_decorators(decorator_list);
 
         self.store_name(name);
@@ -1151,12 +1161,17 @@ impl Compiler {
     fn build_closure(&mut self, code: &CodeObject) {
         if !code.freevars.is_empty() {
             for var in &code.freevars {
-                let symbol = self.symbol_table_stack.last().unwrap().lookup(var).unwrap();
+                let table = self.symbol_table_stack.last().unwrap();
+                let symbol = table.lookup(var).unwrap();
                 let parent_code = self.code_stack.last().unwrap();
                 let vars = match symbol.scope {
                     SymbolScope::Free => &parent_code.freevar_cache,
                     SymbolScope::Cell => &parent_code.cellvar_cache,
-                    _ => unreachable!(),
+                    _ if symbol.is_free_class => &parent_code.freevar_cache,
+                    x => unreachable!(
+                        "var {} in a {:?} should be free or cell but it's {:?}",
+                        var, table.typ, x
+                    ),
                 };
                 let mut idx = vars.get_index_of(var).unwrap();
                 if let SymbolScope::Free = symbol.scope {
@@ -1236,6 +1251,8 @@ impl Compiler {
         keywords: &[ast::Keyword],
         decorator_list: &[ast::Expression],
     ) -> CompileResult<()> {
+        self.prepare_decorators(decorator_list)?;
+
         let prev_ctx = self.ctx;
         self.ctx = CompileContext {
             func: FunctionContext::NoFunction,
@@ -1247,7 +1264,6 @@ impl Compiler {
         let old_qualified_path = self.current_qualified_path.take();
         self.current_qualified_path = Some(qualified_name.clone());
 
-        self.prepare_decorators(decorator_list)?;
         self.emit(Instruction::LoadBuildClass);
         let line_number = self.get_source_line_number();
         self.push_output(CodeObject::new(
@@ -1301,6 +1317,9 @@ impl Compiler {
 
         let code = self.pop_code_object();
 
+        self.current_qualified_path = old_qualified_path;
+        self.ctx = prev_ctx;
+
         self.build_closure(&code);
 
         self.emit_constant(bytecode::ConstantData::Code {
@@ -1350,8 +1369,6 @@ impl Compiler {
         self.apply_decorators(decorator_list);
 
         self.store_name(name);
-        self.current_qualified_path = old_qualified_path;
-        self.ctx = prev_ctx;
         Ok(())
     }
 
@@ -1579,18 +1596,17 @@ impl Compiler {
 
         if let ast::ExpressionType::Identifier { name } = &target.node {
             // Store as dict entry in __annotations__ dict:
-            if !self.ctx.in_func() {
             let annotations = self.name("__annotations__");
             self.emit(Instruction::LoadNameAny(annotations));
             self.emit_constant(bytecode::ConstantData::Str {
                 value: name.to_owned(),
             });
             self.emit(Instruction::StoreSubscript);
-            }
         } else {
             // Drop annotation if not assigned to simple identifier.
             self.emit(Instruction::Pop);
         }
+
         Ok(())
     }
 
@@ -2187,7 +2203,7 @@ impl Compiler {
         let line_number = self.get_source_line_number();
         // Create magnificent function <listcomp>:
         self.push_output(CodeObject::new(
-            Default::default(),
+            bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
             1,
             1,
             0,

src/error.rs

@@ -36,6 +36,7 @@ pub enum CompileErrorType {
     AsyncReturnValue,
     InvalidFuturePlacement,
     InvalidFutureFeature(String),
+    FunctionImportStar,
 }
 
 impl fmt::Display for CompileErrorType {
@@ -66,6 +67,9 @@ impl fmt::Display for CompileErrorType {
             CompileErrorType::InvalidFutureFeature(feat) => {
                 write!(f, "future feature {} is not defined", feat)
             }
+            CompileErrorType::FunctionImportStar => {
+                write!(f, "import * only allowed at module level")
+            }
         }
     }
 }

src/symboltable.rs

@@ -85,7 +85,7 @@ impl fmt::Display for SymbolTableType {
 
 /// Indicator for a single symbol what the scope of this symbol is.
 /// The scope can be unknown, which is unfortunate, but not impossible.
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq)]
 pub enum SymbolScope {
     Unknown,
     Local,
@@ -116,6 +116,16 @@ pub struct Symbol {
     // inidicates that the symbol is used a bound iterator variable. We distinguish this case
     // from normal assignment to detect unallowed re-assignment to iterator variables.
     pub is_iter: bool,
+
+    /// indicates that the symbol is a free variable in a class method from the scope that the
+    /// class is defined in, e.g.:
+    /// ```python
+    /// def foo(x):
+    ///     class A:
+    ///         def method(self):
+    ///             return x // is_free_class
+    /// ```
+    pub is_free_class: bool,
 }
 
 impl Symbol {
@@ -131,6 +141,7 @@ impl Symbol {
             is_imported: false,
             is_assign_namedexpr_in_comprehension: false,
             is_iter: false,
+            is_free_class: false,
         }
     }
 
@@ -142,7 +153,7 @@ impl Symbol {
     }
 
     pub fn is_local(&self) -> bool {
-        matches!(self.scope, SymbolScope::Local)
+        self.scope == SymbolScope::Local
     }
 
     pub fn is_bound(&self) -> bool {
@@ -283,12 +294,10 @@ impl SymbolTableAnalyzer {
     fn analyze_symbol(
         &mut self,
         symbol: &mut Symbol,
-        curr_st_typ: SymbolTableType,
+        st_typ: SymbolTableType,
         sub_tables: &mut [SymbolTable],
     ) -> SymbolTableResult {
-        if symbol.is_assign_namedexpr_in_comprehension
+        if symbol.is_assign_namedexpr_in_comprehension && st_typ == SymbolTableType::Comprehension {
-            && curr_st_typ == SymbolTableType::Comprehension
-        {
             // propagate symbol to next higher level that can hold it,
             // i.e., function or module. Comprehension is skipped and
             // Class is not allowed and detected as error.
@@ -298,10 +307,12 @@ impl SymbolTableAnalyzer {
             match symbol.scope {
                 SymbolScope::Free => {
                     if !self.tables.as_ref().is_empty() {
-                        let scope_depth = self.tables.as_ref().iter().count();
+                        let scope_depth = self.tables.as_ref().len();
                         // check if the name is already defined in any outer scope
                         // therefore
-                        if scope_depth < 2 || !self.found_in_outer_scope(&symbol.name) {
+                        if scope_depth < 2
+                            || self.found_in_outer_scope(&symbol.name) != Some(SymbolScope::Free)
+                        {
                             return Err(SymbolTableError {
                                 error: format!("no binding for nonlocal '{}' found", symbol.name),
                                 // TODO: accurate location info, somehow
@@ -327,58 +338,11 @@ impl SymbolTableAnalyzer {
                 }
                 SymbolScope::Unknown => {
                     // Try hard to figure out what the scope of this symbol is.
-                    self.analyze_unknown_symbol(sub_tables, symbol);
-                }
-            }
-        }
-        Ok(())
-    }
-
-    fn found_in_outer_scope(&mut self, name: &str) -> bool {
-        // Interesting stuff about the __class__ variable:
-        // https://docs.python.org/3/reference/datamodel.html?highlight=__class__#creating-the-class-object
-        if name == "__class__" {
-            return true;
-        }
-        let decl_depth = self.tables.iter().rev().position(|(symbols, typ)| {
-            !matches!(typ, SymbolTableType::Class | SymbolTableType::Module)
-                && symbols.get(name).map_or(false, |sym| sym.is_bound())
-        });
-
-        if let Some(decl_depth) = decl_depth {
-            // decl_depth is the number of tables between the current one and
-            // the one that declared the cell var
-            for (table, _) in self.tables.iter_mut().rev().take(decl_depth) {
-                if !table.contains_key(name) {
-                    let mut symbol = Symbol::new(name);
-                    symbol.scope = SymbolScope::Free;
-                    symbol.is_referenced = true;
-                    table.insert(name.to_owned(), symbol);
-                }
-            }
-        }
-
-        decl_depth.is_some()
-    }
-
-    fn found_in_inner_scope(sub_tables: &mut [SymbolTable], name: &str) -> bool {
-        sub_tables.iter().any(|x| {
-            x.symbols
-                .get(name)
-                .map_or(false, |sym| matches!(sym.scope, SymbolScope::Free))
-        })
-    }
-
-    fn analyze_unknown_symbol(&mut self, sub_tables: &mut [SymbolTable], symbol: &mut Symbol) {
                     let scope = if symbol.is_bound() {
-            if Self::found_in_inner_scope(sub_tables, &symbol.name) {
+                        self.found_in_inner_scope(sub_tables, &symbol.name, st_typ)
-                SymbolScope::Cell
+                            .unwrap_or(SymbolScope::Local)
-            } else {
+                    } else if let Some(scope) = self.found_in_outer_scope(&symbol.name) {
-                SymbolScope::Local
+                        scope
-            }
-        } else if self.found_in_outer_scope(&symbol.name) {
-            // Symbol is in some outer scope.
-            SymbolScope::Free
                     } else if self.tables.is_empty() {
                         // Don't make assumptions when we don't know.
                         SymbolScope::Unknown
@@ -388,6 +352,86 @@ impl SymbolTableAnalyzer {
                     };
                     symbol.scope = scope;
                 }
+            }
+        }
+        Ok(())
+    }
+
+    fn found_in_outer_scope(&mut self, name: &str) -> Option<SymbolScope> {
+        // Interesting stuff about the __class__ variable:
+        // https://docs.python.org/3/reference/datamodel.html?highlight=__class__#creating-the-class-object
+        if name == "__class__" {
+            return Some(SymbolScope::Free);
+        }
+
+        let mut decl_depth = None;
+        for (i, (symbols, typ)) in self.tables.iter().rev().enumerate() {
+            if let SymbolTableType::Class | SymbolTableType::Module = typ {
+                continue;
+            }
+            if let Some(sym) = symbols.get(name) {
+                match sym.scope {
+                    SymbolScope::GlobalExplicit => return Some(SymbolScope::GlobalExplicit),
+                    SymbolScope::GlobalImplicit => {}
+                    _ => {
+                        if sym.is_bound() {
+                            decl_depth = Some(i);
+                            break;
+                        }
+                    }
+                }
+            }
+        }
+
+        if let Some(decl_depth) = decl_depth {
+            // decl_depth is the number of tables between the current one and
+            // the one that declared the cell var
+            for (table, typ) in self.tables.iter_mut().rev().take(decl_depth) {
+                if let SymbolTableType::Class = typ {
+                    if let Some(free_class) = table.get_mut(name) {
+                        free_class.is_free_class = true;
+                    } else {
+                        let mut symbol = Symbol::new(name);
+                        symbol.is_free_class = true;
+                        symbol.scope = SymbolScope::Free;
+                        table.insert(name.to_owned(), symbol);
+                    }
+                } else if !table.contains_key(name) {
+                    let mut symbol = Symbol::new(name);
+                    symbol.scope = SymbolScope::Free;
+                    // symbol.is_referenced = true;
+                    table.insert(name.to_owned(), symbol);
+                }
+            }
+        }
+
+        decl_depth.map(|_| SymbolScope::Free)
+    }
+
+    fn found_in_inner_scope(
+        &self,
+        sub_tables: &[SymbolTable],
+        name: &str,
+        st_typ: SymbolTableType,
+    ) -> Option<SymbolScope> {
+        sub_tables.iter().find_map(|st| {
+            st.symbols.get(name).and_then(|sym| {
+                if sym.scope == SymbolScope::Free || sym.is_free_class {
+                    if st_typ == SymbolTableType::Class && name != "__class__" {
+                        None
+                    } else {
+                        Some(SymbolScope::Cell)
+                    }
+                } else if sym.scope == SymbolScope::GlobalExplicit && self.tables.is_empty() {
+                    // the symbol is defined on the module level, and an inner scope declares
+                    // a global that points to it
+                    Some(SymbolScope::GlobalExplicit)
+                } else {
+                    None
+                }
+            })
+        })
+    }
 
     // Implements the symbol analysis and scope extension for names
     // assigned by a named expression in a comprehension. See:
@@ -431,12 +475,13 @@ impl SymbolTableAnalyzer {
                     if let SymbolScope::Unknown = parent_symbol.scope {
                         // this information is new, as the asignment is done in inner scope
                         parent_symbol.is_assigned = true;
-                        //self.analyze_unknown_symbol(symbol); // not needed, symbol is analyzed anyhow when its scope is analyzed
                     }
 
-                    if !symbol.is_global() {
+                    symbol.scope = if parent_symbol.is_global() {
-                        symbol.scope = SymbolScope::Free;
+                        parent_symbol.scope
-                    }
+                    } else {
+                        SymbolScope::Free
+                    };
                 } else {
                     let mut cloned_sym = symbol.clone();
                     cloned_sym.scope = SymbolScope::Local;
@@ -907,6 +952,7 @@ impl SymbolTableBuilder {
                 // Determine the contextual usage of this symbol:
                 match context {
                     ExpressionContext::Delete => {
+                        self.register_name(name, SymbolUsage::Assigned, location)?;
                         self.register_name(name, SymbolUsage::Used, location)?;
                     }
                     ExpressionContext::Load | ExpressionContext::IterDefinitionExp => {
@@ -1051,6 +1097,7 @@ impl SymbolTableBuilder {
                 SymbolUsage::Global => {
                     if symbol.is_global() {
                         // Ok
+                        symbol.scope = SymbolScope::GlobalExplicit;
                     } else {
                         return Err(SymbolTableError {
                             error: format!("name '{}' is used prior to global declaration", name),
@@ -1140,9 +1187,9 @@ impl SymbolTableBuilder {
             }
             SymbolUsage::Global => {
                 if let SymbolScope::Unknown = symbol.scope {
-                    symbol.scope = SymbolScope::GlobalImplicit;
+                    symbol.scope = SymbolScope::GlobalExplicit;
                 } else if symbol.is_global() {
-                    // Global scope can be set to global
+                    symbol.scope = SymbolScope::GlobalExplicit;
                 } else {
                     return Err(SymbolTableError {
                         error: format!("Symbol {} scope cannot be set to global, since its scope was already determined otherwise.", name),