From da936d6ad8cef5418b7644754ee4bcbf7f6125f8 Mon Sep 17 00:00:00 2001
From: FernandoS27 <fsahmkow27@gmail.com>
Date: Sun, 18 Apr 2021 19:10:55 +0200
Subject: [PATCH] shader: Implement delegation of Exit to dispatcher on CFG

---
 .../frontend/maxwell/control_flow.cpp         | 41 ++++++++++++++++++-
 .../frontend/maxwell/control_flow.h           |  9 +++-
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/src/shader_recompiler/frontend/maxwell/control_flow.cpp b/src/shader_recompiler/frontend/maxwell/control_flow.cpp
index 9811183f12..298faa03e0 100644
--- a/src/shader_recompiler/frontend/maxwell/control_flow.cpp
+++ b/src/shader_recompiler/frontend/maxwell/control_flow.cpp
@@ -185,8 +185,20 @@ Function::Function(ObjectPool<Block>& block_pool, Location start_address)
     label.block->branch_false = nullptr;
 }
 
-CFG::CFG(Environment& env_, ObjectPool<Block>& block_pool_, Location start_address)
-    : env{env_}, block_pool{block_pool_}, program_start{start_address} {
+CFG::CFG(Environment& env_, ObjectPool<Block>& block_pool_, Location start_address,
+         bool exits_to_dispatcher_)
+    : env{env_}, block_pool{block_pool_}, program_start{start_address}, exits_to_dispatcher{
+                                                                            exits_to_dispatcher_} {
+    if (exits_to_dispatcher) {
+        dispatch_block = block_pool.Create(Block{});
+        dispatch_block->begin = {};
+        dispatch_block->end = {};
+        dispatch_block->end_class = EndClass::Exit;
+        dispatch_block->cond = IR::Condition(true);
+        dispatch_block->stack = {};
+        dispatch_block->branch_true = nullptr;
+        dispatch_block->branch_false = nullptr;
+    }
     functions.emplace_back(block_pool, start_address);
     for (FunctionId function_id = 0; function_id < functions.size(); ++function_id) {
         while (!functions[function_id].labels.empty()) {
@@ -196,6 +208,12 @@ CFG::CFG(Environment& env_, ObjectPool<Block>& block_pool_, Location start_addre
             AnalyzeLabel(function_id, label);
         }
     }
+    if (exits_to_dispatcher) {
+        const auto it = functions[0].blocks.rbegin();
+        dispatch_block->begin = it->end + 1;
+        dispatch_block->end = it->end + 1;
+        functions[0].blocks.insert(*dispatch_block);
+    }
 }
 
 void CFG::AnalyzeLabel(FunctionId function_id, Label& label) {
@@ -462,11 +480,22 @@ CFG::AnalysisState CFG::AnalyzeEXIT(Block* block, FunctionId function_id, Locati
         // EXIT will never be taken
         return AnalysisState::Continue;
     }
+    if (exits_to_dispatcher && function_id != 0) {
+        throw NotImplementedException("Dispatch EXIT on external function.");
+    }
     if (pred != Predicate{true} || flow_test != IR::FlowTest::T) {
         if (block->stack.Peek(Token::PEXIT).has_value()) {
             throw NotImplementedException("Conditional EXIT with PEXIT token");
         }
         const IR::Condition cond{flow_test, static_cast<IR::Pred>(pred.index), pred.negated};
+        if (exits_to_dispatcher) {
+            block->end = pc;
+            block->branch_true = dispatch_block;
+            block->end_class = EndClass::Branch;
+            block->cond = cond;
+            block->branch_false = AddLabel(block, block->stack, pc + 1, function_id);
+            return AnalysisState::Branch;
+        }
         AnalyzeCondInst(block, function_id, pc, EndClass::Exit, cond);
         return AnalysisState::Branch;
     }
@@ -477,6 +506,14 @@ CFG::AnalysisState CFG::AnalyzeEXIT(Block* block, FunctionId function_id, Locati
         block->branch_false = nullptr;
         return AnalysisState::Branch;
     }
+    if (exits_to_dispatcher) {
+        block->cond = IR::Condition{true};
+        block->end = pc;
+        block->end_class = EndClass::Branch;
+        block->branch_true = dispatch_block;
+        block->branch_false = nullptr;
+        return AnalysisState::Branch;
+    }
     block->end = pc + 1;
     block->end_class = EndClass::Exit;
     return AnalysisState::Branch;
diff --git a/src/shader_recompiler/frontend/maxwell/control_flow.h b/src/shader_recompiler/frontend/maxwell/control_flow.h
index 89966b16aa..0e515c3b6f 100644
--- a/src/shader_recompiler/frontend/maxwell/control_flow.h
+++ b/src/shader_recompiler/frontend/maxwell/control_flow.h
@@ -111,7 +111,8 @@ class CFG {
     };
 
 public:
-    explicit CFG(Environment& env, ObjectPool<Block>& block_pool, Location start_address);
+    explicit CFG(Environment& env, ObjectPool<Block>& block_pool, Location start_address,
+                 bool exits_to_dispatcher = false);
 
     CFG& operator=(const CFG&) = delete;
     CFG(const CFG&) = delete;
@@ -128,6 +129,10 @@ public:
         return std::span(functions.data(), functions.size());
     }
 
+    [[nodiscard]] bool ExitsToDispatcher() const {
+        return exits_to_dispatcher;
+    }
+
 private:
     void AnalyzeLabel(FunctionId function_id, Label& label);
 
@@ -158,6 +163,8 @@ private:
     boost::container::small_vector<Function, 1> functions;
     FunctionId current_function_id{0};
     Location program_start;
+    bool exits_to_dispatcher{};
+    Block* dispatch_block{};
 };
 
 } // namespace Shader::Maxwell::Flow