1 files changed, 44 insertions, 15 deletions

diff --git a/src/video_core/shader/control_flow.h b/src/video_core/shader/control_flow.h
index 6d0e50d7c..369ca255b 100644
--- a/src/video_core/shader/control_flow.h
+++ b/src/video_core/shader/control_flow.h
@@ -7,6 +7,7 @@
 #include <list>
 #include <optional>
 #include <set>
+#include <variant>
 
 #include "video_core/engines/shader_bytecode.h"
 #include "video_core/shader/ast.h"
@@ -37,29 +38,57 @@ struct Condition {
     }
 };
 
-struct ShaderBlock {
-    struct Branch {
-        Condition cond{};
-        bool kills{};
-        s32 address{};
+class SingleBranch {
+public:
+    SingleBranch() = default;
+    SingleBranch(Condition condition, s32 address, bool kill, bool is_sync, bool is_brk,
+                 bool ignore)
+        : condition{condition}, address{address}, kill{kill}, is_sync{is_sync}, is_brk{is_brk},
+          ignore{ignore} {}
+
+    bool operator==(const SingleBranch& b) const {
+        return std::tie(condition, address, kill, is_sync, is_brk, ignore) ==
+               std::tie(b.condition, b.address, b.kill, b.is_sync, b.is_brk, b.ignore);
+    }
+
+    Condition condition{};
+    s32 address{exit_branch};
+    bool kill{};
+    bool is_sync{};
+    bool is_brk{};
+    bool ignore{};
+};
+
+struct CaseBranch {
+    CaseBranch(u32 cmp_value, u32 address) : cmp_value{cmp_value}, address{address} {}
+    u32 cmp_value;
+    u32 address;
+};
+
+class MultiBranch {
+public:
+    MultiBranch(u32 gpr, std::vector<CaseBranch>& branches)
+        : gpr{gpr}, branches{std::move(branches)} {}
 
-        bool operator==(const Branch& b) const {
-            return std::tie(cond, kills, address) == std::tie(b.cond, b.kills, b.address);
-        }
+    u32 gpr{};
+    std::vector<CaseBranch> branches{};
+};
 
-        bool operator!=(const Branch& b) const {
-            return !operator==(b);
-        }
-    };
+using BranchData = std::variant<SingleBranch, MultiBranch>;
+using BlockBranchInfo = std::shared_ptr<BranchData>;
 
+bool BlockBranchInfoAreEqual(BlockBranchInfo first, BlockBranchInfo second);
+
+struct ShaderBlock {
     u32 start{};
     u32 end{};
     bool ignore_branch{};
-    Branch branch{};
+    BlockBranchInfo branch{};
 
     bool operator==(const ShaderBlock& sb) const {
-        return std::tie(start, end, ignore_branch, branch) ==
-               std::tie(sb.start, sb.end, sb.ignore_branch, sb.branch);
+        return std::tie(start, end, ignore_branch) ==
+                   std::tie(sb.start, sb.end, sb.ignore_branch) &&
+               BlockBranchInfoAreEqual(branch, sb.branch);
     }
 
     bool operator!=(const ShaderBlock& sb) const {