generator: clean up abandonment generators

in order to break cycles we need to traverse all generator owned object at yields.
The interpreter is fine because it stores all objects inside the vregs so they will already get visited.
For the llvm jit pass all owned object to the yield call as vararg argument.

We will currently leak (move objects to gc.garbage) when a generator is involved in a cycle because our PyGen_NeedsFinalizing() is much more conservative.

from_cpython/Lib/test/test_abc.py

-# expected: reffail
-# - generator abandonment
 # Copyright 2007 Google, Inc. All Rights Reserved.
 # Licensed to PSF under a Contributor Agreement.
 

from_cpython/Lib/test/test_abstract_numbers.py

-# expected: reffail
-# - generator abandonment
 """Unit tests for numbers.py."""
 
 import math

from_cpython/Lib/test/test_ast.py

-# expected: reffail
-# - generator abandonment
 import sys, itertools, unittest
 from test import test_support
 import ast

from_cpython/Lib/test/test_collections.py

 # expected: reffail
-# - generator abandonment
 import unittest, doctest, operator
 import inspect
 from test import test_support

from_cpython/Lib/test/test_contextlib.py

-# expected: reffail
-# - generator abandonment
 """Unit tests for contextlib.py, and other context managers."""
 
 import sys

from_cpython/Lib/test/test_cpickle.py

-# expected: reffail
-# - generator abandonment
 import cPickle
 import cStringIO
 import io

from_cpython/Lib/test/test_decimal.py

 # expected: reffail
-# - generator abandonment
 # Copyright (c) 2004 Python Software Foundation.
 # All rights reserved.
 

from_cpython/Lib/test/test_file.py

-# expected: reffail
-# - generator abandonment
 # NOTE: this file tests the new `io` library backported from Python 3.x.
 # Similar tests for the builtin file object can be found in test_file2k.py.
 

from_cpython/Lib/test/test_float.py

-# expected: reffail
-# - generator abandonment
 import unittest, struct
 import os
 from test import test_support

from_cpython/Lib/test/test_fractions.py

-# expected: reffail
-# - generator abandonment
 """Tests for Lib/fractions.py."""
 
 from decimal import Decimal

from_cpython/Lib/test/test_locale.py

-# expected: reffail
-# - generator abandonment
 from test.test_support import run_unittest, verbose
 import unittest
 import locale

from_cpython/Lib/test/test_mailbox.py

-# expected: reffail
-# - generator abandonment
 import os
 import sys
 import time

from_cpython/Lib/test/test_pickle.py

-# expected: reffail
-# - generator abandonment
 import pickle
 from cStringIO import StringIO
 

from_cpython/Lib/test/test_pickletools.py

-# expected: reffail
-# - generator abandonment
 import pickle
 import pickletools
 from test import test_support

from_cpython/Lib/test/test_pkgutil.py

-# expected: reffail
-# - generator abandonment
 from test.test_support import run_unittest
 import unittest
 import sys

from_cpython/Lib/test/test_quopri.py

-# expected: reffail
-# - generator abandonment
 from test import test_support
 import unittest
 

from_cpython/Lib/test/test_robotparser.py

-# expected: reffail
-# - generator abandonment
 import unittest, StringIO, robotparser
 from test import test_support
 from urllib2 import urlopen, HTTPError

from_cpython/Lib/test/test_typechecks.py

-# expected: reffail
-# - generator abandonment
 """Unit tests for __instancecheck__ and __subclasscheck__."""
 
 import unittest

from_cpython/Lib/test/test_urllib.py

-# expected: reffail
-# - generator abandonment
 """Regresssion tests for urllib"""
 
 import urllib

from_cpython/Lib/test/test_urllib2.py

-# expected: reffail
-# - generator abandonment
 import unittest
 from test import test_support
 

from_cpython/Lib/test/test_urllib2net.py

-# expected: reffail
-# - generator abandonment
 import unittest
 from test import test_support
 from test.test_urllib2 import sanepathname2url

from_cpython/Lib/test/test_urlparse.py

-# expected: reffail
-# - generator abandonment
 from test import test_support
 import unittest
 import urlparse

from_cpython/Lib/test/test_with.py

-# expected: reffail
-# - generator abandonment
 """Unit tests for the with statement specified in PEP 343."""
 
 

from_cpython/Lib/test/test_xrange.py

-# expected: reffail
-# - generator abandonment
 # Python test set -- built-in functions
 
 import test.test_support, unittest

src/codegen/ast_interpreter.cpp

@@ -2065,6 +2065,11 @@ extern "C" Box* astInterpretDeoptFromASM(FunctionMetadata* md, AST_expr* after_e
         memset(vregs, 0, sizeof(Box*) * num_vregs);
     }
 
+    // We need to remove the old python frame created in the LLVM tier otherwise we would have a duplicate frame because
+    // the interpreter will set the new state before executing the first statement.
+    RELEASE_ASSERT(cur_thread_state.frame_info == frame_state.frame_info, "");
+    cur_thread_state.frame_info = frame_state.frame_info->back;
+
     ASTInterpreter interpreter(md, vregs, num_vregs, frame_state.frame_info);
 
     for (const auto& p : *frame_state.locals) {
@@ -2132,11 +2137,6 @@ extern "C" Box* astInterpretDeoptFromASM(FunctionMetadata* md, AST_expr* after_e
         assert(starting_statement);
     }
 
-    // We need to remove the old python frame created in the LLVM tier otherwise we would have a duplicate frame because
-    // the interpreter will set the new state before executing the first statement.
-    RELEASE_ASSERT(cur_thread_state.frame_info == frame_state.frame_info, "");
-    cur_thread_state.frame_info = interpreter.getFrameInfo()->back;
-
     Box* v = ASTInterpreter::execute(interpreter, start_block, starting_statement);
     return v ? v : incref(None);
 }

src/codegen/baseline_jit.cpp

@@ -537,7 +537,9 @@ std::vector<RewriterVar*> JitFragmentWriter::emitUnpackIntoArray(RewriterVar* v,
 
 RewriterVar* JitFragmentWriter::emitYield(RewriterVar* v) {
     RewriterVar* generator = getInterp()->getAttr(ASTInterpreterJitInterface::getGeneratorOffset());
-    auto rtn = call(false, (void*)yield, generator, v)->setType(RefType::OWNED);
+    static_assert(sizeof(llvm::ArrayRef<Box*>) == sizeof(void*) * 2,
+                  "we pass two 0ul args to initalize the llvm::ArrayRef");
+    auto rtn = call(false, (void*)yield, generator, v, imm(0ul), imm(0ul))->setType(RefType::OWNED);
     v->refConsumed();
     return rtn;
 }

src/codegen/irgen.cpp

@@ -655,8 +655,10 @@ static void emitBBs(IRGenState* irstate, TypeAnalysis* types, const OSREntryDesc
             if (source->getScopeInfo()->takesClosure())
                 names.insert(source->getInternedStrings().get(PASSED_CLOSURE_NAME));
 
-            if (source->is_generator)
+            if (source->is_generator) {
+                assert(0 && "not sure if this is correct");
                 names.insert(source->getInternedStrings().get(PASSED_GENERATOR_NAME));
+            }
 
             for (const InternedString& s : names) {
                 // printf("adding guessed phi for %s\n", s.c_str());

src/codegen/irgen/irgenerator.cpp

@@ -1540,14 +1540,26 @@ private:
         assert(generator);
         ConcreteCompilerVariable* convertedGenerator = generator->makeConverted(emitter, generator->getBoxType());
 
-
         CompilerVariable* value = node->value ? evalExpr(node->value, unw_info) : emitter.getNone();
         ConcreteCompilerVariable* convertedValue = value->makeConverted(emitter, value->getBoxType());
 
+        std::vector<llvm::Value*> args;
+        args.push_back(convertedGenerator->getValue());
+        args.push_back(convertedValue->getValue());
+        args.push_back(getConstantInt(0, g.i32)); // the refcounting inserter handles yields specially and adds all
+                                                  // owned objects as additional arguments to it
+
+        // put the yield call at the beginning of a new basic block to make it easier for the refcounting inserter.
+        llvm::BasicBlock* yield_block = emitter.createBasicBlock("yield_block");
+        emitter.getBuilder()->CreateBr(yield_block);
+        emitter.setCurrentBasicBlock(yield_block);
+
+        // we do a capi call because it makes it easier for the refcounter to replace the instruction
         llvm::Instruction* rtn
-            = emitter.createCall2(unw_info, g.funcs.yield, convertedGenerator->getValue(), convertedValue->getValue());
+            = emitter.createCall(unw_info, g.funcs.yield_capi, args, CAPI, getNullPtr(g.llvm_value_type_ptr));
         emitter.refConsumed(convertedValue->getValue(), rtn);
         emitter.setType(rtn, RefType::OWNED);
+        emitter.setNullable(rtn, true);
 
         return new ConcreteCompilerVariable(UNKNOWN, rtn);
     }

src/codegen/irgen/refcounts.cpp

@@ -746,8 +746,18 @@ void RefcountTracker::addRefcounts(IRGenState* irstate) {
     }
 
     std::vector<llvm::InvokeInst*> invokes;
+    std::vector<llvm::CallInst*> yields;
     for (auto&& II : llvm::inst_range(f)) {
         llvm::Instruction* inst = &II;
+
+        // is this a yield?
+        if (llvm::isa<llvm::CallInst>(inst)) {
+            llvm::CallInst* call = llvm::cast<llvm::CallInst>(inst);
+            if (call->getCalledValue() == g.funcs.yield_capi)
+                yields.push_back(call);
+        }
+
+        // invoke specific code
         if (!rt->vars.count(inst))
             continue;
         if (auto ii = dyn_cast<InvokeInst>(inst))
@@ -1135,6 +1145,38 @@ void RefcountTracker::addRefcounts(IRGenState* irstate) {
         }
     }
 
+    // yields need to get handled specially
+    // we pass all object which we own at the point of the yield call to the yield so that we can traverse them in
+    // tp_traverse.
+    // we have to create a new call instruction because we can't add arguments to an existing call instruction
+    for (auto&& old_yield : yields) {
+        auto&& state = states[old_yield->getParent()];
+        assert(old_yield->getNumArgOperands() == 3);
+        llvm::Value* yield_value = old_yield->getArgOperand(1);
+
+        llvm::SmallVector<llvm::Value*, 8> args;
+        args.push_back(old_yield->getArgOperand(0)); // generator
+        args.push_back(yield_value);                 // value
+        args.push_back(NULL); // num live values. we replace it with the actual number of varargs after inserting them
+        // we can just traverse state.ending_refs because when generating the yield we make sure that it's at the start
+        // of the BB.
+        for (auto ref : state.ending_refs) {
+            if (rt->vars.lookup(ref.first).reftype == RefType::OWNED) {
+                if (yield_value != ref.first) // ignore this value because yield steals it!
+                    args.push_back(ref.first);
+            }
+        }
+        int num_live_values = args.size() - 3;
+        if (num_live_values == 0)
+            continue; // nothing to do
+
+        args[2] = getConstantInt(num_live_values, g.i32); // replace the dummy value the actual amount
+
+        llvm::CallInst* new_yield = llvm::CallInst::Create(g.funcs.yield_capi, args, llvm::Twine(), old_yield);
+        old_yield->replaceAllUsesWith(new_yield);
+        old_yield->eraseFromParent();
+    }
+
     long us = _t.end();
     static StatCounter us_refcounting("us_compiling_irgen_refcounting");
     us_refcounting.log(us);

src/codegen/runtime_hooks.cpp

@@ -225,7 +225,7 @@ void initGlobalFuncs(GlobalState& g) {
     GET(importStar);
     GET(repr);
     GET(exceptionMatches);
-    GET(yield);
+    GET(yield_capi);
     GET(getiterHelper);
     GET(hasnext);
     GET(apply_slice);

src/codegen/runtime_hooks.h

@@ -38,7 +38,8 @@ struct GlobalFuncs {
         *makePendingCalls, *setFrameExcInfo;
     llvm::Value* getattr, *getattr_capi, *setattr, *delattr, *delitem, *delGlobal, *nonzero, *binop, *compare,
         *augbinop, *unboxedLen, *getitem, *getitem_capi, *getclsattr, *getGlobal, *setitem, *unaryop, *import,
-        *importFrom, *importStar, *repr, *exceptionMatches, *yield, *getiterHelper, *hasnext, *setGlobal, *apply_slice;
+        *importFrom, *importStar, *repr, *exceptionMatches, *yield_capi, *getiterHelper, *hasnext, *setGlobal,
+        *apply_slice;
 
     llvm::Value* unpackIntoArray, *raiseAttributeError, *raiseAttributeErrorStr, *raiseAttributeErrorCapi,
         *raiseAttributeErrorStrCapi, *raiseNotIterableError, *raiseIndexErrorStr, *raiseIndexErrorStrCapi,

src/core/types.h

@@ -23,6 +23,7 @@
 #include <stddef.h>
 #include <vector>
 
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/iterator_range.h"
 #include "Python.h"
 

src/runtime/generator.cpp

@@ -100,7 +100,8 @@ void generatorEntry(BoxedGenerator* g) {
 
             // call body of the generator
             BoxedFunctionBase* func = g->function;
-            KEEP_ALIVE(func);
+            // unnecessary because the generator owns g->function
+            // KEEP_ALIVE(func);
 
             Box** args = g->args ? &g->args->elts[0] : nullptr;
             auto r = callCLFunc<ExceptionStyle::CXX, NOT_REWRITABLE>(func->md, nullptr, func->md->numReceivedArgs(),
@@ -331,8 +332,25 @@ Box* generatorHasnext(Box* s) {
     return boxBool(generatorHasnextUnboxed(s));
 }
 
+extern "C" Box* yield_capi(BoxedGenerator* obj, STOLEN(Box*) value, int num_live_values, ...) noexcept {
+    try {
+        llvm::SmallVector<Box*, 8> live_values;
+        live_values.reserve(num_live_values);
+        va_list ap;
+        va_start(ap, num_live_values);
+        for (int i = 0; i < num_live_values; ++i) {
+            live_values.push_back(va_arg(ap, Box*));
+        }
+        va_end(ap);
+
+        return yield(obj, value, live_values);
+    } catch (ExcInfo e) {
+        setCAPIException(e);
+        return NULL;
+    }
+}
 
-extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value) {
+extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value, llvm::ArrayRef<Box*> live_values) {
     STAT_TIMER(t0, "us_timer_generator_switching", 0);
 
     assert(obj->cls == generator_cls);
@@ -350,9 +368,11 @@ extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value) {
     // reset current frame to the caller tops frame --> removes the frame the generator added
     cur_thread_state.frame_info = self->top_caller_frame_info;
     obj->paused_frame_info = generator_frame_info;
+    obj->live_values = live_values;
     swapContext(&self->context, self->returnContext, 0);
     FrameInfo* top_new_caller_frame_info = (FrameInfo*)cur_thread_state.frame_info;
     obj->paused_frame_info = NULL;
+    obj->live_values = llvm::ArrayRef<Box*>();
 
     // the caller of the generator can change between yield statements that means we can't just restore the top of the
     // frame to the point before the yield instead we have to update it.
@@ -380,7 +400,6 @@ extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value) {
     return r;
 }
 
-
 extern "C" BoxedGenerator* createGenerator(BoxedFunctionBase* function, Box* arg1, Box* arg2, Box* arg3, Box** args) {
     assert(function);
     assert(function->cls == function_cls);
@@ -492,8 +511,12 @@ extern "C" int PyGen_NeedsFinalizing(PyGenObject* gen) noexcept {
     // CPython has some optimizations for not needing to finalize generators that haven't exited, but
     // which are guaranteed to not need any special cleanups.
     // For now just say anything still in-progress needs finalizing.
-    return (bool)self->paused_frame_info;
+    if (!(bool)self->paused_frame_info)
+        return false;
 
+    return true;
+// TODO: is this safe? probably not...
+// return self->paused_frame_info->stmt->type == AST_TYPE::Invoke;
 #if 0
     int i;
     PyFrameObject* f = gen->gi_frame;
@@ -513,17 +536,97 @@ extern "C" int PyGen_NeedsFinalizing(PyGenObject* gen) noexcept {
 #endif
 }
 
+static PyObject* generator_close(PyGenObject* gen, PyObject* args) noexcept {
+    try {
+        return generatorClose((Box*)gen);
+    } catch (ExcInfo e) {
+        setCAPIException(e);
+        return NULL;
+    }
+}
+
+static void generator_del(PyObject* self) noexcept {
+    PyObject* res;
+    PyObject* error_type, *error_value, *error_traceback;
+    BoxedGenerator* gen = (BoxedGenerator*)self;
+
+    // Pyston change:
+    // if (gen->gi_frame == NULL || gen->gi_frame->f_stacktop == NULL)
+    if (!gen->paused_frame_info)
+        /* Generator isn't paused, so no need to close */
+        return;
+
+    /* Temporarily resurrect the object. */
+    assert(self->ob_refcnt == 0);
+    self->ob_refcnt = 1;
+
+    /* Save the current exception, if any. */
+    PyErr_Fetch(&error_type, &error_value, &error_traceback);
+
+    // Pyston change:
+    // res = gen_close(gen, NULL);
+    res = generator_close((PyGenObject*)gen, NULL);
+
+    if (res == NULL)
+        PyErr_WriteUnraisable(self);
+    else
+        Py_DECREF(res);
+
+    /* Restore the saved exception. */
+    PyErr_Restore(error_type, error_value, error_traceback);
+
+    /* Undo the temporary resurrection; can't use DECREF here, it would
+     * cause a recursive call.
+     */
+    assert(self->ob_refcnt > 0);
+    if (--self->ob_refcnt == 0)
+        return; /* this is the normal path out */
+
+    /* close() resurrected it!  Make it look like the original Py_DECREF
+     * never happened.
+     */
+    {
+        Py_ssize_t refcnt = self->ob_refcnt;
+        _Py_NewReference(self);
+        self->ob_refcnt = refcnt;
+    }
+    assert(PyType_IS_GC(self->cls) && _Py_AS_GC(self)->gc.gc_refs != _PyGC_REFS_UNTRACKED);
+
+    /* If Py_REF_DEBUG, _Py_NewReference bumped _Py_RefTotal, so
+     * we need to undo that. */
+    _Py_DEC_REFTOTAL;
+/* If Py_TRACE_REFS, _Py_NewReference re-added self to the object
+ * chain, so no more to do there.
+ * If COUNT_ALLOCS, the original decref bumped tp_frees, and
+ * _Py_NewReference bumped tp_allocs:  both of those need to be
+ * undone.
+ */
+#ifdef COUNT_ALLOCS
+    --self->ob_type->tp_frees;
+    --self->ob_type->tp_allocs;
+#endif
+}
+
 static void generator_dealloc(BoxedGenerator* self) noexcept {
     assert(isSubclass(self->cls, generator_cls));
 
     // Hopefully this never happens:
     assert(!self->running);
 
-    // I don't think this should get hit currently because we don't properly collect the cycle that this would
-    // represent:
-    ASSERT(!self->paused_frame_info, "Can't clean up a generator that is currently paused");
+    _PyObject_GC_UNTRACK(self);
+
+    if (self->weakreflist != NULL)
+        PyObject_ClearWeakRefs(self);
+
+    _PyObject_GC_TRACK(self);
 
-    PyObject_GC_UnTrack(self);
+    if (self->paused_frame_info) {
+        Py_TYPE(self)->tp_del(self);
+        if (self->ob_refcnt > 0)
+            return; /* resurrected.  :( */
+    }
+
+    _PyObject_GC_UNTRACK(self);
 
     freeGeneratorStack(self);
 
@@ -547,8 +650,6 @@ static void generator_dealloc(BoxedGenerator* self) noexcept {
     Py_CLEAR(self->exception.value);
     Py_CLEAR(self->exception.traceback);
 
-    PyObject_ClearWeakRefs(self);
-
     self->cls->tp_free(self);
 }
 
@@ -561,6 +662,10 @@ static int generator_traverse(BoxedGenerator* self, visitproc visit, void* arg)
             return r;
     }
 
+    for (auto v : self->live_values) {
+        Py_VISIT(v);
+    }
+
     int numArgs = self->function->md->numReceivedArgs();
     if (numArgs > 3) {
         numArgs -= 3;
@@ -610,5 +715,6 @@ void setupGenerator() {
 
     generator_cls->freeze();
     generator_cls->tp_iter = PyObject_SelfIter;
+    generator_cls->tp_del = generator_del; // don't do giveAttr("__del__") because it should not be visible from python
 }
 }

src/runtime/generator.h

@@ -29,7 +29,8 @@ void setupGenerator();
 void generatorEntry(BoxedGenerator* g);
 Context* getReturnContextForGeneratorFrame(void* frame_addr);
 
-extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value);
+extern "C" Box* yield(BoxedGenerator* obj, STOLEN(Box*) value, llvm::ArrayRef<Box*> live_values = {});
+extern "C" Box* yield_capi(BoxedGenerator* obj, STOLEN(Box*) value, int num_live_values = 0, ...) noexcept;
 extern "C" BoxedGenerator* createGenerator(BoxedFunctionBase* function, Box* arg1, Box* arg2, Box* arg3, Box** args);
 }
 

src/runtime/inline/link_forcer.cpp

@@ -101,7 +101,7 @@ void force() {
     FORCE(repr);
     FORCE(str);
     FORCE(exceptionMatches);
-    FORCE(yield);
+    FORCE(yield_capi);
     FORCE(getiterHelper);
     FORCE(hasnext);
     FORCE(apply_slice);

src/runtime/types.h

@@ -1257,6 +1257,8 @@ public:
     FrameInfo* paused_frame_info; // The FrameInfo the generator was on when it called yield (or NULL if the generator
                                   // hasn't started or has exited).
 
+    llvm::ArrayRef<Box*> live_values;
+
 #if STAT_TIMERS
     StatTimer* prev_stack;
     StatTimer my_timer;

test/tests/dash_m.py

-# expected: reffail
-# - generator abandonment
 import os
 import sys
 import subprocess

test/tests/decorated_func_line.py

-# expected: reffail
-# - generator abandonment
 def wrapper(f):
     return f
 

test/tests/deopt_generator_tests.py

-# expected: reffail
 # skip-if: '-O' in EXTRA_JIT_ARGS or '-n'  in EXTRA_JIT_ARGS
 # statcheck: 4 <= noninit_count('num_deopt') < 50
 # statcheck: 1 <= stats["num_osr_exits"] <= 2

test/tests/exec_directory.py

-# expected: reffail
-# - generator abandonment
 import sys
 import os
 import subprocess

test/tests/generator_abandonment.py

-# expected: reffail
-
 print any(i == 5 for i in xrange(10))
 
 

test/tests/generator_collection_finally.py

-# expected: fail
 # We currently don't call finalizers when destroying a generator.
 def G():
     try:

test/tests/generator_collection_running.py

-# expected: reffail
 # This test checks if generators which get started but haven't yet stopped (=not raisen a StopIteration exc, etc)
 # get freed when there aren't any references to the generators left.
 

test/tests/generator_cycle.py

-# expected: reffail
 # Test a generator cycle involving an unfinished generator.
+import gc
 
 def f():
     g = (i in (None, g) for i in xrange(2))
     print g.next()
 print f()
+gc.collect()
+# print gc.garbage # XX pyston will put this into garbage currently
+
+def f():
+   g = None
+   def G():
+       c = g
+       g
+       yield 1
+       yield 2
+       yield 3
+   g = G()
+   g.next()
+
+print f()
+gc.collect()
+# print gc.garbage # XX pyston will put this into garbage currently

test/tests/generator_nonrunning_collection.py

-# expected: reffail
 def g():
     l1 = [1]
     l2 = [2]

test/tests/generator_threads.py

-# expected: reffail
 # Pass a started generator to two different threads, and make them both
 # try to run it at the same time.
 

test/tests/generators.py

@@ -8,7 +8,7 @@ g1 = G1()
 for i in range(5):
     print g1.next()
 print g1.__name__
-
+print hasattr(g1, "__del__")
 
 
 def G2():

test/tests/nonzero_exceptions.py

@@ -42,13 +42,10 @@ try:
 except MyException, e:
     print e
 
-# TODO: reenable this once generator abandonment is working again
-"""
 try:
     print list(1 for i in range(5) if C(7))
 except MyException, e:
     print e
-"""
 
 try:
     print 1 if C(8) else 0