5 files changed, 72 insertions, 3 deletions

diff --git a/src/video_core/renderer_opengl/gl_shader_decompiler.cpp b/src/video_core/renderer_opengl/gl_shader_decompiler.cpp
index a311dbcfe..e1730821f 100644
--- a/src/video_core/renderer_opengl/gl_shader_decompiler.cpp
+++ b/src/video_core/renderer_opengl/gl_shader_decompiler.cpp
@@ -751,6 +751,11 @@ private:
 
     Expression Visit(const Node& node) {
         if (const auto operation = std::get_if<OperationNode>(&*node)) {
+            auto amend_index = operation->GetAmendIndex();
+            if (amend_index) {
+                const Node& amend_node = ir.GetAmendNode(*amend_index);
+                Visit(amend_node).CheckVoid();
+            }
             const auto operation_index = static_cast<std::size_t>(operation->GetCode());
             if (operation_index >= operation_decompilers.size()) {
                 UNREACHABLE_MSG("Out of bounds operation: {}", operation_index);
@@ -872,6 +877,11 @@ private:
         }
 
         if (const auto conditional = std::get_if<ConditionalNode>(&*node)) {
+            auto amend_index = conditional->GetAmendIndex();
+            if (amend_index) {
+                const Node& amend_node = ir.GetAmendNode(*amend_index);
+                Visit(amend_node).CheckVoid();
+            }
             // It's invalid to call conditional on nested nodes, use an operation instead
             code.AddLine("if ({}) {{", Visit(conditional->GetCondition()).AsBool());
             ++code.scope;
@@ -884,6 +894,11 @@ private:
         }
 
         if (const auto comment = std::get_if<CommentNode>(&*node)) {
+            auto amend_index = comment->GetAmendIndex();
+            if (amend_index) {
+                const Node& amend_node = ir.GetAmendNode(*amend_index);
+                Visit(amend_node).CheckVoid();
+            }
             code.AddLine("// " + comment->GetText());
             return {};
         }
diff --git a/src/video_core/renderer_vulkan/vk_shader_decompiler.cpp b/src/video_core/renderer_vulkan/vk_shader_decompiler.cpp
index a8baf91de..50feeb003 100644
--- a/src/video_core/renderer_vulkan/vk_shader_decompiler.cpp
+++ b/src/video_core/renderer_vulkan/vk_shader_decompiler.cpp
@@ -954,6 +954,12 @@ private:
 
     Expression Visit(const Node& node) {
         if (const auto operation = std::get_if<OperationNode>(&*node)) {
+            auto amend_index = operation->GetAmendIndex();
+            if (amend_index) {
+                const Node& amend_node = ir.GetAmendNode(*amend_index);
+                [[maybe_unused]] const Type type = Visit(amend_node).type;
+                ASSERT(type == Type::Void);
+            }
             const auto operation_index = static_cast<std::size_t>(operation->GetCode());
             const auto decompiler = operation_decompilers[operation_index];
             if (decompiler == nullptr) {
@@ -1142,6 +1148,12 @@ private:
         }
 
         if (const auto conditional = std::get_if<ConditionalNode>(&*node)) {
+            auto amend_index = conditional->GetAmendIndex();
+            if (amend_index) {
+                const Node& amend_node = ir.GetAmendNode(*amend_index);
+                [[maybe_unused]] const Type type = Visit(amend_node).type;
+                ASSERT(type == Type::Void);
+            }
             // It's invalid to call conditional on nested nodes, use an operation instead
             const Id true_label = OpLabel();
             const Id skip_label = OpLabel();
@@ -1164,6 +1176,12 @@ private:
         }
 
         if (const auto comment = std::get_if<CommentNode>(&*node)) {
+            auto amend_index = comment->GetAmendIndex();
+            if (amend_index) {
+                const Node& amend_node = ir.GetAmendNode(*amend_index);
+                [[maybe_unused]] const Type type = Visit(amend_node).type;
+                ASSERT(type == Type::Void);
+            }
             Name(OpUndef(t_void), comment->GetText());
             return {};
         }
diff --git a/src/video_core/shader/node.h b/src/video_core/shader/node.h
index 4d2f4d6a8..42e82ab74 100644
--- a/src/video_core/shader/node.h
+++ b/src/video_core/shader/node.h
@@ -392,8 +392,30 @@ struct MetaImage {
 using Meta =
     std::variant<MetaArithmetic, MetaTexture, MetaImage, MetaStackClass, Tegra::Shader::HalfType>;
 
+class AmendNode {
+public:
+    std::optional<u32> GetAmendIndex() const {
+        if (amend_index == amend_null_index) {
+            return std::nullopt;
+        }
+        return {amend_index};
+    }
+
+    void SetAmendIndex(u32 index) {
+        amend_index = index;
+    }
+
+    void ClearAmend() {
+        amend_index = amend_null_index;
+    }
+
+private:
+    static constexpr u32 amend_null_index = 0xFFFFFFFF;
+    u32 amend_index{amend_null_index};
+};
+
 /// Holds any kind of operation that can be done in the IR
-class OperationNode final {
+class OperationNode final : public AmendNode {
 public:
     explicit OperationNode(OperationCode code) : OperationNode(code, Meta{}) {}
 
@@ -433,7 +455,7 @@ private:
 };
 
 /// Encloses inside any kind of node that returns a boolean conditionally-executed code
-class ConditionalNode final {
+class ConditionalNode final : public AmendNode {
 public:
     explicit ConditionalNode(Node condition, std::vector<Node>&& code)
         : condition{std::move(condition)}, code{std::move(code)} {}
@@ -630,7 +652,7 @@ private:
 };
 
 /// Commentary, can be dropped
-class CommentNode final {
+class CommentNode final : public AmendNode {
 public:
     explicit CommentNode(std::string text) : text{std::move(text)} {}
 
diff --git a/src/video_core/shader/shader_ir.cpp b/src/video_core/shader/shader_ir.cpp
index 1d9825c76..49678767c 100644
--- a/src/video_core/shader/shader_ir.cpp
+++ b/src/video_core/shader/shader_ir.cpp
@@ -446,4 +446,10 @@ Node ShaderIR::BitfieldInsert(Node base, Node insert, u32 offset, u32 bits) {
                      Immediate(bits));
 }
 
+u32 ShaderIR::DeclareAmend(Node new_amend) {
+    const u32 id = static_cast<u32>(amend_code.size());
+    amend_code.push_back(new_amend);
+    return id;
+}
+
 } // namespace VideoCommon::Shader
diff --git a/src/video_core/shader/shader_ir.h b/src/video_core/shader/shader_ir.h
index baed06ccd..52f130e1b 100644
--- a/src/video_core/shader/shader_ir.h
+++ b/src/video_core/shader/shader_ir.h
@@ -176,6 +176,10 @@ public:
     /// Returns a condition code evaluated from internal flags
     Node GetConditionCode(Tegra::Shader::ConditionCode cc) const;
 
+    const Node& GetAmendNode(u32 index) const {
+        return amend_code[index];
+    }
+
 private:
     friend class ASTDecoder;
 
@@ -392,6 +396,9 @@ private:
                                                                Tegra::Shader::Instruction instr,
                                                                bool is_write);
 
+    /// Amends
+    u32 DeclareAmend(Node new_amend);
+
     const ProgramCode& program_code;
     const u32 main_offset;
     const CompilerSettings settings;
@@ -406,6 +413,7 @@ private:
     std::map<u32, NodeBlock> basic_blocks;
     NodeBlock global_code;
     ASTManager program_manager{true, true};
+    NodeBlock amend_code;
 
     std::set<u32> used_registers;
     std::set<Tegra::Shader::Pred> used_predicates;