Unwind generator stacks correctly

before we stopped at the generatorEntry func and did not unwind
into the caller.

src/codegen/unwinding.cpp

@@ -29,6 +29,7 @@
 #include "codegen/irgen/hooks.h"
 #include "codegen/stackmaps.h"
 #include "core/util.h"
+#include "runtime/generator.h"
 #include "runtime/traceback.h"
 #include "runtime/types.h"
 
@@ -328,25 +329,17 @@ public:
         return rtn;
     }
 
-    bool isMainInterpreterFunc(unw_word_t ip) {
-        // Remember the addr of the end of the interpreter function because unw_get_proc_info is slow and if we know
-        // the bounds we can replace it with pointer comparisons.
-        static intptr_t interpreter_instr_end = 0; // first ip NOT covered by interpreter func
-        if (interpreter_instr_end == 0) {
-            unw_proc_info_t pip;
-            int code = unw_get_proc_info(&this->cursor, &pip);
-            RELEASE_ASSERT(code == 0, "%d", code);
-            if (pip.start_ip == (intptr_t)interpreter_instr_addr) {
-                interpreter_instr_end = pip.end_ip;
-                return true;
-            }
-        } else if (ip >= (intptr_t)interpreter_instr_addr && ip < interpreter_instr_end) {
-            return true;
-        }
-        return false;
+    unw_word_t getFunctionEnd(unw_word_t ip) {
+        unw_proc_info_t pip;
+        int ret = unw_get_proc_info_by_ip(unw_local_addr_space, ip, &pip, NULL);
+        RELEASE_ASSERT(ret == 0 && pip.end_ip, "");
+        return pip.end_ip;
     }
 
     bool incr() {
+        static unw_word_t interpreter_instr_end = getFunctionEnd((unw_word_t)interpreter_instr_addr);
+        static unw_word_t generator_entry_end = getFunctionEnd((unw_word_t)generatorEntry);
+
         bool was_osr = cur_is_osr;
 
         while (true) {
@@ -377,7 +370,7 @@ public:
                 return true;
             }
 
-            if (isMainInterpreterFunc(ip)) {
+            if ((unw_word_t)interpreter_instr_addr <= ip && ip < interpreter_instr_end) {
                 unw_word_t bp;
                 unw_get_reg(&this->cursor, UNW_TDEP_BP, &bp);
 
@@ -395,6 +388,15 @@ public:
                 return true;
             }
 
+            if ((unw_word_t)generatorEntry <= ip && ip < generator_entry_end) {
+                // for generators continue unwinding in the context in which the generator got called
+                static_assert(sizeof(ucontext_t) == sizeof(unw_context_t), "");
+                unw_word_t bp;
+                unw_get_reg(&this->cursor, UNW_TDEP_BP, &bp);
+                ucontext_t* remote_ctx = getReturnContextForGeneratorFrame((void*)bp);
+                unw_init_local(&cursor, remote_ctx);
+            }
+
             // keep unwinding
         }
     }

src/runtime/generator.cpp

@@ -41,13 +41,38 @@ static std::vector<uint64_t> available_addrs;
 #define STACK_REDZONE_SIZE PAGE_SIZE
 #define MAX_STACK_SIZE (4 * 1024 * 1024)
 
-static void generatorEntry(BoxedGenerator* g) {
+static std::unordered_map<void*, BoxedGenerator*> s_generator_map;
+static_assert(THREADING_USE_GIL, "have to make the generator map thread safe!");
+
+class RegisterHelper {
+private:
+    void* frame_addr;
+
+public:
+    RegisterHelper(BoxedGenerator* generator, void* frame_addr) : frame_addr(frame_addr) {
+        s_generator_map[frame_addr] = generator;
+    }
+    ~RegisterHelper() {
+        assert(s_generator_map.count(frame_addr));
+        s_generator_map.erase(frame_addr);
+    }
+};
+
+ucontext* getReturnContextForGeneratorFrame(void* frame_addr) {
+    BoxedGenerator* generator = s_generator_map[frame_addr];
+    assert(generator);
+    return &generator->returnContext;
+}
+
+void generatorEntry(BoxedGenerator* g) {
     assert(g->cls == generator_cls);
     assert(g->function->cls == function_cls);
 
     threading::pushGenerator(g, g->stack_begin, (void*)g->returnContext.uc_mcontext.gregs[REG_RSP]);
 
     try {
+        RegisterHelper context_registerer(g, __builtin_frame_address(0));
+
         // call body of the generator
         BoxedFunctionBase* func = g->function;
 

src/runtime/generator.h

@@ -18,11 +18,15 @@
 #include "core/types.h"
 #include "runtime/types.h"
 
+struct ucontext;
+
 namespace pyston {
 
 extern BoxedClass* generator_cls;
 
 void setupGenerator();
+void generatorEntry(BoxedGenerator* g);
+ucontext* getReturnContextForGeneratorFrame(void* frame_addr);
 
 extern "C" Box* yield(BoxedGenerator* obj, Box* value);
 extern "C" BoxedGenerator* createGenerator(BoxedFunctionBase* function, Box* arg1, Box* arg2, Box* arg3, Box** args);