Register owned attributes from rearrangeArguments

The Rewriter will automatically refcount objects represented as RewriterVar's,
but if you store one as an attribute somewhere, it can't keep track of that.
So add a registerOwnedAttribute() function that lets you say that there is
a reference that lives in an attribute (stored in a memory offset) of another RewriterVar.

from_cpython/Lib/test/test_decimal.py

 # expected: reffail
-# - takes too long; leaks refs
+# - takes too long
 # Copyright (c) 2004 Python Software Foundation.
 # All rights reserved.
 

from_cpython/Lib/test/test_file2k.py

-# expected: reffail
-# - negative ref
 import sys
 import os
 import unittest

src/asm_writing/rewriter.cpp

@@ -1229,8 +1229,7 @@ std::vector<Location> Rewriter::getDecrefLocations() {
             if (l.type == Location::Scratch) {
                 // convert to stack based location because later on we may not know the offset of the scratch area from
                 // the SP.
-                assert(indirectFor(l).offset % 8 == 0);
-                decref_infos.emplace_back(Location::Stack, indirectFor(l).offset / 8);
+                decref_infos.emplace_back(Location::Stack, indirectFor(l).offset);
             } else if (l.type == Location::Register) {
                 // CSRs shouldn't be getting allocated, and we should only be calling this at a callsite:
                 RELEASE_ASSERT(0, "we shouldn't be trying to decref anything in a register");
@@ -1238,6 +1237,20 @@ std::vector<Location> Rewriter::getDecrefLocations() {
                 RELEASE_ASSERT(0, "not implemented");
         }
     }
+
+    for (auto&& p : owned_attrs) {
+        RewriterVar* var = p.first;
+
+        assert(var->locations.size() == 1);
+        Location l = *var->locations.begin();
+
+        assert(l.type == Location::Scratch);
+
+        int offset1 = indirectFor(l).offset;
+        int offset2 = p.second;
+        decref_infos.emplace_back(Location::StackIndirect, offset1, offset2);
+    }
+
     return decref_infos;
 }
 
@@ -1323,6 +1336,24 @@ bool RewriterVar::needsDecref() {
     return reftype == RefType::OWNED && !this->refHandedOff();
 }
 
+void RewriterVar::registerOwnedAttr(int byte_offset) {
+    rewriter->addAction([=]() {
+        auto p = std::make_pair(this, byte_offset);
+        assert(!this->rewriter->owned_attrs.count(p));
+        this->rewriter->owned_attrs.insert(p);
+        this->bumpUse();
+    }, { this }, ActionType::NORMAL);
+}
+
+void RewriterVar::deregisterOwnedAttr(int byte_offset) {
+    rewriter->addAction([=]() {
+        auto p = std::make_pair(this, byte_offset);
+        assert(this->rewriter->owned_attrs.count(p));
+        this->rewriter->owned_attrs.erase(p);
+        this->bumpUse();
+    }, { this }, ActionType::NORMAL);
+}
+
 void RewriterVar::bumpUse() {
     rewriter->assertPhaseEmitting();
 
@@ -1351,6 +1382,11 @@ void RewriterVar::releaseIfNoUses() {
 void Rewriter::commit() {
     STAT_TIMER(t0, "us_timer_rewriter", 10);
 
+    // The rewriter could add decrefs here, but for now let's make the user add them explicitly
+    // and then call deregisterOwnedAttr().  Making people call it explicitly reduces the chances
+    // of bugs only in the exceptional path, from forgetting to call deregisterOwnedAttr.
+    assert(!owned_attrs.size() && "missing a call to deregisterOwnedAttr");
+
     assert(!finished);
     initPhaseEmitting();
 
@@ -1438,7 +1474,7 @@ void Rewriter::commit() {
         // add increfs if required
         for (auto&& var : actions[i].consumed_refs) {
             if (var->refHandedOff()) {
-                // if this action is the one which the variable gets handed off we don't need todo anything
+                // if this action is the one which the variable gets handed off we don't need to do anything
                 assert(var->last_refconsumed_numuses > 0 && var->last_refconsumed_numuses <= var->uses.size());
                 int last_used_action_id = var->uses[var->last_refconsumed_numuses - 1];
                 if (last_used_action_id == i)

src/asm_writing/rewriter.h

@@ -23,6 +23,7 @@
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallSet.h"
 
 #include "asm_writing/assembler.h"
@@ -184,6 +185,12 @@ public:
 
     void refUsed();
 
+    // registerOwnedAttr tells the refcounter that a certain memory location holds a pointer
+    // to an owned reference.  This must be paired with a call to deregisterOwnedAttr
+    // Call these right before emitting the store (for register) or decref (for deregister).
+    void registerOwnedAttr(int byte_offset);
+    void deregisterOwnedAttr(int byte_offset);
+
 
     template <typename Src, typename Dst> inline RewriterVar* getAttrCast(int offset, Location loc = Location::any());
 
@@ -344,11 +351,13 @@ private:
         }
     };
 
+private:
     // This needs to be the first member:
     RegionAllocator allocator;
 
-    // Increfs that we thought of doing during the guarding phase that we were able to put off.
-    std::vector<std::pair<RewriterVar*, int>> pending_increfs;
+    // The rewriter has refcounting logic for handling owned RewriterVars.  If we own memory inside those RewriterVars,
+    // we need to register that via registerOwnedAttr, which ends up here:
+    llvm::DenseSet<std::pair<RewriterVar*, int /* offset */>> owned_attrs;
 
 protected:
     // Allocates `bytes` bytes of data.  The allocation will get freed when the rewriter gets freed.

src/asm_writing/types.h

@@ -15,6 +15,8 @@
 #ifndef PYSTON_ASMWRITING_TYPES_H
 #define PYSTON_ASMWRITING_TYPES_H
 
+#include <climits>
+
 #include "core/common.h"
 
 namespace pyston {
@@ -185,6 +187,8 @@ public:
         Stack,
         Scratch, // stack location, relative to the scratch start
 
+        StackIndirect, // A location like $rsp[offset1][offset2]
+
         // For representing constants that fit in 32-bits, that can be encoded as immediates
         AnyReg,        // special type for use when specifying a location as a destination
         None,          // special type that represents the lack of a location, ex where a "ret void" gets returned
@@ -199,11 +203,19 @@ public:
         // also valid if type==XMMRegister
         int32_t regnum;
         // only valid if type==Stack; this is the offset from bottom of the original frame.
-        // ie argument #6 will have a stack_offset of 0, #7 will have a stack offset of 8, etc
+        // ie argument #6 will have a stack_offset of 0, #7 will have a stack offset of 8, etc.
+        // Measured in bytes
         int32_t stack_offset;
-        // only valid if type == Scratch; offset from the beginning of the scratch area
+        // only valid if type == Scratch; offset from the beginning of the scratch area.
+        // Measured in bytes
         int32_t scratch_offset;
 
+        // Only valid if type == StackIndirect:
+        struct {
+            int16_t stack_first_offset;
+            int16_t stack_second_offset;
+        };
+
         int32_t _data;
     };
 
@@ -212,6 +224,12 @@ public:
     Location& operator=(const Location& r) = default;
 
     constexpr Location(LocationType type, int32_t data) : type(type), _data(data) {}
+    Location(LocationType type, int64_t offset1, int64_t offset2)
+        : type(type), stack_first_offset(offset1), stack_second_offset(offset2) {
+        assert(type == StackIndirect);
+        assert(SHRT_MIN <= offset1 && offset1 <= SHRT_MAX);
+        assert(SHRT_MIN <= offset2 && offset2 <= SHRT_MAX);
+    }
 
     constexpr Location(assembler::Register reg) : type(Register), regnum(reg.regnum) {}
 

src/codegen/unwinding.cpp

@@ -584,7 +584,14 @@ public:
                         Box* b = NULL;
                         if (l.type == Location::Stack) {
                             unw_word_t sp = get_cursor_sp(cursor);
-                            b = ((Box**)sp)[l.stack_offset];
+                            assert(l.stack_offset % 8 == 0);
+                            b = ((Box**)sp)[l.stack_offset / 8];
+                        } else if (l.type == Location::StackIndirect) {
+                            unw_word_t sp = get_cursor_sp(cursor);
+                            assert(l.stack_first_offset % 8 == 0);
+                            Box** b_ptr = ((Box***)sp)[l.stack_first_offset / 8];
+                            assert(l.stack_second_offset % 8 == 0);
+                            b = b_ptr[l.stack_second_offset / 8];
                         } else if (l.type == Location::Register) {
                             RELEASE_ASSERT(0, "untested");
                             // This branch should never get hit since we shouldn't generate Register locations,

src/runtime/objmodel.cpp

@@ -4096,8 +4096,10 @@ public:
 
 void decrefOargs(RewriterVar* oargs, bool* oargs_owned, int num_oargs) {
     for (int i = 0; i < num_oargs; i++) {
-        if (oargs_owned[i])
+        if (oargs_owned[i]) {
+            oargs->deregisterOwnedAttr(i * sizeof(Box*));
             oargs->getAttr(i * sizeof(Box*))->setType(RefType::OWNED);
+        }
     }
 }
 
@@ -4352,12 +4354,16 @@ void rearrangeArgumentsInternal(ParamReceiveSpec paramspec, const ParamNames* pa
                 if (varargs_idx == 2)
                     rewrite_args->arg3 = varargs_val;
                 if (varargs_idx >= 3) {
-                    rewrite_args->args->setAttr(
-                        (varargs_idx - 3) * sizeof(Box*), varargs_val,
-                        (is_owned ? RewriterVar::SetattrType::HANDED_OFF : RewriterVar::SetattrType::UNKNOWN));
                     if (is_owned) {
+                        rewrite_args->args->registerOwnedAttr((varargs_idx - 3) * sizeof(Box*));
+
+                        rewrite_args->args->setAttr((varargs_idx - 3) * sizeof(Box*), varargs_val,
+                                                    RewriterVar::SetattrType::HANDED_OFF);
+
                         oargs_owned[varargs_idx - 3] = true;
                         varargs_val->refConsumed();
+                    } else {
+                        rewrite_args->args->setAttr((varargs_idx - 3) * sizeof(Box*), varargs_val);
                     }
                 }
             }

test/tests/oargs_decref.py

+# Regression test: make sure that, if an exception is thrown after
+# rearrange arguments, that we still decref any owned oargs.
+
+def f(a, b, c, d, e, *args):
+    print a, b, c, d, e, args
+    1/a
+
+for i in xrange(-100, 100):
+    try:
+        f(i, 0, 0, 0, 0, 0)
+    except ZeroDivisionError:
+        pass