fixup! RegAlloc: Simplify code

src/core/arm/jit_x64/interface.cpp

@@ -88,13 +88,14 @@ struct BlockOfCode : Gen::XCodeBlock {
 
 namespace JitX64 {
 
-static Gen::RunJittedCode run_jit = {};
+struct ARM_Jit::Impl {
-static Gen::BlockOfCode block_of_code = {};
+    Gen::RunJittedCode run_jit = {};
-static JitX64 compiler { &block_of_code };
+    Gen::BlockOfCode block_of_code = {};
+    JitX64 compiler{ &block_of_code };
+};
 
-ARM_Jit::ARM_Jit(PrivilegeMode initial_mode) {
+ARM_Jit::ARM_Jit(PrivilegeMode initial_mode) : impl(std::make_unique<Impl>()), state(Common::make_unique<JitState>()) {
     ASSERT_MSG(initial_mode == PrivilegeMode::USER32MODE, "Unimplemented");
-    state = Common::make_unique<JitState>();
     ClearCache();
 }
 
@@ -173,9 +174,9 @@ void ARM_Jit::ExecuteInstructions(int num_instructions) {
         else
             state->cpu_state.Reg[15] &= 0xfffffffc;
 
-        u8 *ptr = compiler.GetBB(state->cpu_state.Reg[15], state->cpu_state.TFlag, EFlag);
+        u8 *ptr = impl->compiler.GetBB(state->cpu_state.Reg[15], state->cpu_state.TFlag, EFlag);
 
-        unsigned ticks_executed = run_jit.CallCode(state.get(), ptr, num_instructions, state->cpu_state.Reg[15]);
+        unsigned ticks_executed = impl->run_jit.CallCode(state.get(), ptr, num_instructions, state->cpu_state.Reg[15]);
         num_instructions -= ticks_executed;
         AddTicks(ticks_executed);
     } while (!reschedule && num_instructions > 0);
@@ -223,14 +224,14 @@ void ARM_Jit::PrepareReschedule() {
 }
 
 void ARM_Jit::ClearCache() {
-    compiler.ClearCache();
+    impl->compiler.ClearCache();
-    block_of_code.ClearCodeSpace();
+    impl->block_of_code.ClearCodeSpace();
     state->cpu_state.instruction_cache.clear();
 }
 
 void ARM_Jit::FastClearCache() {
-    compiler.ClearCache();
+    impl->compiler.ClearCache();
-    block_of_code.ResetCodePtr();
+    impl->block_of_code.ResetCodePtr();
     state->cpu_state.instruction_cache.clear();
 }
 

src/core/arm/jit_x64/interface.h

@@ -44,7 +44,10 @@ public:
     void FastClearCache();
 
 private:
-    bool reschedule;
+    struct Impl;
+
+    std::unique_ptr<Impl> impl;
+    bool reschedule = false;
     std::unique_ptr<JitState> state;
 };
 

src/core/arm/jit_x64/reg_alloc.cpp

@@ -12,22 +12,24 @@
 
 namespace JitX64 {
 
-static const std::map<Gen::X64Reg, size_t> x64_reg_to_index = {
+const std::map<Gen::X64Reg, size_t> RegAlloc::x64_reg_to_index = {
-    { Gen::RAX, 0 },
+    // Ordered such that the caller saved registers are on top,
-    { Gen::RBX, 1 },
+    // and registers often flushed at the bottom.
-    { Gen::RCX, 2 },
+    { Gen::R10, 0 },
-    { Gen::RDX, 3 },
+    { Gen::R11, 1 },
-    { Gen::RBP, 4 },
+    { Gen::RBX, 2 },
-    { Gen::RSI, 5 },
+    { Gen::RBP, 3 },
-    { Gen::RDI, 6 },
+    { Gen::R12, 4 },
-    { Gen::RSP, 7 },
+    { Gen::R13, 5 },
-    { Gen::R8,  8 },
+    { Gen::R14, 6 },
-    { Gen::R9,  9 },
+    { Gen::RAX, 7 },
-    { Gen::R10, 10 },
+    { Gen::R9,  8 },
-    { Gen::R11, 11 },
+    { Gen::R8,  9 },
-    { Gen::R12, 12 },
+    { Gen::RSI, 10 },
-    { Gen::R13, 13 },
+    { Gen::RDX, 11 },
-    { Gen::R14, 14 },
+    { Gen::RDI, 12 },
+    { Gen::RCX, 13 },
+    { Gen::RSP, 14 },
 };
 
 constexpr Gen::X64Reg jit_state_reg = Gen::R15;
@@ -50,40 +52,41 @@ void RegAlloc::Init(Gen::XEmitter* emitter) {
     code = emitter;
 
     for (size_t i = 0; i < arm_gpr.size(); i++) {
-        arm_gpr[i].locked = false;
+        const ArmReg arm_reg = static_cast<ArmReg>(i);
-        arm_gpr[i].location = MJitStateCpuReg(static_cast<ArmReg>(i));
+        arm_gpr[i] = { arm_reg, MJitStateCpuReg(arm_reg), false };
     }
 
-    for (size_t i = 0; i < x64_gpr.size(); i++) {
+    for (const auto& entry : x64_reg_to_index) {
-        x64_gpr[i].locked = false;
+        const Gen::X64Reg x64_reg = entry.first;
-        x64_gpr[i].state = X64State::State::Free;
+        const size_t i = entry.second;
-        x64_gpr[i].arm_reg = ArmReg::INVALID_REG;
+        x64_gpr[i] = { x64_reg, false, X64State::State::Free, ArmReg::INVALID_REG };
     }
 }
 
 void RegAlloc::FlushX64(Gen::X64Reg x64_reg) {
-    const size_t x64_index = x64_reg_to_index.at(x64_reg);
+    X64State& x64_state = GetState(x64_reg);
-    X64State& state = x64_gpr[x64_index];
 
-    ASSERT(!state.locked);
+    ASSERT(!x64_state.locked);
 
-    switch (state.state) {
+    switch (x64_state.state) {
     case X64State::State::Free:
     case X64State::State::Temp:
-        state.state = X64State::State::Free;
+        x64_state.state = X64State::State::Free;
         break;
     case X64State::State::CleanArmReg: {
-        state.state = X64State::State::Free;
+        ArmState& arm_state = GetState(x64_state.arm_reg);
-        ArmState& arm_state = arm_gpr[static_cast<unsigned>(state.arm_reg)];
+        AssertStatesConsistent(x64_state, arm_state);
-        arm_state.location = MJitStateCpuReg(state.arm_reg);
+
+        // We ignore the value in the x64 register since it's not dirty.
+        x64_state.state = X64State::State::Free;
+        arm_state.location = MJitStateCpuReg(x64_state.arm_reg);
         break;
     }
     case X64State::State::DirtyArmReg: {
-        ArmState& arm_state = arm_gpr[static_cast<unsigned>(state.arm_reg)];
+        AssertStatesConsistent(x64_state, GetState(x64_state.arm_reg));
-        ASSERT(arm_state.location.IsSimpleReg());
+
-        ASSERT(arm_state.location.GetSimpleReg() == x64_reg);
+        // Flush the value in the x64 register back into ARMul_State since it's dirty.
-        FlushArm(state.arm_reg);
+        FlushArm(x64_state.arm_reg);
-        ASSERT(state.state == X64State::State::Free);
         break;
     }
     default:
@@ -91,113 +94,101 @@ void RegAlloc::FlushX64(Gen::X64Reg x64_reg) {
         break;
     }
 
-    ASSERT(state.state == X64State::State::Free);
+    ASSERT(x64_state.state == X64State::State::Free);
-    ASSERT(!state.locked);
+    ASSERT(!x64_state.locked);
 }
 
 void RegAlloc::LockX64(Gen::X64Reg x64_reg) {
-    X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
+    X64State& x64_state = GetState(x64_reg);
+
+    ASSERT(!x64_state.locked && x64_state.state == X64State::State::Free);
 
-    ASSERT(!x64_state.locked);
-    ASSERT(x64_state.state == X64State::State::Free);
     x64_state.locked = true;
     x64_state.state = X64State::State::UserManuallyLocked;
 }
 
 void RegAlloc::UnlockX64(Gen::X64Reg x64_reg) {
-    X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
+    X64State& x64_state = GetState(x64_reg);
+
+    ASSERT(x64_state.locked && x64_state.state == X64State::State::UserManuallyLocked);
 
-    ASSERT(x64_state.locked);
-    ASSERT(x64_state.state == X64State::State::UserManuallyLocked);
     x64_state.locked = false;
     x64_state.state = X64State::State::Free;
 }
 
 void RegAlloc::FlushArm(ArmReg arm_reg) {
-    ASSERT(IsValidArmReg(arm_reg));
+    ArmState& arm_state = GetState(arm_reg);
 
-    ArmState& arm_state = arm_gpr[static_cast<unsigned>(arm_reg)];
     ASSERT(!arm_state.locked);
-    if (!arm_state.location.IsSimpleReg()) {
-        return;
-    }
 
-    Gen::X64Reg x64_reg = GetX64For(arm_reg);
+    if (!arm_state.IsInX64Register())
-    X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
+        return; // Nothing to do
+
+    X64State& x64_state = GetState(GetX64For(arm_reg));
 
     ASSERT(!x64_state.locked);
-    ASSERT(x64_state.state == X64State::State::CleanArmReg || x64_state.state == X64State::State::DirtyArmReg);
-    ASSERT(x64_state.arm_reg == arm_reg);
 
     if (x64_state.state == X64State::State::DirtyArmReg) {
-        code->MOV(32, MJitStateCpuReg(arm_reg), R(x64_reg));
+        code->MOV(32, MJitStateCpuReg(arm_reg), R(x64_state.x64_reg));
     }
+
     x64_state.state = X64State::State::Free;
     arm_state.location = MJitStateCpuReg(arm_reg);
 }
 
 Gen::OpArg RegAlloc::LockArmForRead(ArmReg arm_reg) {
-    ASSERT(IsValidArmReg(arm_reg) && arm_reg != ArmReg::PC); // Not valid for R15 (cannot read from it)
+    ASSERT(arm_reg != ArmReg::PC); // Not valid for R15 (cannot read from it)
 
-    ArmState& arm_state = arm_gpr[static_cast<unsigned>(arm_reg)];
+    ArmState& arm_state = GetState(arm_reg);
+
+    if (arm_state.IsInX64Register()) {
+        X64State& x64_state = GetState(GetX64For(arm_reg));
+        ASSERT(!x64_state.locked);
+        x64_state.locked = true;
+    }
 
     ASSERT(!arm_state.locked);
     arm_state.locked = true;
 
-    if (arm_state.location.IsSimpleReg()) {
-        Gen::X64Reg x64_reg = arm_state.location.GetSimpleReg();
-        X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
-        ASSERT(!x64_state.locked);
-        ASSERT(x64_state.state == X64State::State::CleanArmReg || x64_state.state == X64State::State::DirtyArmReg);
-        ASSERT(x64_state.arm_reg == arm_reg);
-
-        x64_state.locked = true;
-    }
-
     return arm_state.location;
 }
 
 Gen::OpArg RegAlloc::LockArmForWrite(ArmReg arm_reg) {
-    ASSERT(IsValidArmReg(arm_reg)); // Valid for R15 (write-only)
+    // Valid for R15 (write-only)
 
-    ArmState& arm_state = arm_gpr[static_cast<unsigned>(arm_reg)];
+    ArmState& arm_state = GetState(arm_reg);
 
-    ASSERT(!arm_state.locked);
+    if (arm_state.IsInX64Register()) {
-    arm_state.locked = true;
+        X64State& x64_state = GetState(GetX64For(arm_reg));
-
-    if (arm_state.location.IsSimpleReg()) {
-        Gen::X64Reg x64_reg = arm_state.location.GetSimpleReg();
-        X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
         ASSERT(!x64_state.locked);
-        ASSERT(x64_state.state == X64State::State::CleanArmReg || x64_state.state == X64State::State::DirtyArmReg);
-        ASSERT(x64_state.arm_reg == arm_reg);
-
         x64_state.locked = true;
         x64_state.state = X64State::State::DirtyArmReg;
     }
 
+    ASSERT(!arm_state.locked);
+    arm_state.locked = true;
+
     return arm_state.location;
 }
 
 Gen::X64Reg RegAlloc::BindArmToX64(ArmReg arm_reg, bool load) {
-    ArmState& arm_state = arm_gpr[static_cast<unsigned>(arm_reg)];
+    ArmState& arm_state = GetState(arm_reg);
 
     ASSERT(!arm_state.locked);
-    arm_state.locked = true;
 
-    if (arm_state.location.IsSimpleReg()) {
+    if (arm_state.IsInX64Register()) {
-        const Gen::X64Reg x64_reg = arm_state.location.GetSimpleReg();
+        X64State& x64_state = GetState(GetX64For(arm_reg));
-        X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
         ASSERT(!x64_state.locked);
-        ASSERT(x64_state.state == X64State::State::CleanArmReg || x64_state.state == X64State::State::DirtyArmReg);
+        arm_state.locked = true;
-        ASSERT(x64_state.arm_reg == arm_reg);
-
         x64_state.locked = true;
-        return x64_reg;
+        return x64_state.x64_reg;
     }
 
     const Gen::X64Reg x64_reg = AllocReg();
-    X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
+    X64State& x64_state = GetState(x64_reg);
+
+    ASSERT(!x64_state.locked && x64_state.state == X64State::State::Free);
+
     x64_state.locked = true;
     x64_state.arm_reg = arm_reg;
     x64_state.state = X64State::State::CleanArmReg;
@@ -205,21 +196,20 @@ Gen::X64Reg RegAlloc::BindArmToX64(ArmReg arm_reg, bool load) {
     if (load)
         code->MOV(32, R(x64_reg), MJitStateCpuReg(arm_reg));
 
+    arm_state.locked = true;
     arm_state.location = R(x64_reg);
 
     return x64_reg;
 }
 
 Gen::X64Reg RegAlloc::BindArmForRead(ArmReg arm_reg) {
-    ASSERT(IsValidArmReg(arm_reg) && arm_reg != ArmReg::PC); // Not valid for R15 (cannot read from it)
+    ASSERT(arm_reg != ArmReg::PC); // Not valid for R15 (cannot read from it)
 
-    const Gen::X64Reg x64_reg = BindArmToX64(arm_reg, true);
+    return BindArmToX64(arm_reg, true);
-
-    return x64_reg;
 }
 
 Gen::X64Reg RegAlloc::BindArmForWrite(ArmReg arm_reg) {
-    ASSERT(IsValidArmReg(arm_reg)); // Valid for R15 (we're not reading from it)
+    // Valid for R15 (we're not reading from it)
 
     const Gen::X64Reg x64_reg = BindArmToX64(arm_reg, false);
 
@@ -229,55 +219,45 @@ Gen::X64Reg RegAlloc::BindArmForWrite(ArmReg arm_reg) {
 }
 
 void RegAlloc::UnlockArm(ArmReg arm_reg) {
-    ASSERT(IsValidArmReg(arm_reg));
+    ArmState& arm_state = GetState(arm_reg);
 
-    ArmState& arm_state = arm_gpr[static_cast<unsigned>(arm_reg)];
+    if (arm_state.IsInX64Register()) {
+        X64State& x64_state = GetState(GetX64For(arm_reg));
+        ASSERT(x64_state.locked);
+        x64_state.locked = false;
+    }
 
     ASSERT(arm_state.locked);
     arm_state.locked = false;
-
-    if (arm_state.location.IsSimpleReg()) {
-        const Gen::X64Reg x64_reg = arm_state.location.GetSimpleReg();
-        X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
-        ASSERT(x64_state.locked);
-        ASSERT(x64_state.state == X64State::State::CleanArmReg || x64_state.state == X64State::State::DirtyArmReg);
-        ASSERT(x64_state.arm_reg == arm_reg);
-
-        x64_state.locked = false;
-    }
 }
 
 void RegAlloc::MarkDirty(ArmReg arm_reg) {
-    const ArmState& arm_state = arm_gpr[static_cast<unsigned>(arm_reg)];
+    const ArmState& arm_state = GetState(arm_reg);
 
-    ASSERT(arm_state.locked);
+    ASSERT(arm_state.locked && arm_state.IsInX64Register());
-    ASSERT(arm_state.location.IsSimpleReg());
 
-    const Gen::X64Reg x64_reg = arm_state.location.GetSimpleReg();
+    X64State& x64_state = GetState(GetX64For(arm_reg));
-    X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
 
     ASSERT(x64_state.locked);
-    ASSERT(x64_state.state == X64State::State::CleanArmReg || x64_state.state == X64State::State::DirtyArmReg);
-    ASSERT(x64_state.arm_reg == arm_reg);
 
     x64_state.state = X64State::State::DirtyArmReg;
 }
 
 void RegAlloc::FlushEverything() {
-    for (auto i : x64_reg_to_index) {
+    for (const X64State& x64_state : x64_gpr) {
-        X64State& x64_state = x64_gpr[i.second];
+        ASSERT(!x64_state.locked);
-        ASSERT(!x64_state.locked || x64_state.state == X64State::State::UserManuallyLocked);
+        FlushX64(x64_state.x64_reg);
-        FlushX64(i.first);
         ASSERT(x64_state.state == X64State::State::Free);
     }
 }
 
 void RegAlloc::FlushABICallerSaved() {
-    for (auto i : x64_reg_to_index) {
+    for (const X64State& x64_state : x64_gpr) {
-        X64State& x64_state = x64_gpr[i.second];
+        if (!ABI_ALL_CALLER_SAVED[x64_state.x64_reg])
+            continue;
         if (x64_state.state != X64State::State::UserManuallyLocked) {
             ASSERT(!x64_state.locked);
-            FlushX64(i.first);
+            FlushX64(x64_state.x64_reg);
             ASSERT(x64_state.state == X64State::State::Free);
         } else {
             ASSERT(x64_state.locked);
@@ -287,29 +267,62 @@ void RegAlloc::FlushABICallerSaved() {
     ASSERT(!ABI_ALL_CALLER_SAVED[JitStateReg()]);
 }
 
-Gen::X64Reg RegAlloc::GetX64For(ArmReg arm_reg) {
+RegAlloc::X64State& RegAlloc::GetState(Gen::X64Reg x64_reg) {
-    const ArmState& arm_state = arm_gpr[static_cast<unsigned>(arm_reg)];
+    ASSERT(x64_reg_to_index.find(x64_reg) != x64_reg_to_index.end());
+    const size_t x64_index = x64_reg_to_index.at(x64_reg);
+    ASSERT(x64_gpr[x64_index].x64_reg == x64_reg);
+    return x64_gpr[x64_index];
+}
 
-    ASSERT(arm_state.location.IsSimpleReg());
+const RegAlloc::X64State& RegAlloc::GetState(Gen::X64Reg x64_reg) const {
+    ASSERT(x64_reg_to_index.find(x64_reg) != x64_reg_to_index.end());
+    const size_t x64_index = x64_reg_to_index.at(x64_reg);
+    ASSERT(x64_gpr[x64_index].x64_reg == x64_reg);
+    return x64_gpr[x64_index];
+}
+
+RegAlloc::ArmState& RegAlloc::GetState(ArmReg arm_reg) {
+    ASSERT(IsValidArmReg(arm_reg));
+    const size_t arm_index = static_cast<size_t>(arm_reg);
+    ASSERT(arm_gpr[arm_index].arm_reg == arm_reg);
+    return arm_gpr[arm_index];
+}
+
+const RegAlloc::ArmState& RegAlloc::GetState(ArmReg arm_reg) const {
+    ASSERT(IsValidArmReg(arm_reg));
+    const size_t arm_index = static_cast<size_t>(arm_reg);
+    ASSERT(arm_gpr[arm_index].arm_reg == arm_reg);
+    return arm_gpr[arm_index];
+}
+
+void RegAlloc::AssertStatesConsistent(const X64State& x64_state, const ArmState& arm_state) const {
+    ASSERT(arm_state.locked == x64_state.locked);
+
+    ASSERT(arm_state.IsInX64Register());
+    ASSERT(arm_state.location.GetSimpleReg() == x64_state.x64_reg);
+
+    ASSERT(x64_state.arm_reg == arm_state.arm_reg);
+    ASSERT(x64_state.state == X64State::State::CleanArmReg || x64_state.state == X64State::State::DirtyArmReg);
+}
+
+Gen::X64Reg RegAlloc::GetX64For(ArmReg arm_reg) const {
+    const ArmState& arm_state = GetState(arm_reg);
+
+    ASSERT(arm_state.IsInX64Register());
 
     const Gen::X64Reg x64_reg = arm_state.location.GetSimpleReg();
-    const X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
 
-    ASSERT(x64_state.state == X64State::State::CleanArmReg || x64_state.state == X64State::State::DirtyArmReg);
+    AssertStatesConsistent(GetState(x64_reg), arm_state);
-    ASSERT(x64_state.arm_reg == arm_reg);
 
     return x64_reg;
 }
 
-bool RegAlloc::IsBoundToX64(ArmReg arm_reg) {
-    const ArmState& arm_state = arm_gpr[static_cast<unsigned>(arm_reg)];
-
-    return arm_state.location.IsSimpleReg();
-}
-
 Gen::X64Reg RegAlloc::AllocTemp() {
     const Gen::X64Reg x64_reg = AllocReg();
-    X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
+    X64State& x64_state = GetState(x64_reg);
+
+    ASSERT(!x64_state.locked && x64_state.state == X64State::State::Free);
+
     x64_state.locked = true;
     x64_state.state = X64State::State::Temp;
 
@@ -317,7 +330,7 @@ Gen::X64Reg RegAlloc::AllocTemp() {
 }
 
 void RegAlloc::UnlockTemp(Gen::X64Reg x64_reg) {
-    X64State& x64_state = x64_gpr[x64_reg_to_index.at(x64_reg)];
+    X64State& x64_state = GetState(x64_reg);
 
     ASSERT(x64_state.locked);
     ASSERT(x64_state.state == X64State::State::Temp);
@@ -326,49 +339,41 @@ void RegAlloc::UnlockTemp(Gen::X64Reg x64_reg) {
     x64_state.state = X64State::State::Free;
 }
 
-void RegAlloc::AssertNoLocked() {
+void RegAlloc::AssertNoLocked() const {
-    for (size_t i = 0; i < arm_gpr.size(); i++) {
+    ASSERT(std::all_of(arm_gpr.begin(), arm_gpr.end(), [&](const ArmState& arm_state) {
-        ArmState& arm_state = arm_gpr[i];
+        if (arm_state.IsInX64Register()) {
-        ASSERT(!arm_state.locked);
+            const X64State& x64_state = GetState(GetX64For(arm_state.arm_reg));
-        if (arm_state.location.IsSimpleReg()) {
+            AssertStatesConsistent(x64_state, arm_state);
-            X64State& x64_state = x64_gpr[x64_reg_to_index.at(arm_state.location.GetSimpleReg())];
-            ASSERT(x64_state.state == X64State::State::CleanArmReg || x64_state.state == X64State::State::DirtyArmReg);
-            ASSERT(x64_state.arm_reg == static_cast<ArmReg>(i));
         }
-    }
+        return !arm_state.locked;
+    }));
 
-    for (auto i : x64_reg_to_index) {
+    ASSERT(std::all_of(x64_gpr.begin(), x64_gpr.end(), [](const X64State& x64_state) {
-        X64State& x64_state = x64_gpr[i.second];
+        ASSERT(x64_state.state != X64State::State::Temp); // Temp is never unlocked.
-        ASSERT(!x64_state.locked);
+        return !x64_state.locked;
-    }
+    }));
 }
 
 Gen::X64Reg RegAlloc::AllocReg() {
     // TODO: Improve with an actual register allocator as this is terrible.
 
     // First check to see if there anything free.
-    for (auto i : x64_reg_to_index) {
+    auto free_x64 = std::find_if(x64_gpr.begin(), x64_gpr.end(), [](const auto& x64_state) {
-        X64State& x64_state = x64_gpr[i.second];
+        return !x64_state.locked && x64_state.state == X64State::State::Free;
+    });
 
-        if (x64_state.locked)
+    if (free_x64 != x64_gpr.end()) {
-            continue;
+        return free_x64->x64_reg;
-
-        ASSERT(x64_state.state != X64State::State::Temp); // This can never happen.
-
-        if (x64_state.state == X64State::State::Free) {
-            return i.first;
-        }
     }
 
     // Otherwise flush something.
-    for (auto i : x64_reg_to_index) {
+    auto flushable_x64 = std::find_if(x64_gpr.begin(), x64_gpr.end(), [](const auto& x64_state) {
-        X64State& x64_state = x64_gpr[i.second];
+        return !x64_state.locked;
+    });
 
-        if (x64_state.locked)
+    if (flushable_x64 != x64_gpr.end()) {
-            continue;
+        FlushX64(flushable_x64->x64_reg);
-
+        return flushable_x64->x64_reg;
-        FlushX64(i.first);
-        return i.first;
     }
 
     ASSERT_MSG(false, "Ran out of x64 registers");

src/core/arm/jit_x64/reg_alloc.h

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <array>
+#include <map>
 
 #include "common/common_types.h"
 #include "common/common_funcs.h"
@@ -20,49 +21,6 @@ namespace JitX64 {
 //       (Designed so it should be simple to implement later.)
 
 class RegAlloc final {
-private:
-    struct ArmState {
-        /**
-         * Where is the current value of this register? There are two cases:
-         * - In an x64 register, in which case location.IsSimpleReg() == true.
-         * - In memory in ARMul_State, in which case location == MJitStateCpuReg(arm_reg).
-         */
-        Gen::OpArg location = { 0, 0, Gen::INVALID_REG, Gen::INVALID_REG };
-        bool locked = false;
-    };
-
-    /**
-     * Possible states of X64State:
-     *
-     * Free (locked must be false):    This x64 reg is free to be allocated for any purpose.
-     * Temp (locked must be true):     This x64 reg is being used as a temporary in a calculation.
-     * DirtyArmReg (arm_reg is valid): This x64 reg is bound to an ARM reg.
-     *                                 It is marked as dirty (value has changed).
-     *                                 This value MUST be flushed back to memory.
-     * CleanArmReg (arm_reg is valid): This x64 reg is bound to an ARM reg.
-     *                                 It hasn't been written to (i.e.: value is still the same as the in-memory version).
-     *                                 This value WILL NOT be flushed back to memory.
-     * UserManuallyLocked:             User has called LockX64 on this register. User must call UnlockX64 to unlock.
-     */
-    struct X64State {
-        enum class State {
-            Free,
-            Temp,
-            DirtyArmReg,
-            CleanArmReg,
-            UserManuallyLocked
-        };
-
-        bool locked = false;
-        State state = State::Free;
-        ArmReg arm_reg = ArmReg::INVALID_REG; ///< Only holds a valid value when state == DirtyArmReg / CleanArmReg
-    };
-
-    std::array<ArmState, 16> arm_gpr;
-    std::array<X64State, 16> x64_gpr;
-
-    Gen::XEmitter* code = nullptr;
-
 public:
     RegAlloc() { Init(nullptr); }
 
@@ -95,7 +53,7 @@ public:
      */
     Gen::OpArg LockArmForReadWrite(ArmReg arm_reg) {
         Gen::OpArg ret = LockArmForRead(arm_reg);
-        if (IsBoundToX64(arm_reg)) {
+        if (GetState(arm_reg).IsInX64Register()) {
             MarkDirty(arm_reg);
         }
         return ret;
@@ -171,13 +129,78 @@ public:
 
     // Debug:
 
-    void AssertNoLocked();
+    void AssertNoLocked() const;
 
 private:
+    struct ArmState {
+        ArmState() = default;
+        ArmState(ArmReg arm_reg, Gen::OpArg location, bool locked) : arm_reg(arm_reg), location(location), locked(locked) {}
+
+        /// Which register is this?
+        ArmReg arm_reg = ArmReg::INVALID_REG;
+
+        /**
+         * Where is the current value of this register? There are two cases:
+         * - In an x64 register, in which case location.IsSimpleReg() == true.
+         * - In memory in ARMul_State, in which case location == MJitStateCpuReg(arm_reg).
+         */
+        Gen::OpArg location = { 0, 0, Gen::INVALID_REG, Gen::INVALID_REG };
+
+        /// Are we currently in-use?
+        bool locked = false;
+
+        bool IsInX64Register() const { return location.IsSimpleReg(); }
+    };
+
+    /**
+     * Possible states of X64State:
+     *
+     * Free (locked must be false):    This x64 reg is free to be allocated for any purpose.
+     * Temp (locked must be true):     This x64 reg is being used as a temporary in a calculation.
+     * DirtyArmReg (arm_reg is valid): This x64 reg is bound to an ARM reg.
+     *                                 It is marked as dirty (value has changed).
+     *                                 This value MUST be flushed back to memory.
+     * CleanArmReg (arm_reg is valid): This x64 reg is bound to an ARM reg.
+     *                                 It hasn't been written to (i.e.: value is still the same as the in-memory version).
+     *                                 This value WILL NOT be flushed back to memory.
+     * UserManuallyLocked:             User has called LockX64 on this register. User must call UnlockX64 to unlock.
+     */
+    struct X64State {
+        enum class State {
+            Free,
+            Temp,
+            DirtyArmReg,
+            CleanArmReg,
+            UserManuallyLocked
+        };
+
+        X64State() = default;
+        X64State(Gen::X64Reg x64_reg, bool locked, State state, ArmReg arm_reg) : x64_reg(x64_reg), locked(locked), state(state), arm_reg(arm_reg) {}
+
+        Gen::X64Reg x64_reg = Gen::INVALID_REG;
+        bool locked = false;
+        State state = State::Free;
+        ArmReg arm_reg = ArmReg::INVALID_REG; ///< Only holds a valid value when state == DirtyArmReg / CleanArmReg
+    };
+
+    std::array<ArmState, 16> arm_gpr;
+    std::array<X64State, 15> x64_gpr;
+
+    Gen::XEmitter* code = nullptr;
+
+private:
+    static const std::map<Gen::X64Reg, size_t> x64_reg_to_index;
+
+    /// INTERNAL: Gets the X64State that corresponds to x64_reg.
+    X64State& GetState(Gen::X64Reg x64_reg);
+    const X64State& GetState(Gen::X64Reg x64_reg) const;
+    /// INTERNAL: Gets the ArmState that corresponds to arm_reg.
+    ArmState& GetState(ArmReg x64_reg);
+    const ArmState& GetState(ArmReg x64_reg) const;
+    /// INTERNAL: Checks consistency of the two states
+    void AssertStatesConsistent(const X64State& x64_state, const ArmState& arm_state) const;
     /// INTERNAL: Gets the x64 register this ArmReg is currently bound to.
-    Gen::X64Reg GetX64For(ArmReg arm_reg);
+    Gen::X64Reg GetX64For(ArmReg arm_reg) const;
-    /// INTERNAL: Is this ARM register currently in an x64 register?
-    bool IsBoundToX64(ArmReg arm_reg);
     /// INTERNAL: Marks register as dirty. Ensures that it is written back to memory if it's in a x64 register.
     void MarkDirty(ArmReg arm_reg);
     /// INTERNAL: Allocates a register that is free. Flushes registers that are not locked if necessary.