Merge pull request #296 from undingen/gen_traceback

Generator tracebacks

src/analysis/type_analysis.cpp

@@ -372,6 +372,7 @@ private:
                 return UNKNOWN;
             case AST_LangPrimitive::NONE:
             case AST_LangPrimitive::SET_EXC_INFO:
+            case AST_LangPrimitive::UNCACHE_EXC_INFO:
                 return NONE;
             case AST_LangPrimitive::NONZERO:
                 return BOOL;

src/codegen/ast_interpreter.cpp

@@ -562,6 +562,10 @@ Value ASTInterpreter::visit_langPrimitive(AST_LangPrimitive* node) {
 
         getFrameInfo()->exc = ExcInfo(type.o, value.o, traceback.o);
         v = None;
+    } else if (node->opcode == AST_LangPrimitive::UNCACHE_EXC_INFO) {
+        assert(node->args.empty());
+        getFrameInfo()->exc = ExcInfo(NULL, NULL, NULL);
+        v = None;
     } else
         RELEASE_ASSERT(0, "not implemented");
     return v;

src/codegen/irgen/irgenerator.cpp

@@ -638,6 +638,22 @@ private:
 
                 return getNone();
             }
+            case AST_LangPrimitive::UNCACHE_EXC_INFO: {
+                assert(node->args.empty());
+
+                auto* builder = emitter.getBuilder();
+
+                llvm::Value* frame_info = irstate->getFrameInfoVar();
+                llvm::Value* exc_info = builder->CreateConstInBoundsGEP2_32(frame_info, 0, 0);
+                assert(exc_info->getType() == g.llvm_excinfo_type->getPointerTo());
+
+                llvm::Constant* v = embedConstantPtr(0, g.llvm_value_type_ptr);
+                builder->CreateStore(v, builder->CreateConstInBoundsGEP2_32(exc_info, 0, 0));
+                builder->CreateStore(v, builder->CreateConstInBoundsGEP2_32(exc_info, 0, 1));
+                builder->CreateStore(v, builder->CreateConstInBoundsGEP2_32(exc_info, 0, 2));
+
+                return getNone();
+            }
             default:
                 RELEASE_ASSERT(0, "%d", node->opcode);
         }

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/core/ast.cpp

@@ -1454,6 +1454,9 @@ bool PrintVisitor::visit_langprimitive(AST_LangPrimitive* node) {
         case AST_LangPrimitive::SET_EXC_INFO:
             printf("SET_EXC_INFO");
             break;
+        case AST_LangPrimitive::UNCACHE_EXC_INFO:
+            printf("UNCACHE_EXC_INFO");
+            break;
         default:
             RELEASE_ASSERT(0, "%d", node->opcode);
     }

src/core/ast.h

@@ -997,6 +997,7 @@ public:
         NONE,
         NONZERO,
         SET_EXC_INFO,
+        UNCACHE_EXC_INFO,
     } opcode;
     std::vector<AST_expr*> args;
 

src/core/cfg.cpp

@@ -1006,7 +1006,13 @@ private:
         rtn->lineno = node->lineno;
         rtn->col_offset = node->col_offset;
         rtn->value = remapExpr(node->value);
-        return rtn;
+
+        InternedString node_name(nodeName(rtn));
+        pushAssign(node_name, rtn);
+
+        push_back(makeExpr(new AST_LangPrimitive(AST_LangPrimitive::UNCACHE_EXC_INFO)));
+
+        return makeName(node_name, AST_TYPE::Load, node->lineno);
     }
 
     AST_expr* remapExpr(AST_expr* node, bool wrap_with_assign = true) {

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);

test/tests/generator_tracebacks.py

-# expected: fail
-# - wip
-
 # Generators participate in the notional Python stack just like normal function calls do,
 # even if we implement them using separate C stacks.
 #
@@ -8,6 +5,7 @@
 # get inherited when we go into a generator
 
 # exc_info gets passed into generators (at both begin and send()) and cleared like normal on the way out:
+import sys
 def f12():
     print
     print "f12"
@@ -25,6 +23,7 @@ def f12():
         print "after KeyError:", sys.exc_info()[0]
         yield 2
 
+    print list(g())
     try:
         raise AttributeError()
     except:
@@ -32,6 +31,7 @@ def f12():
 
     i = g()
     i.next()
+    print list(g())
     try:
         1/0
     except: