5 files changed, 74 insertions, 49 deletions

diff --git a/src/shader_recompiler/backend/spirv/emit_context.cpp b/src/shader_recompiler/backend/spirv/emit_context.cpp
index 9759591bd..a98e08392 100644
--- a/src/shader_recompiler/backend/spirv/emit_context.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_context.cpp
@@ -463,7 +463,6 @@ EmitContext::EmitContext(const Profile& profile_, IR::Program& program, Bindings
     DefineImages(program.info, image_binding);
     DefineAttributeMemAccess(program.info);
     DefineGlobalMemoryFunctions(program.info);
-    DefineLabels(program);
 }
 
 EmitContext::~EmitContext() = default;
@@ -487,8 +486,6 @@ Id EmitContext::Def(const IR::Value& value) {
         return Const(value.F32());
     case IR::Type::F64:
         return Constant(F64[1], value.F64());
-    case IR::Type::Label:
-        return value.Label()->Definition<Id>();
     default:
         throw NotImplementedException("Immediate type {}", value.Type());
     }
@@ -1139,12 +1136,6 @@ void EmitContext::DefineImages(const Info& info, u32& binding) {
     }
 }
 
-void EmitContext::DefineLabels(IR::Program& program) {
-    for (IR::Block* const block : program.blocks) {
-        block->SetDefinition(OpLabel());
-    }
-}
-
 void EmitContext::DefineInputs(const Info& info) {
     if (info.uses_workgroup_id) {
         workgroup_id = DefineInput(*this, U32[3], false, spv::BuiltIn::WorkgroupId);
diff --git a/src/shader_recompiler/backend/spirv/emit_context.h b/src/shader_recompiler/backend/spirv/emit_context.h
index 8b000f1ec..d2b79f6c1 100644
--- a/src/shader_recompiler/backend/spirv/emit_context.h
+++ b/src/shader_recompiler/backend/spirv/emit_context.h
@@ -296,7 +296,6 @@ private:
     void DefineImages(const Info& info, u32& binding);
     void DefineAttributeMemAccess(const Info& info);
     void DefineGlobalMemoryFunctions(const Info& info);
-    void DefineLabels(IR::Program& program);
 
     void DefineInputs(const Info& info);
     void DefineOutputs(const IR::Program& program);
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv.cpp b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
index 2dad87e87..c22edfec2 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
@@ -41,8 +41,6 @@ ArgType Arg(EmitContext& ctx, const IR::Value& arg) {
         return arg;
     } else if constexpr (std::is_same_v<ArgType, u32>) {
         return arg.U32();
-    } else if constexpr (std::is_same_v<ArgType, IR::Block*>) {
-        return arg.Label();
     } else if constexpr (std::is_same_v<ArgType, IR::Attribute>) {
         return arg.Attribute();
     } else if constexpr (std::is_same_v<ArgType, IR::Patch>) {
@@ -109,15 +107,74 @@ Id TypeId(const EmitContext& ctx, IR::Type type) {
     }
 }
 
+void Traverse(EmitContext& ctx, IR::Program& program) {
+    IR::Block* current_block{};
+    for (const IR::AbstractSyntaxNode& node : program.syntax_list) {
+        switch (node.type) {
+        case IR::AbstractSyntaxNode::Type::Block:
+            const Id label{node.block->Definition<Id>()};
+            if (current_block) {
+                ctx.OpBranch(label);
+            }
+            current_block = node.block;
+            ctx.AddLabel(label);
+            for (IR::Inst& inst : node.block->Instructions()) {
+                EmitInst(ctx, &inst);
+            }
+            break;
+        case IR::AbstractSyntaxNode::Type::If: {
+            const Id if_label{node.if_node.body->Definition<Id>()};
+            const Id endif_label{node.if_node.merge->Definition<Id>()};
+            ctx.OpSelectionMerge(endif_label, spv::SelectionControlMask::MaskNone);
+            ctx.OpBranchConditional(ctx.Def(node.if_node.cond), if_label, endif_label);
+            break;
+        }
+        case IR::AbstractSyntaxNode::Type::Loop: {
+            const Id body_label{node.loop.body->Definition<Id>()};
+            const Id continue_label{node.loop.continue_block->Definition<Id>()};
+            const Id endloop_label{node.loop.merge->Definition<Id>()};
+
+            ctx.OpLoopMerge(endloop_label, continue_label, spv::LoopControlMask::MaskNone);
+            ctx.OpBranch(node.loop.body->Definition<Id>());
+            break;
+        }
+        case IR::AbstractSyntaxNode::Type::Break: {
+            const Id break_label{node.break_node.merge->Definition<Id>()};
+            const Id skip_label{node.break_node.skip->Definition<Id>()};
+            ctx.OpBranchConditional(ctx.Def(node.break_node.cond), break_label, skip_label);
+            break;
+        }
+        case IR::AbstractSyntaxNode::Type::EndIf:
+            if (current_block) {
+                ctx.OpBranch(node.end_if.merge->Definition<Id>());
+            }
+            break;
+        case IR::AbstractSyntaxNode::Type::Repeat: {
+            const Id loop_header_label{node.repeat.loop_header->Definition<Id>()};
+            const Id merge_label{node.repeat.merge->Definition<Id>()};
+            ctx.OpBranchConditional(ctx.Def(node.repeat.cond), loop_header_label, merge_label);
+            break;
+        }
+        case IR::AbstractSyntaxNode::Type::Return:
+            ctx.OpReturn();
+            break;
+        case IR::AbstractSyntaxNode::Type::Unreachable:
+            ctx.OpUnreachable();
+            break;
+        }
+        if (node.type != IR::AbstractSyntaxNode::Type::Block) {
+            current_block = nullptr;
+        }
+    }
+}
+
 Id DefineMain(EmitContext& ctx, IR::Program& program) {
     const Id void_function{ctx.TypeFunction(ctx.void_id)};
     const Id main{ctx.OpFunction(ctx.void_id, spv::FunctionControlMask::MaskNone, void_function)};
     for (IR::Block* const block : program.blocks) {
-        ctx.AddLabel(block->Definition<Id>());
-        for (IR::Inst& inst : block->Instructions()) {
-            EmitInst(ctx, &inst);
-        }
+        block->SetDefinition(ctx.OpLabel());
     }
+    Traverse(ctx, program);
     ctx.OpFunctionEnd();
     return main;
 }
@@ -411,6 +468,8 @@ Id EmitIdentity(EmitContext& ctx, const IR::Value& value) {
     return id;
 }
 
+void EmitBranchConditionRef(EmitContext&) {}
+
 void EmitGetZeroFromOp(EmitContext&) {
     throw LogicError("Unreachable instruction");
 }
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv_control_flow.cpp b/src/shader_recompiler/backend/spirv/emit_spirv_control_flow.cpp
index 6154c46be..d33486f28 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv_control_flow.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv_control_flow.cpp
@@ -7,40 +7,21 @@
 
 namespace Shader::Backend::SPIRV {
 
-void EmitBranch(EmitContext& ctx, Id label) {
-    ctx.OpBranch(label);
-}
-
-void EmitBranchConditional(EmitContext& ctx, Id condition, Id true_label, Id false_label) {
-    ctx.OpBranchConditional(condition, true_label, false_label);
-}
-
-void EmitLoopMerge(EmitContext& ctx, Id merge_label, Id continue_label) {
-    ctx.OpLoopMerge(merge_label, continue_label, spv::LoopControlMask::MaskNone);
-}
-
-void EmitSelectionMerge(EmitContext& ctx, Id merge_label) {
-    ctx.OpSelectionMerge(merge_label, spv::SelectionControlMask::MaskNone);
-}
-
-void EmitReturn(EmitContext& ctx) {
-    ctx.OpReturn();
-}
-
 void EmitJoin(EmitContext&) {
     throw NotImplementedException("Join shouldn't be emitted");
 }
 
-void EmitUnreachable(EmitContext& ctx) {
-    ctx.OpUnreachable();
-}
-
-void EmitDemoteToHelperInvocation(EmitContext& ctx, Id continue_label) {
+void EmitDemoteToHelperInvocation(EmitContext& ctx) {
     if (ctx.profile.support_demote_to_helper_invocation) {
         ctx.OpDemoteToHelperInvocationEXT();
-        ctx.OpBranch(continue_label);
     } else {
+        const Id kill_label{ctx.OpLabel()};
+        const Id impossible_label{ctx.OpLabel()};
+        ctx.OpSelectionMerge(impossible_label, spv::SelectionControlMask::MaskNone);
+        ctx.OpBranchConditional(ctx.true_value, kill_label, impossible_label);
+        ctx.AddLabel(kill_label);
         ctx.OpKill();
+        ctx.AddLabel(impossible_label);
     }
 }
 
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv_instructions.h b/src/shader_recompiler/backend/spirv/emit_spirv_instructions.h
index a1ca3f43d..2f4f6e59e 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv_instructions.h
+++ b/src/shader_recompiler/backend/spirv/emit_spirv_instructions.h
@@ -23,14 +23,9 @@ class EmitContext;
 Id EmitPhi(EmitContext& ctx, IR::Inst* inst);
 void EmitVoid(EmitContext& ctx);
 Id EmitIdentity(EmitContext& ctx, const IR::Value& value);
-void EmitBranch(EmitContext& ctx, Id label);
-void EmitBranchConditional(EmitContext& ctx, Id condition, Id true_label, Id false_label);
-void EmitLoopMerge(EmitContext& ctx, Id merge_label, Id continue_label);
-void EmitSelectionMerge(EmitContext& ctx, Id merge_label);
-void EmitReturn(EmitContext& ctx);
+void EmitBranchConditionRef(EmitContext&);
 void EmitJoin(EmitContext& ctx);
-void EmitUnreachable(EmitContext& ctx);
-void EmitDemoteToHelperInvocation(EmitContext& ctx, Id continue_label);
+void EmitDemoteToHelperInvocation(EmitContext& ctx);
 void EmitBarrier(EmitContext& ctx);
 void EmitWorkgroupMemoryBarrier(EmitContext& ctx);
 void EmitDeviceMemoryBarrier(EmitContext& ctx);