static closures cleanup

src/analysis/scoping_analysis.cpp

@@ -328,6 +328,8 @@ public:
 
         size_t parentCounter = 0;
         // Casting to a ScopeInfoBase* is okay because only a ScopeInfoBase can have a closure.
+        // We just walk up the scopes until we find the scope with this name. Count the number
+        // of parent links we follow, and then get the offset of the name.
         for (ScopeInfoBase* parent = static_cast<ScopeInfoBase*>(this->parent); parent != NULL;
              parent = static_cast<ScopeInfoBase*>(parent->parent)) {
             if (parent->createsClosure()) {
@@ -347,7 +349,10 @@ public:
         return closure_offsets[name];
     }
 
-    size_t getClosureSize() override { return closure_offsets.size(); }
+    size_t getClosureSize() override {
+        assert(createsClosure());
+        return closure_offsets.size();
+    }
 
     InternedString mangleName(const InternedString id) override {
         return pyston::mangleName(id, usage->private_name, usage->scoping->getInternedStrings());
@@ -359,6 +364,10 @@ public:
         if (!allDerefVarsAndInfoCached) {
             allDerefVarsAndInfoCached = true;
 
+            // TODO this could probably be implemented faster
+
+            // Get all the variables that we need to return: any variable from the
+            // passed-in closure that is accessed in this scope or in a child scope.
             StrSet allDerefs = usage->got_from_closure;
             for (InternedString name : usage->passthrough_accesses) {
                 if (allDerefs.find(name) != allDerefs.end()) {
@@ -366,10 +375,13 @@ public:
                 }
             }
 
+            // Call `getDerefInfo` on all of these variables and put the results in
+            // `allDerefVarsAndInfo`
             for (InternedString name : allDerefs) {
                 allDerefVarsAndInfo.push_back({ name, getDerefInfo(name) });
             }
 
+            // Sort in order of `num_parents_from_passed_closure`
             std::sort(allDerefVarsAndInfo.begin(), allDerefVarsAndInfo.end(), derefComparator);
         }
         return allDerefVarsAndInfo;

src/analysis/scoping_analysis.h

@@ -25,6 +25,11 @@ class AST_Module;
 class AST_Expression;
 class AST_Suite;
 
+// Each closure has an array (fixed-size for that particular scope) of variables
+// and a parent pointer to a parent closure. To look up a variable from the passed-in
+// closure (i.e., DEREF), you just need to know (i) how many parents up to go and
+// (ii) what offset into the array to find the variable. This struct stores that
+// information. You can query the ScopeInfo with a name to get this info.
 struct DerefInfo {
     size_t num_parents_from_passed_closure;
     size_t offset;
@@ -88,9 +93,50 @@ public:
 
     virtual bool areLocalsFromModule() = 0;
 
+    // For a variable with DEREF lookup, return the DerefInfo used to lookup
+    // the variable in a passed closure.
     virtual DerefInfo getDerefInfo(InternedString name) = 0;
+
+    // Gets the DerefInfo for each DEREF variable accessible in the scope.
+    // The returned vector is in SORTED ORDER by the `num_parents_from_passed_closure` field.
+    // This allows the caller to iterate through the vector while also walking up
+    // the closure chain to collect all the DEREF variable values. This is useful, for example,
+    // in the implementation of locals().
+    //
+    // Note that:
+    //      (a) This may not return a variable even if it is in the passed-in scope,
+    //          if the variable is not actually used in this scope or any child
+    //          scopes. This can happen, because the variable
+    //          could be in the closure to be accessed by a different function, e.g.
+    //
+    //                  def f();
+    //                      a = 0
+    //                      b = 0
+    //                      def g():
+    //                           print a
+    //                      def h():
+    //                           print b
+    //                           # locals() should not contain `a` even though `h` is
+    //                           # passed a closure object with `a` in it
+    //                           print locals()
+    //
+    //      (b) This can contain a variable even if it is not access in this scope,
+    //          if it used in a child scope instead. For example:
+    //
+    //                  def f():
+    //                       a = 0
+    //                       def g():
+    //                           def h():
+    //                               print a
+    //                           print locals() # should contain `a`
     virtual const std::vector<std::pair<InternedString, DerefInfo>>& getAllDerefVarsAndInfo() = 0;
+
+    // For a variable with CLOSURE lookup, returns the offset within the `elts`
+    // array of a closure that this variable is stored.
     virtual size_t getClosureOffset(InternedString name) = 0;
+
+    // Returns the size of the `elts` array for a closure created by this scope.
+    // Should only be called if this scope creates a closure.
     virtual size_t getClosureSize() = 0;
 
     virtual InternedString mangleName(InternedString id) = 0;

src/codegen/irgen/irgenerator.cpp

@@ -72,6 +72,21 @@ llvm::Value* IRGenState::getScratchSpace(int min_bytes) {
     return scratch_space;
 }
 
+static llvm::Value* getClosureParentGep(IREmitter& emitter, llvm::Value* closure) {
+    static_assert(sizeof(Box) == offsetof(BoxedClosure, parent), "");
+    static_assert(offsetof(BoxedClosure, parent) + sizeof(BoxedClosure*) == offsetof(BoxedClosure, nelts), "");
+    return emitter.getBuilder()->CreateConstInBoundsGEP2_32(closure, 0, 1);
+}
+
+static llvm::Value* getClosureElementGep(IREmitter& emitter, llvm::Value* closure, size_t index) {
+    static_assert(sizeof(Box) == offsetof(BoxedClosure, parent), "");
+    static_assert(offsetof(BoxedClosure, parent) + sizeof(BoxedClosure*) == offsetof(BoxedClosure, nelts), "");
+    static_assert(offsetof(BoxedClosure, nelts) + sizeof(size_t) == offsetof(BoxedClosure, elts), "");
+    return emitter.getBuilder()->CreateGEP(
+        closure,
+        { llvm::ConstantInt::get(g.i32, 0), llvm::ConstantInt::get(g.i32, 3), llvm::ConstantInt::get(g.i32, index) });
+}
+
 static llvm::Value* getBoxedLocalsGep(llvm::IRBuilder<true>& builder, llvm::Value* v) {
     static_assert(offsetof(FrameInfo, exc) == 0, "");
     static_assert(sizeof(ExcInfo) == 24, "");
@@ -898,27 +913,27 @@ private:
             assert(!is_kill);
             assert(scope_info->takesClosure());
 
+            // This is the information on how to look up the variable in the closure object.
             DerefInfo deref_info = scope_info->getDerefInfo(node->id);
 
-            static_assert(sizeof(Box) == offsetof(BoxedClosure, parent), "");
-            static_assert(offsetof(BoxedClosure, parent) + sizeof(BoxedClosure*) == offsetof(BoxedClosure, nelts), "");
-            static_assert(offsetof(BoxedClosure, nelts) + sizeof(size_t) == offsetof(BoxedClosure, elts), "");
-
+            // This code is basically:
+            // closure = created_closure;
+            // closure = closure->parent;
+            // [...]
+            // closure = closure->parent;
+            // closure->elts[deref_info.offset]
+            // Where the parent lookup is done `deref_info.num_parents_from_passed_closure` times
             CompilerVariable* closure = symbol_table[internString(PASSED_CLOSURE_NAME)];
             llvm::Value* closureValue = closure->makeConverted(emitter, CLOSURE)->getValue();
             closure->decvref(emitter);
-            llvm::Value* gep;
             for (int i = 0; i < deref_info.num_parents_from_passed_closure; i++) {
-                gep = emitter.getBuilder()->CreateConstInBoundsGEP2_32(closureValue, 0, 1);
-                closureValue = emitter.getBuilder()->CreateLoad(gep);
+                closureValue = emitter.getBuilder()->CreateLoad(getClosureParentGep(emitter, closureValue));
             }
-            gep = emitter.getBuilder()->CreateGEP(closureValue,
-                                                  { llvm::ConstantInt::get(g.i32, 0), llvm::ConstantInt::get(g.i32, 3),
-                                                    llvm::ConstantInt::get(g.i32, deref_info.offset) });
-            llvm::Value* lookupResult = emitter.getBuilder()->CreateLoad(gep);
+            llvm::Value* lookupResult
+                = emitter.getBuilder()->CreateLoad(getClosureElementGep(emitter, closureValue, deref_info.offset));
 
-            // If the value is NULL, it is undefined.
-            // Create a branch on if the value is NULL
+            // If the value is NULL, the variable is undefined.
+            // Create a branch on if the value is NULL.
             llvm::BasicBlock* success_bb
                 = llvm::BasicBlock::Create(g.context, "deref_defined", irstate->getLLVMFunction());
             success_bb->moveAfter(curblock);
@@ -929,7 +944,7 @@ private:
                 = emitter.getBuilder()->CreateICmpEQ(lookupResult, embedConstantPtr(NULL, g.llvm_value_type_ptr));
             llvm::BranchInst* non_null_check = emitter.getBuilder()->CreateCondBr(check_val, fail_bb, success_bb);
 
-            // In the case that it failed, call the assert fail function
+            // Case that it is undefined: call the assert fail function.
             curblock = fail_bb;
             emitter.getBuilder()->SetInsertPoint(curblock);
 
@@ -938,7 +953,7 @@ private:
             call.setDoesNotReturn();
             emitter.getBuilder()->CreateUnreachable();
 
-            // Carry on in the case that it succeeded
+            // Case that it is defined: carry on in with the retrieved value.
             curblock = success_bb;
             emitter.getBuilder()->SetInsertPoint(curblock);
 
@@ -1474,12 +1489,11 @@ private:
             if (vst == ScopeInfo::VarScopeType::CLOSURE) {
                 size_t offset = scope_info->getClosureOffset(name);
 
+                // This is basically `closure->elts[offset] = val;`
                 CompilerVariable* closure = symbol_table[internString(CREATED_CLOSURE_NAME)];
                 llvm::Value* closureValue = closure->makeConverted(emitter, CLOSURE)->getValue();
                 closure->decvref(emitter);
-                llvm::Value* gep = emitter.getBuilder()->CreateGEP(
-                    closureValue, { llvm::ConstantInt::get(g.i32, 0), llvm::ConstantInt::get(g.i32, 3),
-                                    llvm::ConstantInt::get(g.i32, offset) });
+                llvm::Value* gep = getClosureElementGep(emitter, closureValue, offset);
                 emitter.getBuilder()->CreateStore(val->makeConverted(emitter, UNKNOWN)->getValue(), gep);
             }
         }