4 files changed, 55 insertions, 6 deletions

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