Fix memory leak of RewriterVars when the rewrite doesn't commit

I'm not sure I like the current approach, but it seems to be working.

src/asm_writing/rewriter.cpp

@@ -325,7 +325,7 @@ assembler::Immediate RewriterVar::tryGetAsImmediate(bool* is_immediate) {
     return assembler::Immediate((uint64_t)0);
 }
 
-assembler::Register RewriterVar::getInReg(Location dest) {
+assembler::Register RewriterVar::getInReg(Location dest, bool allow_constant_in_reg) {
     assert(dest.type == Location::Register || dest.type == Location::AnyReg);
 
     // assembler::Register reg = var->rewriter->allocReg(l);
@@ -333,8 +333,10 @@ assembler::Register RewriterVar::getInReg(Location dest) {
     // return reg;
     assert(locations.size());
 #ifndef NDEBUG
-    for (Location l : locations) {
-        ASSERT(l.type != Location::Constant, "why do you want this in a register?");
+    if (!allow_constant_in_reg) {
+        for (Location l : locations) {
+            ASSERT(l.type != Location::Constant, "why do you want this in a register?");
+        }
     }
 #endif
 
@@ -363,15 +365,20 @@ assembler::Register RewriterVar::getInReg(Location dest) {
 
     assert(locations.size() == 1);
     Location l(*locations.begin());
-    assert(l.type == Location::Scratch || l.type == Location::Stack);
-
 
     assembler::Register reg = rewriter->allocReg(dest);
     assert(rewriter->vars_by_location.count(reg) == 0);
 
-    assembler::Indirect mem = rewriter->indirectFor(l);
-    rewriter->assembler->mov(mem, reg);
+    if (l.type == Location::Constant) {
+        rewriter->assembler->mov(assembler::Immediate(l.constant_val), reg);
+    } else if (l.type == Location::Scratch || l.type == Location::Stack) {
+        assembler::Indirect mem = rewriter->indirectFor(l);
+        rewriter->assembler->mov(mem, reg);
+    } else {
+        abort();
+    }
     rewriter->addLocationToVar(this, reg);
+
     return reg;
 }
 
@@ -639,7 +646,29 @@ RewriterVarUsage Rewriter::call(bool can_call_into_python, void* func_addr, std:
     return RewriterVarUsage(var);
 }
 
+void Rewriter::abort() {
+    assert(!finished);
+    finished = true;
+
+    // This feels hacky: are we really guaranteed to find all of the things we need to delete?
+    std::unordered_set<RewriterVar*> found;
+
+    for (const auto& p : vars_by_location) {
+        found.insert(p.second);
+    }
+    for (const auto v : live_outs) {
+        found.insert(v);
+    }
+
+    for (auto v : found) {
+        delete v;
+    }
+}
+
 void Rewriter::commit() {
+    assert(!finished);
+    finished = true;
+
     static StatCounter rewriter2_commits("rewriter2_commits");
     rewriter2_commits.log();
 
@@ -687,7 +716,7 @@ void Rewriter::finishAssembly(int continue_offset) {
 
 void Rewriter::commitReturning(RewriterVarUsage usage) {
     // assert(usage.var->isInLocation(getReturnDestination()));
-    usage.var->getInReg(getReturnDestination());
+    usage.var->getInReg(getReturnDestination(), true /* allow_constant_in_reg */);
 
     usage.setDoneUsing();
     commit();
@@ -931,6 +960,10 @@ TypeRecorder* Rewriter::getTypeRecorder() {
 Rewriter::Rewriter(ICSlotRewrite* rewrite, int num_args, const std::vector<int>& live_outs)
     : rewrite(rewrite), assembler(rewrite->getAssembler()), return_location(rewrite->returnRegister()),
       done_guarding(false), ndecisions(0), decision_path(1) {
+#ifndef NDEBUG
+    start_vars = RewriterVar::nvars;
+#endif
+    finished = false;
     // assembler->trap();
 
     for (int i = 0; i < num_args; i++) {
@@ -994,4 +1027,8 @@ Rewriter* Rewriter::createRewriter(void* rtn_addr, int num_args, const char* deb
 RewriterVarUsage RewriterVarUsage::addUse() {
     return RewriterVarUsage(var);
 }
+
+#ifndef NDEBUG
+int RewriterVar::nvars = 0;
+#endif
 }

src/asm_writing/rewriter.h

@@ -186,7 +186,7 @@ private:
     // Gets a copy of this variable in a register, spilling/reloading if necessary.
     // TODO have to be careful with the result since the interface doesn't guarantee
     // that the register will still contain your value when you go to use it
-    assembler::Register getInReg(Location l = Location::any());
+    assembler::Register getInReg(Location l = Location::any(), bool allow_constant_in_reg = false);
     assembler::XMMRegister getInXMMReg(Location l = Location::any());
 
     // If this is an immediate, try getting it as one
@@ -195,14 +195,24 @@ private:
     void dump();
 
 public:
+#ifndef NDEBUG
+    static int nvars;
+#endif
     void incUse();
     void decUse();
 
     RewriterVar(Rewriter* rewriter, Location location) : rewriter(rewriter), num_uses(0) {
+#ifndef NDEBUG
+        nvars++;
+#endif
         assert(rewriter);
         locations.insert(location);
     }
 
+#ifndef NDEBUG
+    ~RewriterVar() { nvars--; }
+#endif
+
     friend class RewriterVarUsage;
     friend class Rewriter;
 };
@@ -216,6 +226,11 @@ private:
 
     bool done_guarding;
 
+    bool finished; // committed or aborted
+#ifndef NDEBUG
+    int start_vars;
+#endif
+
     std::vector<int> live_out_regs;
 
     std::unordered_map<Location, RewriterVar*> vars_by_location;
@@ -258,6 +273,14 @@ public:
     // This should be called exactly once for each argument
     RewriterVarUsage getArg(int argnum);
 
+    ~Rewriter() {
+        if (!finished)
+            this->abort();
+        assert(finished);
+
+        ASSERT(RewriterVar::nvars == start_vars, "%d %d", RewriterVar::nvars, start_vars);
+    }
+
     Location getReturnDestination();
 
     bool isDoneGuarding() { return done_guarding; }
@@ -275,6 +298,7 @@ public:
     RewriterVarUsage allocateAndCopyPlus1(RewriterVarUsage first_elem, RewriterVarUsage rest, int n_rest);
     void deallocateStack(int nbytes);
 
+    void abort();
     void commit();
     void commitReturning(RewriterVarUsage rtn);
 

src/core/common.h

@@ -56,10 +56,10 @@ void _printStacktrace();
 #define RELEASE_ASSERT(condition, fmt, ...)                                                                            \
     do {                                                                                                               \
         if (!(condition)) {                                                                                            \
-            fprintf(stderr, __FILE__ ":" STRINGIFY(__LINE__) ": %s: Assertion `" #condition "' failed: " fmt "\n",     \
-                    __PRETTY_FUNCTION__, ##__VA_ARGS__);                                                               \
-            pyston::_printStacktrace();                                                                                \
-            abort();                                                                                                   \
+            ::fprintf(stderr, __FILE__ ":" STRINGIFY(__LINE__) ": %s: Assertion `" #condition "' failed: " fmt "\n",   \
+                      __PRETTY_FUNCTION__, ##__VA_ARGS__);                                                             \
+            ::pyston::_printStacktrace();                                                                              \
+            ::abort();                                                                                                 \
         }                                                                                                              \
     } while (false)
 #ifndef NDEBUG

src/gc/gc_alloc.cpp

@@ -18,6 +18,7 @@
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <dlfcn.h>
 #include <stdint.h>
 #include <sys/mman.h>
 
@@ -44,16 +45,36 @@ extern "C" void gc_compat_free(void* ptr) {
     gc_free(ptr);
 }
 
+int nallocs = 0;
+bool recursive = false;
+
 // We may need to hook malloc as well:
-/*
+#if 0
 extern "C" void* malloc(size_t sz) {
-    abort();
+    static void *(*libc_malloc)(size_t) = (void* (*)(size_t))dlsym(RTLD_NEXT, "malloc");
+    nallocs++;
+    void* r = libc_malloc(sz);;
+    if (!recursive && nallocs > 4000000) {
+        recursive = true;
+        printf("malloc'd: %p\n", r);
+        raise(SIGTRAP);
+        recursive = false;
+    }
+    return r;
 }
 
 extern "C" void free(void* p) {
-    abort();
+    static void (*libc_free)(void*) = (void (*)(void*))dlsym(RTLD_NEXT, "free");
+    if (!recursive && nallocs > 4000000) {
+        recursive = true;
+        printf("free: %p\n", p);
+        raise(SIGTRAP);
+        recursive = false;
+    }
+    nallocs--;
+    libc_free(p);
 }
-*/
+#endif
 
 } // namespace gc
 } // namespace pyston

src/runtime/objmodel.cpp

@@ -1493,6 +1493,8 @@ extern "C" bool nonzero(Box* obj) {
     } else if (obj->cls == none_cls) {
         if (rewriter.get()) {
             r_obj.setDoneUsing();
+            RewriterVarUsage b = rewriter->loadConst(0, rewriter->getReturnDestination());
+            rewriter->commitReturning(std::move(b));
         }
         return false;
     }

test/tests/sre_compile_test.py

@@ -9,3 +9,5 @@ print r.match("ac")
 print r.match("abc").groups()
 for i in xrange(100000):
     r.match("abbc").groups()
+    if i % 1000 == 0:
+        print i