From 6b122751febaef6f6ab39346ce63681754da3b0b Mon Sep 17 00:00:00 2001
From: MerryMage <MerryMage@users.noreply.github.com>
Date: Mon, 27 Nov 2017 19:38:22 +0000
Subject: [PATCH] reg_alloc: Remove reliance on
 IR::Inst::DecrementRemainingUses

---
 src/backend_x64/reg_alloc.cpp | 26 +++++++++++++++++++++-----
 src/backend_x64/reg_alloc.h   |  5 +++++
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/backend_x64/reg_alloc.cpp b/src/backend_x64/reg_alloc.cpp
index 20eb147d..6edb7d73 100644
--- a/src/backend_x64/reg_alloc.cpp
+++ b/src/backend_x64/reg_alloc.cpp
@@ -5,6 +5,7 @@
  */
 
 #include <algorithm>
+#include <numeric>
 
 #include <xbyak.h>
 
@@ -80,7 +81,7 @@ bool HostLocInfo::IsEmpty() const {
 }
 
 bool HostLocInfo::IsLastUse() const {
-    return !is_being_used && std::all_of(values.begin(), values.end(), [](const auto& inst) { return !inst->HasUses(); });
+    return !is_being_used && current_references == 1;
 }
 
 bool HostLocInfo::ContainsValue(const IR::Inst* inst) const {
@@ -100,11 +101,25 @@ void HostLocInfo::WriteLock() {
 
 void HostLocInfo::AddValue(IR::Inst* inst) {
     values.push_back(inst);
+    total_uses += inst->UseCount();
+}
+
+void HostLocInfo::AddArgReference() {
+    current_references++;
+    ASSERT(accumulated_uses + current_references <= total_uses);
 }
 
 void HostLocInfo::EndOfAllocScope() {
-    const auto to_erase = std::remove_if(values.begin(), values.end(), [](const auto& inst) { return !inst->HasUses(); });
-    values.erase(to_erase, values.end());
+    accumulated_uses += current_references;
+    current_references = 0;
+
+    if (total_uses == accumulated_uses) {
+        values.clear();
+        accumulated_uses = 0;
+        total_uses = 0;
+    }
+
+    ASSERT(total_uses == std::accumulate(values.begin(), values.end(), size_t(0), [](size_t sum, IR::Inst* inst) { return sum + inst->UseCount(); }));
 
     is_being_used = false;
     is_scratch = false;
@@ -159,10 +174,11 @@ bool Argument::IsInMemory() const {
 std::array<Argument, 3> RegAlloc::GetArgumentInfo(IR::Inst* inst) {
     std::array<Argument, 3> ret = { Argument{*this}, Argument{*this}, Argument{*this} };
     for (size_t i = 0; i < inst->NumArgs(); i++) {
-        IR::Value arg = inst->GetArg(i);
+        const IR::Value& arg = inst->GetArg(i);
         ret[i].value = arg;
         if (!arg.IsImmediate()) {
-            arg.GetInst()->DecrementRemainingUses();
+            ASSERT_MSG(ValueLocation(arg.GetInst()), "argument must already been defined");
+            LocInfo(*ValueLocation(arg.GetInst())).AddArgReference();
         }
     }
     return ret;
diff --git a/src/backend_x64/reg_alloc.h b/src/backend_x64/reg_alloc.h
index 6a54819f..eea9887b 100644
--- a/src/backend_x64/reg_alloc.h
+++ b/src/backend_x64/reg_alloc.h
@@ -37,12 +37,17 @@ public:
 
     void AddValue(IR::Inst* inst);
 
+    void AddArgReference();
     void EndOfAllocScope();
 
 private:
     std::vector<IR::Inst*> values;
     bool is_being_used = false;
     bool is_scratch = false;
+
+    size_t current_references = 0;
+    size_t accumulated_uses = 0;
+    size_t total_uses = 0;
 };
 
 struct Argument {