9 files changed, 130 insertions, 33 deletions

diff --git a/src/shader_recompiler/backend/glasm/emit_glasm_instructions.h b/src/shader_recompiler/backend/glasm/emit_glasm_instructions.h
index c9f4826ce..fef9ff9be 100644
--- a/src/shader_recompiler/backend/glasm/emit_glasm_instructions.h
+++ b/src/shader_recompiler/backend/glasm/emit_glasm_instructions.h
@@ -42,6 +42,8 @@ void EmitSetGotoVariable(EmitContext& ctx);
 void EmitGetGotoVariable(EmitContext& ctx);
 void EmitSetIndirectBranchVariable(EmitContext& ctx);
 void EmitGetIndirectBranchVariable(EmitContext& ctx);
+void EmitSetLoopSafetyVariable(EmitContext& ctx);
+void EmitGetLoopSafetyVariable(EmitContext& ctx);
 void EmitGetCbufU8(EmitContext& ctx, IR::Inst& inst, const IR::Value& binding, ScalarU32 offset);
 void EmitGetCbufS8(EmitContext& ctx, IR::Inst& inst, const IR::Value& binding, ScalarU32 offset);
 void EmitGetCbufU16(EmitContext& ctx, IR::Inst& inst, const IR::Value& binding, ScalarU32 offset);
diff --git a/src/shader_recompiler/backend/glasm/emit_glasm_not_implemented.cpp b/src/shader_recompiler/backend/glasm/emit_glasm_not_implemented.cpp
index 95bcbd750..60735fe31 100644
--- a/src/shader_recompiler/backend/glasm/emit_glasm_not_implemented.cpp
+++ b/src/shader_recompiler/backend/glasm/emit_glasm_not_implemented.cpp
@@ -153,6 +153,14 @@ void EmitGetIndirectBranchVariable(EmitContext& ctx) {
     NotImplemented();
 }
 
+void EmitSetLoopSafetyVariable(EmitContext& ctx) {
+    NotImplemented();
+}
+
+void EmitGetLoopSafetyVariable(EmitContext& ctx) {
+    NotImplemented();
+}
+
 void EmitGetZFlag(EmitContext& ctx) {
     NotImplemented();
 }
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp b/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
index 442a958a5..42fff74e3 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
@@ -163,35 +163,43 @@ Id GetCbufElement(EmitContext& ctx, Id vector, const IR::Value& offset, u32 inde
 } // Anonymous namespace
 
 void EmitGetRegister(EmitContext&) {
-    throw NotImplementedException("SPIR-V Instruction");
+    throw LogicError("Unreachable instruction");
 }
 
 void EmitSetRegister(EmitContext&) {
-    throw NotImplementedException("SPIR-V Instruction");
+    throw LogicError("Unreachable instruction");
 }
 
 void EmitGetPred(EmitContext&) {
-    throw NotImplementedException("SPIR-V Instruction");
+    throw LogicError("Unreachable instruction");
 }
 
 void EmitSetPred(EmitContext&) {
-    throw NotImplementedException("SPIR-V Instruction");
+    throw LogicError("Unreachable instruction");
 }
 
 void EmitSetGotoVariable(EmitContext&) {
-    throw NotImplementedException("SPIR-V Instruction");
+    throw LogicError("Unreachable instruction");
 }
 
 void EmitGetGotoVariable(EmitContext&) {
-    throw NotImplementedException("SPIR-V Instruction");
+    throw LogicError("Unreachable instruction");
 }
 
 void EmitSetIndirectBranchVariable(EmitContext&) {
-    throw NotImplementedException("SPIR-V Instruction");
+    throw LogicError("Unreachable instruction");
 }
 
 void EmitGetIndirectBranchVariable(EmitContext&) {
-    throw NotImplementedException("SPIR-V Instruction");
+    throw LogicError("Unreachable instruction");
+}
+
+void EmitSetLoopSafetyVariable(EmitContext&) {
+    throw LogicError("Unreachable instruction");
+}
+
+void EmitGetLoopSafetyVariable(EmitContext&) {
+    throw LogicError("Unreachable instruction");
 }
 
 Id EmitGetCbufU8(EmitContext& ctx, const IR::Value& binding, const IR::Value& offset) {
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv_instructions.h b/src/shader_recompiler/backend/spirv/emit_spirv_instructions.h
index 1181e7b4f..e3e5b03fe 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv_instructions.h
+++ b/src/shader_recompiler/backend/spirv/emit_spirv_instructions.h
@@ -43,6 +43,8 @@ void EmitSetGotoVariable(EmitContext& ctx);
 void EmitGetGotoVariable(EmitContext& ctx);
 void EmitSetIndirectBranchVariable(EmitContext& ctx);
 void EmitGetIndirectBranchVariable(EmitContext& ctx);
+void EmitSetLoopSafetyVariable(EmitContext& ctx);
+void EmitGetLoopSafetyVariable(EmitContext& ctx);
 Id EmitGetCbufU8(EmitContext& ctx, const IR::Value& binding, const IR::Value& offset);
 Id EmitGetCbufS8(EmitContext& ctx, const IR::Value& binding, const IR::Value& offset);
 Id EmitGetCbufU16(EmitContext& ctx, const IR::Value& binding, const IR::Value& offset);
diff --git a/src/shader_recompiler/frontend/ir/ir_emitter.cpp b/src/shader_recompiler/frontend/ir/ir_emitter.cpp
index e9fd41237..6c37af5e7 100644
--- a/src/shader_recompiler/frontend/ir/ir_emitter.cpp
+++ b/src/shader_recompiler/frontend/ir/ir_emitter.cpp
@@ -125,6 +125,12 @@ U1 IREmitter::GetPred(IR::Pred pred, bool is_negated) {
     }
 }
 
+void IREmitter::SetPred(IR::Pred pred, const U1& value) {
+    if (pred != IR::Pred::PT) {
+        Inst(Opcode::SetPred, pred, value);
+    }
+}
+
 U1 IREmitter::GetGotoVariable(u32 id) {
     return Inst<U1>(Opcode::GetGotoVariable, id);
 }
@@ -141,8 +147,12 @@ void IREmitter::SetIndirectBranchVariable(const U32& value) {
     Inst(Opcode::SetIndirectBranchVariable, value);
 }
 
-void IREmitter::SetPred(IR::Pred pred, const U1& value) {
-    Inst(Opcode::SetPred, pred, value);
+U32 IREmitter::GetLoopSafetyVariable(u32 id) {
+    return Inst<U32>(Opcode::GetLoopSafetyVariable, id);
+}
+
+void IREmitter::SetLoopSafetyVariable(u32 id, const U32& counter) {
+    Inst(Opcode::SetLoopSafetyVariable, id, counter);
 }
 
 U32 IREmitter::GetCbuf(const U32& binding, const U32& byte_offset) {
diff --git a/src/shader_recompiler/frontend/ir/ir_emitter.h b/src/shader_recompiler/frontend/ir/ir_emitter.h
index bb3500c54..7caab1f61 100644
--- a/src/shader_recompiler/frontend/ir/ir_emitter.h
+++ b/src/shader_recompiler/frontend/ir/ir_emitter.h
@@ -55,6 +55,9 @@ public:
     [[nodiscard]] U32 GetIndirectBranchVariable();
     void SetIndirectBranchVariable(const U32& value);
 
+    [[nodiscard]] U32 GetLoopSafetyVariable(u32 id);
+    void SetLoopSafetyVariable(u32 id, const U32& counter);
+
     [[nodiscard]] U32 GetCbuf(const U32& binding, const U32& byte_offset);
     [[nodiscard]] Value GetCbuf(const U32& binding, const U32& byte_offset, size_t bitsize,
                                 bool is_signed);
diff --git a/src/shader_recompiler/frontend/ir/opcodes.inc b/src/shader_recompiler/frontend/ir/opcodes.inc
index 8a8d0d759..e87aeddd5 100644
--- a/src/shader_recompiler/frontend/ir/opcodes.inc
+++ b/src/shader_recompiler/frontend/ir/opcodes.inc
@@ -32,6 +32,8 @@ OPCODE(GetGotoVariable,                                     U1,             U32,
 OPCODE(SetGotoVariable,                                     Void,           U32,            U1,                                                             )
 OPCODE(GetIndirectBranchVariable,                           U32,                                                                                            )
 OPCODE(SetIndirectBranchVariable,                           Void,           U32,                                                                            )
+OPCODE(GetLoopSafetyVariable,                               U32,            U32,                                                                            )
+OPCODE(SetLoopSafetyVariable,                               Void,           U32,            U32,                                                            )
 OPCODE(GetCbufU8,                                           U32,            U32,            U32,                                                            )
 OPCODE(GetCbufS8,                                           U32,            U32,            U32,                                                            )
 OPCODE(GetCbufU16,                                          U32,            U32,            U32,                                                            )
diff --git a/src/shader_recompiler/frontend/maxwell/structured_control_flow.cpp b/src/shader_recompiler/frontend/maxwell/structured_control_flow.cpp
index c1e0646e6..b2b8c492a 100644
--- a/src/shader_recompiler/frontend/maxwell/structured_control_flow.cpp
+++ b/src/shader_recompiler/frontend/maxwell/structured_control_flow.cpp
@@ -9,11 +9,13 @@
 #include <unordered_map>
 #include <utility>
 #include <vector>
+#include <version>
 
 #include <fmt/format.h>
 
 #include <boost/intrusive/list.hpp>
 
+#include "common/settings.h"
 #include "shader_recompiler/environment.h"
 #include "shader_recompiler/frontend/ir/basic_block.h"
 #include "shader_recompiler/frontend/ir/ir_emitter.h"
@@ -739,8 +741,25 @@ private:
             }
             case StatementType::Loop: {
                 IR::Block* const loop_header_block{block_pool.Create(inst_pool)};
-                if (current_block) {
-                    current_block->AddBranch(loop_header_block);
+                const u32 this_loop_id{loop_id++};
+
+                if (Settings::values.disable_shader_loop_safety_checks) {
+                    if (current_block) {
+                        current_block->AddBranch(loop_header_block);
+                    }
+                } else {
+                    IR::Block* const init_block{block_pool.Create(inst_pool)};
+                    IR::IREmitter ir{*init_block};
+                    ir.SetLoopSafetyVariable(this_loop_id, ir.Imm32(0x2000));
+
+                    if (current_block) {
+                        current_block->AddBranch(init_block);
+                    }
+                    init_block->AddBranch(loop_header_block);
+
+                    auto& init_node{syntax_list.emplace_back()};
+                    init_node.type = IR::AbstractSyntaxNode::Type::Block;
+                    init_node.data.block = init_block;
                 }
                 auto& header_node{syntax_list.emplace_back()};
                 header_node.type = IR::AbstractSyntaxNode::Type::Block;
@@ -758,7 +777,16 @@ private:
 
                 // The continue block is located at the end of the loop
                 IR::IREmitter ir{*continue_block};
-                const IR::U1 cond{ir.ConditionRef(VisitExpr(ir, *stmt.cond))};
+                IR::U1 cond{VisitExpr(ir, *stmt.cond)};
+                if (!Settings::values.disable_shader_loop_safety_checks) {
+                    const IR::U32 old_counter{ir.GetLoopSafetyVariable(this_loop_id)};
+                    const IR::U32 new_counter{ir.ISub(old_counter, ir.Imm32(1))};
+                    ir.SetLoopSafetyVariable(this_loop_id, new_counter);
+
+                    const IR::U1 safety_cond{ir.INotEqual(new_counter, ir.Imm32(0))};
+                    cond = ir.LogicalAnd(cond, safety_cond);
+                }
+                cond = ir.ConditionRef(cond);
 
                 IR::Block* const body_block{syntax_list.at(body_block_index).data.block};
                 loop_header_block->AddBranch(body_block);
@@ -863,8 +891,14 @@ private:
     ObjectPool<IR::Block>& block_pool;
     Environment& env;
     IR::AbstractSyntaxList& syntax_list;
-    // TODO: Make this constexpr when std::vector is constexpr
+    u32 loop_id{};
+
+// TODO: C++20 Remove this when all compilers support constexpr std::vector
+#if __cpp_lib_constexpr_vector >= 201907
+    static constexpr Flow::Block dummy_flow_block;
+#else
     const Flow::Block dummy_flow_block;
+#endif
 };
 } // Anonymous namespace
 
diff --git a/src/shader_recompiler/ir_opt/ssa_rewrite_pass.cpp b/src/shader_recompiler/ir_opt/ssa_rewrite_pass.cpp
index e54499ba5..a4ba393ef 100644
--- a/src/shader_recompiler/ir_opt/ssa_rewrite_pass.cpp
+++ b/src/shader_recompiler/ir_opt/ssa_rewrite_pass.cpp
@@ -48,73 +48,91 @@ struct GotoVariable : FlagTag {
     u32 index;
 };
 
+struct LoopSafetyVariable {
+    LoopSafetyVariable() = default;
+    explicit LoopSafetyVariable(u32 index_) : index{index_} {}
+
+    auto operator<=>(const LoopSafetyVariable&) const noexcept = default;
+
+    u32 index;
+};
+
 struct IndirectBranchVariable {
     auto operator<=>(const IndirectBranchVariable&) const noexcept = default;
 };
 
-using Variant = std::variant<IR::Reg, IR::Pred, ZeroFlagTag, SignFlagTag, CarryFlagTag,
-                             OverflowFlagTag, GotoVariable, IndirectBranchVariable>;
-using ValueMap = boost::container::flat_map<IR::Block*, IR::Value, std::less<IR::Block*>>;
+using Variant =
+    std::variant<IR::Reg, IR::Pred, ZeroFlagTag, SignFlagTag, CarryFlagTag, OverflowFlagTag,
+                 GotoVariable, LoopSafetyVariable, IndirectBranchVariable>;
+using ValueMap = boost::container::flat_map<IR::Block*, IR::Value>;
 
 struct DefTable {
-    const IR::Value& Def(IR::Block* block, IR::Reg variable) noexcept {
+    const IR::Value& Def(IR::Block* block, IR::Reg variable) {
         return block->SsaRegValue(variable);
     }
-    void SetDef(IR::Block* block, IR::Reg variable, const IR::Value& value) noexcept {
+    void SetDef(IR::Block* block, IR::Reg variable, const IR::Value& value) {
         block->SetSsaRegValue(variable, value);
     }
 
-    const IR::Value& Def(IR::Block* block, IR::Pred variable) noexcept {
+    const IR::Value& Def(IR::Block* block, IR::Pred variable) {
         return preds[IR::PredIndex(variable)][block];
     }
-    void SetDef(IR::Block* block, IR::Pred variable, const IR::Value& value) noexcept {
+    void SetDef(IR::Block* block, IR::Pred variable, const IR::Value& value) {
         preds[IR::PredIndex(variable)].insert_or_assign(block, value);
     }
 
-    const IR::Value& Def(IR::Block* block, GotoVariable variable) noexcept {
+    const IR::Value& Def(IR::Block* block, GotoVariable variable) {
         return goto_vars[variable.index][block];
     }
-    void SetDef(IR::Block* block, GotoVariable variable, const IR::Value& value) noexcept {
+    void SetDef(IR::Block* block, GotoVariable variable, const IR::Value& value) {
         goto_vars[variable.index].insert_or_assign(block, value);
     }
 
-    const IR::Value& Def(IR::Block* block, IndirectBranchVariable) noexcept {
+    const IR::Value& Def(IR::Block* block, LoopSafetyVariable variable) {
+        return loop_safety_vars[variable.index][block];
+    }
+    void SetDef(IR::Block* block, LoopSafetyVariable variable, const IR::Value& value) {
+        loop_safety_vars[variable.index].insert_or_assign(block, value);
+    }
+
+    const IR::Value& Def(IR::Block* block, IndirectBranchVariable) {
         return indirect_branch_var[block];
     }
-    void SetDef(IR::Block* block, IndirectBranchVariable, const IR::Value& value) noexcept {
+    void SetDef(IR::Block* block, IndirectBranchVariable, const IR::Value& value) {
         indirect_branch_var.insert_or_assign(block, value);
     }
 
-    const IR::Value& Def(IR::Block* block, ZeroFlagTag) noexcept {
+    const IR::Value& Def(IR::Block* block, ZeroFlagTag) {
         return zero_flag[block];
     }
-    void SetDef(IR::Block* block, ZeroFlagTag, const IR::Value& value) noexcept {
+    void SetDef(IR::Block* block, ZeroFlagTag, const IR::Value& value) {
         zero_flag.insert_or_assign(block, value);
     }
 
-    const IR::Value& Def(IR::Block* block, SignFlagTag) noexcept {
+    const IR::Value& Def(IR::Block* block, SignFlagTag) {
         return sign_flag[block];
     }
-    void SetDef(IR::Block* block, SignFlagTag, const IR::Value& value) noexcept {
+    void SetDef(IR::Block* block, SignFlagTag, const IR::Value& value) {
         sign_flag.insert_or_assign(block, value);
     }
 
-    const IR::Value& Def(IR::Block* block, CarryFlagTag) noexcept {
+    const IR::Value& Def(IR::Block* block, CarryFlagTag) {
         return carry_flag[block];
     }
-    void SetDef(IR::Block* block, CarryFlagTag, const IR::Value& value) noexcept {
+    void SetDef(IR::Block* block, CarryFlagTag, const IR::Value& value) {
         carry_flag.insert_or_assign(block, value);
     }
 
-    const IR::Value& Def(IR::Block* block, OverflowFlagTag) noexcept {
+    const IR::Value& Def(IR::Block* block, OverflowFlagTag) {
         return overflow_flag[block];
     }
-    void SetDef(IR::Block* block, OverflowFlagTag, const IR::Value& value) noexcept {
+    void SetDef(IR::Block* block, OverflowFlagTag, const IR::Value& value) {
         overflow_flag.insert_or_assign(block, value);
     }
 
     std::array<ValueMap, IR::NUM_USER_PREDS> preds;
     boost::container::flat_map<u32, ValueMap> goto_vars;
+    boost::container::flat_map<u32, ValueMap> loop_safety_vars;
     ValueMap indirect_branch_var;
     ValueMap zero_flag;
     ValueMap sign_flag;
@@ -134,6 +152,10 @@ IR::Opcode UndefOpcode(const FlagTag&) noexcept {
     return IR::Opcode::UndefU1;
 }
 
+IR::Opcode UndefOpcode(const LoopSafetyVariable&) noexcept {
+    return IR::Opcode::UndefU32;
+}
+
 IR::Opcode UndefOpcode(IndirectBranchVariable) noexcept {
     return IR::Opcode::UndefU32;
 }
@@ -315,6 +337,9 @@ void VisitInst(Pass& pass, IR::Block* block, IR::Inst& inst) {
     case IR::Opcode::SetGotoVariable:
         pass.WriteVariable(GotoVariable{inst.Arg(0).U32()}, block, inst.Arg(1));
         break;
+    case IR::Opcode::SetLoopSafetyVariable:
+        pass.WriteVariable(LoopSafetyVariable{inst.Arg(0).U32()}, block, inst.Arg(0));
+        break;
     case IR::Opcode::SetIndirectBranchVariable:
         pass.WriteVariable(IndirectBranchVariable{}, block, inst.Arg(0));
         break;
@@ -343,6 +368,9 @@ void VisitInst(Pass& pass, IR::Block* block, IR::Inst& inst) {
     case IR::Opcode::GetGotoVariable:
         inst.ReplaceUsesWith(pass.ReadVariable(GotoVariable{inst.Arg(0).U32()}, block));
         break;
+    case IR::Opcode::GetLoopSafetyVariable:
+        inst.ReplaceUsesWith(pass.ReadVariable(LoopSafetyVariable{inst.Arg(0).U32()}, block));
+        break;
     case IR::Opcode::GetIndirectBranchVariable:
         inst.ReplaceUsesWith(pass.ReadVariable(IndirectBranchVariable{}, block));
         break;