Have rewrites keep track of how many stack frames are inside them

This is kind of hacky due to our method of first doing the rewrite,
and then picking the slot to rewrite to.  This means that at the time
that we create the rewrite, we don't know the exact location of the
"num_inside" counter, so we have to go back at the end and rewrite it.

We could also switch things to pick the rewrite slot first, but that
is also complicated but perhaps more subtly, since during the course
of the rewrite the chosen slot could have gotten rewritten!

This changeset adds this functionality but doesn't use it for anything.

src/asm_writing/assembler.cpp

@@ -616,12 +616,52 @@ void Assembler::sub(Immediate imm, Register reg) {
     emitArith(imm, reg, OPCODE_SUB);
 }
 
-void Assembler::inc(Register reg) {
-    UNIMPLEMENTED();
+void Assembler::incl(Indirect mem) {
+    int src_idx = mem.base.regnum;
+
+    int rex = 0;
+    if (src_idx >= 8) {
+        rex |= REX_B;
+        src_idx -= 8;
+    }
+
+    assert(src_idx >= 0 && src_idx < 8);
+
+    if (rex)
+        emitRex(rex);
+    emitByte(0xff);
+
+    assert(-0x80 <= mem.offset && mem.offset < 0x80);
+    if (mem.offset == 0) {
+        emitModRM(0b00, 0, src_idx);
+    } else {
+        emitModRM(0b01, 0, src_idx);
+        emitByte(mem.offset);
+    }
 }
 
-void Assembler::inc(Indirect mem) {
-    UNIMPLEMENTED();
+void Assembler::decl(Indirect mem) {
+    int src_idx = mem.base.regnum;
+
+    int rex = 0;
+    if (src_idx >= 8) {
+        rex |= REX_B;
+        src_idx -= 8;
+    }
+
+    assert(src_idx >= 0 && src_idx < 8);
+
+    if (rex)
+        emitRex(rex);
+    emitByte(0xff);
+
+    assert(-0x80 <= mem.offset && mem.offset < 0x80);
+    if (mem.offset == 0) {
+        emitModRM(0b00, 1, src_idx);
+    } else {
+        emitModRM(0b01, 1, src_idx);
+        emitByte(mem.offset);
+    }
 }
 
 

src/asm_writing/assembler.h

@@ -125,8 +125,8 @@ public:
 
     void add(Immediate imm, Register reg);
     void sub(Immediate imm, Register reg);
-    void inc(Register reg);
-    void inc(Indirect mem);
+    void incl(Indirect mem);
+    void decl(Indirect mem);
 
     void callq(Register reg);
     void retq();

src/asm_writing/icinfo.cpp

@@ -99,9 +99,10 @@ void ICSlotRewrite::commit(uint64_t decision_path, CommitHook* hook) {
     uint8_t* slot_start = (uint8_t*)ic->start_addr + ic_entry->idx * ic->getSlotSize();
     uint8_t* continue_point = (uint8_t*)ic->continue_addr;
 
-    hook->finishAssembly(continue_point - slot_start);
+    hook->finishAssembly(ic_entry, continue_point - slot_start);
 
     assert(assembler->isExactlyFull());
+    assert(!assembler->hasFailed());
 
     // if (VERBOSITY()) printf("Commiting to %p-%p\n", start, start + ic->slot_size);
     memcpy(slot_start, buf, ic->getSlotSize());

src/asm_writing/icinfo.h

@@ -32,10 +32,11 @@ class ICInvalidator;
 
 struct ICSlotInfo {
 public:
-    ICSlotInfo(ICInfo* ic, int idx) : ic(ic), idx(idx) {}
+    ICSlotInfo(ICInfo* ic, int idx) : ic(ic), idx(idx), num_inside(0) {}
 
     ICInfo* ic;
-    int idx;
+    int idx;        // the index inside the ic
+    int num_inside; // the number of stack frames that are currently inside this slot
 
     void clear();
 };
@@ -45,7 +46,7 @@ public:
     class CommitHook {
     public:
         virtual ~CommitHook() {}
-        virtual void finishAssembly(int fastpath_offset) = 0;
+        virtual void finishAssembly(ICSlotInfo* picked_slot, int fastpath_offset) = 0;
     };
 
 private:

src/asm_writing/rewriter.cpp

@@ -558,6 +558,29 @@ void Rewriter::_call(RewriterVar* result, bool can_call_into_python, void* func_
     // assert(!can_call_into_python);
     assert(done_guarding);
 
+    // we've been ignoring these annotations for long enough that I'm not sure they can be trusted,
+    // so just be pessimistic:
+    can_call_into_python = true;
+
+    if (can_call_into_python) {
+        if (!marked_inside_ic) {
+            // assembler->trap();
+
+            // TODO this is super hacky: we don't know the address that we want to inc/dec, since
+            // it depends on the slot that we end up picking, so just write out an arbitrary
+            // constant an we'll rewrite it later
+            // TODO if we can guarantee that the mark_addr will fit in 32 bits,
+            // we can use a more compact instruction encoding
+            mark_addr_addrs.push_back((void**)(assembler->curInstPointer() + 2));
+            assembler::Register reg = allocReg(Location::any());
+            assembler->mov(assembler::Immediate(0x1234567890abcdefL), reg);
+            assembler->incl(assembler::Indirect(reg, 0));
+
+            assertConsistent();
+            marked_inside_ic = true;
+        }
+    }
+
     // RewriterVarUsage scratch = createNewVar(Location::any());
     assembler::Register r = allocReg(assembler::R11);
 
@@ -731,16 +754,18 @@ void Rewriter::commit() {
     assert(!finished);
     initPhaseEmitting();
 
-    if (assembler->hasFailed()) {
-        static StatCounter rewriter_assemblyfail("rewriter_assemblyfail");
-        rewriter_assemblyfail.log();
+    static StatCounter rewriter_assemblyfail("rewriter_assemblyfail");
 
+    auto on_assemblyfail = [&]() {
+        rewriter_assemblyfail.log();
         this->abort();
+    };
+
+    if (assembler->hasFailed()) {
+        on_assemblyfail();
         return;
     }
 
-    finished = true;
-
     static StatCounter rewriter_commits("rewriter_commits");
     rewriter_commits.log();
 
@@ -760,7 +785,8 @@ void Rewriter::commit() {
     // at each guard in the var's `uses` list.
 
     // First: check if we're done guarding before we even begin emitting.
-    if (last_guard_action == -1) {
+
+    auto on_done_guarding = [&]() {
         done_guarding = true;
         for (RewriterVar* arg : args) {
             if (arg->next_use == arg->uses.size()) {
@@ -770,9 +796,12 @@ void Rewriter::commit() {
                 arg->locations.clear();
             }
         }
-    }
+        assertConsistent();
+    };
 
-    assertConsistent();
+    if (last_guard_action == -1) {
+        on_done_guarding();
+    }
 
     // Now, start emitting assembly; check if we're dong guarding after each.
     for (int i = 0; i < actions.size(); i++) {
@@ -780,17 +809,20 @@ void Rewriter::commit() {
 
         assertConsistent();
         if (i == last_guard_action) {
-            done_guarding = true;
-            for (RewriterVar* arg : args) {
-                if (arg->next_use == arg->uses.size()) {
-                    for (Location loc : arg->locations) {
-                        vars_by_location.erase(loc);
-                    }
-                    arg->locations.clear();
-                }
-            }
+            on_done_guarding();
         }
-        assertConsistent();
+    }
+
+    if (marked_inside_ic) {
+        // TODO this is super hacky: we don't know the address that we want to inc/dec, since
+        // it depends on the slot that we end up picking, so just write out an arbitrary
+        // constant an we'll rewrite it later
+        // TODO if we can guarantee that the mark_addr will fit in 32 bits,
+        // we can use a more compact instruction encoding
+        mark_addr_addrs.push_back((void**)(assembler->curInstPointer() + 2));
+        assembler::Register reg = allocReg(Location::any(), getReturnDestination());
+        assembler->mov(assembler::Immediate(0x1234567890abcdefL), reg);
+        assembler->decl(assembler::Indirect(reg, 0));
     }
 
 // Make sure that we have been calling bumpUse correctly.
@@ -896,10 +928,30 @@ void Rewriter::commit() {
     }
 #endif
 
+    if (assembler->hasFailed()) {
+        on_assemblyfail();
+        return;
+    }
+
+    finished = true;
+
+    // TODO: have to check that we have enough room to write the final jmp
     rewrite->commit(decision_path, this);
+
+    assert(!assembler->hasFailed());
 }
 
-void Rewriter::finishAssembly(int continue_offset) {
+void Rewriter::finishAssembly(ICSlotInfo* picked_slot, int continue_offset) {
+    if (marked_inside_ic) {
+        void* mark_addr = &picked_slot->num_inside;
+
+        // Go back and rewrite the faked constants to point to the correct address:
+        for (void** mark_addr_addr : mark_addr_addrs) {
+            assert(*mark_addr_addr == (void*)0x1234567890abcdefL);
+            *mark_addr_addr = mark_addr;
+        }
+    }
+
     assembler->jmp(assembler::JumpDestination::fromStart(continue_offset));
 
     assembler->fillWithNops();
@@ -1299,7 +1351,8 @@ 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()),
-      added_changing_action(false), last_guard_action(-1), done_guarding(false), ndecisions(0), decision_path(1) {
+      added_changing_action(false), marked_inside_ic(false), last_guard_action(-1), done_guarding(false), ndecisions(0),
+      decision_path(1) {
     initPhaseCollecting();
 
 #ifndef NDEBUG

src/asm_writing/rewriter.h

@@ -346,6 +346,9 @@ private:
         actions.emplace_back(action);
     }
     bool added_changing_action;
+    bool marked_inside_ic;
+    std::vector<void**> mark_addr_addrs;
+
     int last_guard_action;
 
     bool done_guarding;
@@ -374,7 +377,7 @@ private:
     // Do the bookkeeping to say that var is no longer in location l
     void removeLocationFromVar(RewriterVar* var, Location l);
 
-    void finishAssembly(int continue_offset) override;
+    void finishAssembly(ICSlotInfo* picked_slot, int continue_offset) override;
 
     int ndecisions;
     uint64_t decision_path;