Reduce malloc() calls in rewriter

Allocating things in-line, using malloc-friendly data structures, etc.

src/asm_writing/icinfo.cpp

@@ -56,17 +56,15 @@ void ICSlotInfo::clear() {
     ic->clear(this);
 }
 
-ICSlotRewrite::ICSlotRewrite(ICInfo* ic, const char* debug_name) : ic(ic), debug_name(debug_name) {
-    buf = (uint8_t*)malloc(ic->getSlotSize());
-    assembler = new Assembler(buf, ic->getSlotSize());
-    assembler->nop();
+ICSlotRewrite::ICSlotRewrite(ICInfo* ic, const char* debug_name)
+    : ic(ic), debug_name(debug_name), buf((uint8_t*)malloc(ic->getSlotSize())), assembler(buf, ic->getSlotSize()) {
+    assembler.nop();
 
     if (VERBOSITY() >= 4)
         printf("starting %s icentry\n", debug_name);
 }
 
 ICSlotRewrite::~ICSlotRewrite() {
-    delete assembler;
     free(buf);
 }
 
@@ -109,7 +107,7 @@ void ICSlotRewrite::commit(CommitHook* hook) {
     if (!do_commit)
         return;
 
-    assert(!assembler->hasFailed());
+    assert(!assembler.hasFailed());
 
     for (int i = 0; i < dependencies.size(); i++) {
         ICInvalidator* invalidator = dependencies[i].first;

src/asm_writing/icinfo.h

@@ -21,6 +21,7 @@
 
 #include "llvm/IR/CallingConv.h"
 
+#include "asm_writing/assembler.h"
 #include "asm_writing/types.h"
 
 namespace pyston {
@@ -53,12 +54,13 @@ public:
 
 private:
     ICInfo* ic;
-    assembler::Assembler* assembler;
     const char* debug_name;
 
     uint8_t* buf;
 
-    std::vector<std::pair<ICInvalidator*, int64_t>> dependencies;
+    assembler::Assembler assembler;
+
+    llvm::SmallVector<std::pair<ICInvalidator*, int64_t>, 4> dependencies;
 
     ICSlotInfo* ic_entry;
 
@@ -66,7 +68,7 @@ public:
     ICSlotRewrite(ICInfo* ic, const char* debug_name);
     ~ICSlotRewrite();
 
-    assembler::Assembler* getAssembler() { return assembler; }
+    assembler::Assembler* getAssembler() { return &assembler; }
     int getSlotSize();
     int getScratchRspOffset();
     int getScratchSize();

src/asm_writing/rewriter.cpp

@@ -185,9 +185,11 @@ void Rewriter::ConstLoader::moveImmediate(uint64_t val, assembler::Register dst_
 assembler::Register Rewriter::ConstLoader::findConst(uint64_t val, bool& found_value) {
     assert(rewriter->phase_emitting);
 
-    auto it = constToVar.find(val);
-    if (it != constToVar.end()) {
-        RewriterVar* var = it->second;
+    for (auto& p : consts) {
+        if (p.first != val)
+            continue;
+
+        RewriterVar* var = p.second;
         for (Location l : var->locations) {
             if (l.type == Location::Register) {
                 found_value = true;
@@ -216,11 +218,11 @@ assembler::Register Rewriter::ConstLoader::loadConst(uint64_t val, Location othe
     assert(rewriter->phase_emitting);
 
     bool found_value = false;
-    assembler::Register reg = findConst(val, found_value);
+    assembler::Register /*reg = findConst(val, found_value);
     if (found_value)
-        return reg;
+        return reg;*/
 
-    reg = rewriter->allocReg(Location::any(), otherThan);
+        reg = rewriter->allocReg(Location::any(), otherThan);
     if (tryLea(val, reg))
         return reg;
 
@@ -711,10 +713,15 @@ void Rewriter::_trap() {
 RewriterVar* Rewriter::loadConst(int64_t val, Location dest) {
     STAT_TIMER(t0, "us_timer_rewriter", 10);
 
-    RewriterVar*& const_loader_var = const_loader.constToVar[val];
-    if (!const_loader_var) {
-        const_loader_var = createNewConstantVar(val);
+    for (auto& p : const_loader.consts) {
+        if (p.first != val)
+            continue;
+
+        return p.second;
     }
+
+    RewriterVar* const_loader_var = createNewConstantVar(val);
+    const_loader.consts.push_back(std::make_pair(val, const_loader_var));
     return const_loader_var;
 }
 
@@ -1099,11 +1106,11 @@ void Rewriter::commit() {
     for (int i = 0; i < live_outs.size(); i++) {
         live_outs[i]->uses.push_back(actions.size());
     }
-    for (RewriterVar* var : vars) {
+    for (RewriterVar& var : vars) {
         // Add a use for every constant. This helps make constants available for the lea stuff
         // But since "spilling" a constant has no cost, it shouldn't add register pressure.
-        if (var->is_constant) {
-            var->uses.push_back(actions.size());
+        if (var.is_constant) {
+            var.uses.push_back(actions.size());
         }
     }
 
@@ -1176,22 +1183,22 @@ void Rewriter::commit() {
 // Make sure that we have been calling bumpUse correctly.
 // All uses should have been accounted for, other than the live outs
 #ifndef NDEBUG
-    for (RewriterVar* var : vars) {
+    for (RewriterVar& var : vars) {
         int num_as_live_out = 0;
         for (RewriterVar* live_out : live_outs) {
-            if (live_out == var) {
+            if (live_out == &var) {
                 num_as_live_out++;
             }
         }
-        assert(var->next_use + num_as_live_out + (var->is_constant ? 1 : 0) == var->uses.size());
+        assert(var.next_use + num_as_live_out + (var.is_constant ? 1 : 0) == var.uses.size());
     }
 #endif
 
     assert(live_out_regs.size() == live_outs.size());
 
-    for (RewriterVar* var : vars) {
-        if (var->is_constant) {
-            var->bumpUse();
+    for (RewriterVar& var : vars) {
+        if (var.is_constant) {
+            var.bumpUse();
         }
     }
 
@@ -1686,9 +1693,8 @@ void Rewriter::removeLocationFromVar(RewriterVar* var, Location l) {
 RewriterVar* Rewriter::createNewVar() {
     assertPhaseCollecting();
 
-    RewriterVar* var = new RewriterVar(this);
-    vars.push_back(var);
-    return var;
+    vars.emplace_back(this);
+    return &vars.back();
 }
 
 RewriterVar* Rewriter::createNewConstantVar(uint64_t val) {

src/asm_writing/rewriter.h

@@ -15,6 +15,7 @@
 #ifndef PYSTON_ASMWRITING_REWRITER_H
 #define PYSTON_ASMWRITING_REWRITER_H
 
+#include <deque>
 #include <map>
 #include <memory>
 #include <tuple>
@@ -101,6 +102,8 @@ public:
 
     bool operator!=(const Location rhs) const { return !(*this == rhs); }
 
+    bool operator<(const Location& rhs) const { return this->asInt() < rhs.asInt(); }
+
     uint64_t asInt() const { return (int)type + ((uint64_t)_data << 4); }
 
     void dump() const;
@@ -224,7 +227,7 @@ public:
 private:
     Rewriter* rewriter;
 
-    std::unordered_set<Location> locations;
+    std::set<Location> locations;
     bool isInLocation(Location l);
 
     // uses is a vector of the indices into the Rewriter::actions vector
@@ -317,7 +320,7 @@ protected:
         // Loads the constant into any register or if already in a register just return it
         assembler::Register loadConst(uint64_t val, Location otherThan = Location::any());
 
-        std::unordered_map<uint64_t, RewriterVar*> constToVar;
+        std::vector<std::pair<uint64_t, RewriterVar*>> consts;
     };
 
 
@@ -326,7 +329,7 @@ protected:
     ICSlotInfo* picked_slot;
 
     ConstLoader const_loader;
-    std::vector<RewriterVar*> vars;
+    std::deque<RewriterVar> vars;
 
     const Location return_location;
 
@@ -346,11 +349,11 @@ protected:
     void assertPhaseEmitting() {}
 #endif
 
-    std::vector<int> live_out_regs;
+    llvm::SmallVector<int, 8> live_out_regs;
 
     LocMap<RewriterVar*> vars_by_location;
-    std::vector<RewriterVar*> args;
-    std::vector<RewriterVar*> live_outs;
+    llvm::SmallVector<RewriterVar*, 8> args;
+    llvm::SmallVector<RewriterVar*, 8> live_outs;
 
     Rewriter(std::unique_ptr<ICSlotRewrite> rewrite, int num_args, const std::vector<int>& live_outs);
 
@@ -443,15 +446,22 @@ protected:
 
     void assertConsistent() {
 #ifndef NDEBUG
-        for (RewriterVar* var : vars) {
-            for (Location l : var->locations) {
-                assert(vars_by_location[l] == var);
+        for (RewriterVar& var : vars) {
+            for (Location l : var.locations) {
+                assert(vars_by_location[l] == &var);
             }
         }
         for (std::pair<Location, RewriterVar*> p : vars_by_location.getAsMap()) {
             assert(p.second != NULL);
             if (p.second != LOCATION_PLACEHOLDER) {
-                assert(std::find(vars.begin(), vars.end(), p.second) != vars.end());
+                bool found = false;
+                for (auto& v : vars) {
+                    if (&v == p.second) {
+                        found = true;
+                        break;
+                    }
+                }
+                assert(found);
                 assert(p.second->locations.count(p.first) == 1);
             }
         }
@@ -471,10 +481,6 @@ public:
         if (!finished)
             this->abort();
         assert(finished);
-
-        for (RewriterVar* var : vars) {
-            delete var;
-        }
     }
 
     Location getReturnDestination();