Improve interpreter perf by caching the name lookup scope

src/codegen/ast_interpreter.cpp

@@ -1126,23 +1126,55 @@ Value ASTInterpreter::visit_str(AST_Str* node) {
 }
 
 Value ASTInterpreter::visit_name(AST_Name* node) {
-    if (scope_info->refersToGlobal(node->id))
-        return getGlobal(source_info->parent_module, &node->id.str());
-    else if (scope_info->refersToClosure(node->id)) {
-        return getattr(passed_closure, node->id.c_str());
-    } else {
-        SymMap::iterator it = sym_table.find(node->id);
-        if (it != sym_table.end()) {
-            Box* value = it->second;
-            return value;
-        }
+    switch (node->lookup_type) {
+        case AST_Name::UNKNOWN: {
+            if (scope_info->refersToGlobal(node->id)) {
+                node->lookup_type = AST_Name::GLOBAL;
+                return getGlobal(source_info->parent_module, &node->id.str());
+            } else if (scope_info->refersToClosure(node->id)) {
+                node->lookup_type = AST_Name::CLOSURE;
+                return getattr(passed_closure, node->id.c_str());
+            } else {
+                bool is_old_local = (source_info->ast->type == AST_TYPE::ClassDef);
+                node->lookup_type = is_old_local ? AST_Name::LOCAL : AST_Name::FAST_LOCAL;
+
+                SymMap::iterator it = sym_table.find(node->id);
+                if (it != sym_table.end()) {
+                    Box* value = it->second;
+                    return value;
+                }
 
-        // classdefs have different scoping rules than functions:
-        if (source_info->ast->type == AST_TYPE::ClassDef)
+                // classdefs have different scoping rules than functions:
+                if (source_info->ast->type == AST_TYPE::ClassDef)
+                    return getGlobal(source_info->parent_module, &node->id.str());
+
+                assertNameDefined(0, node->id.c_str(), UnboundLocalError, true);
+                return Value();
+            }
+        }
+        case AST_Name::GLOBAL:
             return getGlobal(source_info->parent_module, &node->id.str());
+        case AST_Name::CLOSURE:
+            return getattr(passed_closure, node->id.c_str());
+        case AST_Name::FAST_LOCAL: {
+            SymMap::iterator it = sym_table.find(node->id);
+            if (it != sym_table.end())
+                return it->second;
+
+            assertNameDefined(0, node->id.c_str(), UnboundLocalError, true);
+            return Value();
+        }
+        case AST_Name::LOCAL: {
+            SymMap::iterator it = sym_table.find(node->id);
+            if (it != sym_table.end()) {
+                Box* value = it->second;
+                return value;
+            }
 
-        assertNameDefined(0, node->id.c_str(), UnboundLocalError, true);
-        return Value();
+            return getGlobal(source_info->parent_module, &node->id.str());
+        }
+        default:
+            abort();
     }
 }
 

src/core/ast.h

@@ -662,11 +662,23 @@ public:
     AST_TYPE::AST_TYPE ctx_type;
     InternedString id;
 
+    // The resolved scope of this name.  Kind of hacky to be storing it in the AST node;
+    // in CPython it ends up getting "cached" by being translated into one of a number of
+    // different bytecodes.
+    // We don't have a separate bytecode representation, so just store it in here for now.
+    enum LookupType {
+        UNKNOWN,
+        GLOBAL,
+        CLOSURE,
+        FAST_LOCAL,
+        LOCAL,
+    } lookup_type;
+
     virtual void accept(ASTVisitor* v);
     virtual void* accept_expr(ExprVisitor* v);
 
     AST_Name(InternedString id, AST_TYPE::AST_TYPE ctx_type, int lineno, int col_offset = 0)
-        : AST_expr(AST_TYPE::Name, lineno, col_offset), ctx_type(ctx_type), id(id) {}
+        : AST_expr(AST_TYPE::Name, lineno, col_offset), ctx_type(ctx_type), id(id), lookup_type(UNKNOWN) {}
 
     static const AST_TYPE::AST_TYPE TYPE = AST_TYPE::Name;
 };