7 files changed, 214 insertions, 104 deletions

diff --git a/src/video_core/shader/ast.cpp b/src/video_core/shader/ast.cpp
index d521a7b52..0bf289f98 100644
--- a/src/video_core/shader/ast.cpp
+++ b/src/video_core/shader/ast.cpp
@@ -363,6 +363,71 @@ std::string ASTManager::Print() {
     return printer.GetResult();
 }
 
+ASTManager::ASTManager() = default;
+
+ASTManager::~ASTManager() {
+    Clear();
+}
+
+void ASTManager::Init() {
+    main_node = ASTBase::Make<ASTProgram>(ASTNode{});
+    program = std::get_if<ASTProgram>(main_node->GetInnerData());
+    true_condition = MakeExpr<ExprBoolean>(true);
+}
+
+ASTManager::ASTManager(ASTManager&& other)
+    : labels_map(std::move(other.labels_map)), labels_count{other.labels_count},
+      gotos(std::move(other.gotos)), labels(std::move(other.labels)), variables{other.variables},
+      program{other.program}, main_node{other.main_node}, true_condition{other.true_condition} {
+    other.main_node.reset();
+}
+
+ASTManager& ASTManager::operator=(ASTManager&& other) {
+    labels_map = std::move(other.labels_map);
+    labels_count = other.labels_count;
+    gotos = std::move(other.gotos);
+    labels = std::move(other.labels);
+    variables = other.variables;
+    program = other.program;
+    main_node = other.main_node;
+    true_condition = other.true_condition;
+
+    other.main_node.reset();
+    return *this;
+}
+
+void ASTManager::DeclareLabel(u32 address) {
+    const auto pair = labels_map.emplace(address, labels_count);
+    if (pair.second) {
+        labels_count++;
+        labels.resize(labels_count);
+    }
+}
+
+void ASTManager::InsertLabel(u32 address) {
+    u32 index = labels_map[address];
+    ASTNode label = ASTBase::Make<ASTLabel>(main_node, index);
+    labels[index] = label;
+    program->nodes.PushBack(label);
+}
+
+void ASTManager::InsertGoto(Expr condition, u32 address) {
+    u32 index = labels_map[address];
+    ASTNode goto_node = ASTBase::Make<ASTGoto>(main_node, condition, index);
+    gotos.push_back(goto_node);
+    program->nodes.PushBack(goto_node);
+}
+
+void ASTManager::InsertBlock(u32 start_address, u32 end_address) {
+    ASTNode block = ASTBase::Make<ASTBlockEncoded>(main_node, start_address, end_address);
+    program->nodes.PushBack(block);
+}
+
+void ASTManager::InsertReturn(Expr condition, bool kills) {
+    ASTNode node = ASTBase::Make<ASTReturn>(main_node, condition, kills);
+    program->nodes.PushBack(node);
+}
+
 void ASTManager::Decompile() {
     auto it = gotos.begin();
     while (it != gotos.end()) {
@@ -460,7 +525,6 @@ void ASTManager::SanityCheck() {
 
 void ASTManager::EncloseDoWhile(ASTNode goto_node, ASTNode label) {
     // ShowCurrentState("Before DoWhile Enclose");
-    enclose_count++;
     ASTZipper& zipper = goto_node->GetManager();
     ASTNode loop_start = label->GetNext();
     if (loop_start == goto_node) {
@@ -481,7 +545,6 @@ void ASTManager::EncloseDoWhile(ASTNode goto_node, ASTNode label) {
 
 void ASTManager::EncloseIfThen(ASTNode goto_node, ASTNode label) {
     // ShowCurrentState("Before IfThen Enclose");
-    enclose_count++;
     ASTZipper& zipper = goto_node->GetManager();
     ASTNode if_end = label->GetPrevious();
     if (if_end == goto_node) {
@@ -514,7 +577,6 @@ void ASTManager::EncloseIfThen(ASTNode goto_node, ASTNode label) {
 
 void ASTManager::MoveOutward(ASTNode goto_node) {
     // ShowCurrentState("Before MoveOutward");
-    outward_count++;
     ASTZipper& zipper = goto_node->GetManager();
     ASTNode parent = goto_node->GetParent();
     ASTZipper& zipper2 = parent->GetManager();
@@ -582,4 +644,75 @@ void ASTManager::MoveOutward(ASTNode goto_node) {
     // ShowCurrentState("After MoveOutward");
 }
 
+class ASTClearer {
+public:
+    ASTClearer() = default;
+
+    void operator()(ASTProgram& ast) {
+        ASTNode current = ast.nodes.GetFirst();
+        while (current) {
+            Visit(current);
+            current = current->GetNext();
+        }
+    }
+
+    void operator()(ASTIfThen& ast) {
+        ASTNode current = ast.nodes.GetFirst();
+        while (current) {
+            Visit(current);
+            current = current->GetNext();
+        }
+    }
+
+    void operator()(ASTIfElse& ast) {
+        ASTNode current = ast.nodes.GetFirst();
+        while (current) {
+            Visit(current);
+            current = current->GetNext();
+        }
+    }
+
+    void operator()(ASTBlockEncoded& ast) {}
+
+    void operator()(ASTBlockDecoded& ast) {
+        ast.nodes.clear();
+    }
+
+    void operator()(ASTVarSet& ast) {}
+
+    void operator()(ASTLabel& ast) {}
+
+    void operator()(ASTGoto& ast) {}
+
+    void operator()(ASTDoWhile& ast) {
+        ASTNode current = ast.nodes.GetFirst();
+        while (current) {
+            Visit(current);
+            current = current->GetNext();
+        }
+    }
+
+    void operator()(ASTReturn& ast) {}
+
+    void operator()(ASTBreak& ast) {}
+
+    void Visit(ASTNode& node) {
+        std::visit(*this, *node->GetInnerData());
+        node->Clear();
+    }
+};
+
+void ASTManager::Clear() {
+    if (!main_node) {
+        return;
+    }
+    ASTClearer clearer{};
+    clearer.Visit(main_node);
+    main_node.reset();
+    program = nullptr;
+    labels_map.clear();
+    labels.clear();
+    gotos.clear();
+}
+
 } // namespace VideoCommon::Shader
diff --git a/src/video_core/shader/ast.h b/src/video_core/shader/ast.h
index 4276f66a9..958989bcd 100644
--- a/src/video_core/shader/ast.h
+++ b/src/video_core/shader/ast.h
@@ -30,8 +30,8 @@ class ASTDoWhile;
 class ASTReturn;
 class ASTBreak;
 
-using ASTData = std::variant<ASTProgram, ASTIfThen, ASTIfElse, ASTBlockEncoded, ASTBlockDecoded, ASTVarSet, ASTGoto,
-                             ASTLabel, ASTDoWhile, ASTReturn, ASTBreak>;
+using ASTData = std::variant<ASTProgram, ASTIfThen, ASTIfElse, ASTBlockEncoded, ASTBlockDecoded,
+                             ASTVarSet, ASTGoto, ASTLabel, ASTDoWhile, ASTReturn, ASTBreak>;
 
 using ASTNode = std::shared_ptr<ASTBase>;
 
@@ -261,6 +261,13 @@ public:
         return nullptr;
     }
 
+    void Clear() {
+        next.reset();
+        previous.reset();
+        parent.reset();
+        manager = nullptr;
+    }
+
 private:
     friend class ASTZipper;
 
@@ -273,43 +280,26 @@ private:
 
 class ASTManager final {
 public:
-    explicit ASTManager() {
-        main_node = ASTBase::Make<ASTProgram>(ASTNode{});
-        program = std::get_if<ASTProgram>(main_node->GetInnerData());
-        true_condition = MakeExpr<ExprBoolean>(true);
-    }
+    ASTManager();
+    ~ASTManager();
 
-    void DeclareLabel(u32 address) {
-        const auto pair = labels_map.emplace(address, labels_count);
-        if (pair.second) {
-            labels_count++;
-            labels.resize(labels_count);
-        }
-    }
+    ASTManager(const ASTManager& o) = delete;
+    ASTManager& operator=(const ASTManager& other) = delete;
 
-    void InsertLabel(u32 address) {
-        u32 index = labels_map[address];
-        ASTNode label = ASTBase::Make<ASTLabel>(main_node, index);
-        labels[index] = label;
-        program->nodes.PushBack(label);
-    }
+    ASTManager(ASTManager&& other);
+    ASTManager& operator=(ASTManager&& other);
 
-    void InsertGoto(Expr condition, u32 address) {
-        u32 index = labels_map[address];
-        ASTNode goto_node = ASTBase::Make<ASTGoto>(main_node, condition, index);
-        gotos.push_back(goto_node);
-        program->nodes.PushBack(goto_node);
-    }
+    void Init();
 
-    void InsertBlock(u32 start_address, u32 end_address) {
-        ASTNode block = ASTBase::Make<ASTBlockEncoded>(main_node, start_address, end_address);
-        program->nodes.PushBack(block);
-    }
+    void DeclareLabel(u32 address);
 
-    void InsertReturn(Expr condition, bool kills) {
-        ASTNode node = ASTBase::Make<ASTReturn>(main_node, condition, kills);
-        program->nodes.PushBack(node);
-    }
+    void InsertLabel(u32 address);
+
+    void InsertGoto(Expr condition, u32 address);
+
+    void InsertBlock(u32 start_address, u32 end_address);
+
+    void InsertReturn(Expr condition, bool kills);
 
     std::string Print();
 
@@ -323,6 +313,12 @@ public:
         return gotos.size() == 0;
     }
 
+    ASTNode GetProgram() {
+        return main_node;
+    }
+
+    void Clear();
+
 private:
     bool IndirectlyRelated(ASTNode first, ASTNode second);
 
@@ -332,7 +328,7 @@ private:
 
     void EncloseIfThen(ASTNode goto_node, ASTNode label);
 
-    void MoveOutward(ASTNode goto_node) ;
+    void MoveOutward(ASTNode goto_node);
 
     u32 NewVariable() {
         u32 new_var = variables;
@@ -345,11 +341,9 @@ private:
     std::vector<ASTNode> labels{};
     std::list<ASTNode> gotos{};
     u32 variables{};
-    ASTProgram* program;
-    ASTNode main_node;
-    Expr true_condition;
-    u32 outward_count{};
-    u32 enclose_count{};
+    ASTProgram* program{};
+    ASTNode main_node{};
+    Expr true_condition{};
 };
 
 } // namespace VideoCommon::Shader
diff --git a/src/video_core/shader/control_flow.cpp b/src/video_core/shader/control_flow.cpp
index 7a21d870f..deb3d3ebd 100644
--- a/src/video_core/shader/control_flow.cpp
+++ b/src/video_core/shader/control_flow.cpp
@@ -57,8 +57,8 @@ struct BlockInfo {
 
 struct CFGRebuildState {
     explicit CFGRebuildState(const ProgramCode& program_code, const std::size_t program_size,
-                             const u32 start)
-        : start{start}, program_code{program_code}, program_size{program_size} {}
+                             const u32 start, ASTManager& manager)
+        : program_code{program_code}, program_size{program_size}, start{start}, manager{manager} {}
 
     u32 start{};
     std::vector<BlockInfo> block_info{};
@@ -71,6 +71,7 @@ struct CFGRebuildState {
     std::unordered_map<u32, BlockStack> stacks{};
     const ProgramCode& program_code;
     const std::size_t program_size;
+    ASTManager& manager;
 };
 
 enum class BlockCollision : u32 { None, Found, Inside };
@@ -455,29 +456,28 @@ void InsertBranch(ASTManager& mm, const BlockBranchInfo& branch) {
 }
 
 void DecompileShader(CFGRebuildState& state) {
-    ASTManager manager{};
+    state.manager.Init();
     for (auto label : state.labels) {
-        manager.DeclareLabel(label);
+        state.manager.DeclareLabel(label);
     }
     for (auto& block : state.block_info) {
         if (state.labels.count(block.start) != 0) {
-            manager.InsertLabel(block.start);
+            state.manager.InsertLabel(block.start);
         }
         u32 end = block.branch.ignore ? block.end + 1 : block.end;
-        manager.InsertBlock(block.start, end);
+        state.manager.InsertBlock(block.start, end);
         if (!block.branch.ignore) {
-            InsertBranch(manager, block.branch);
+            InsertBranch(state.manager, block.branch);
         }
     }
-    //manager.ShowCurrentState("Before Decompiling");
-    manager.Decompile();
-    //manager.ShowCurrentState("After Decompiling");
+    // state.manager.ShowCurrentState("Before Decompiling");
+    state.manager.Decompile();
+    // state.manager.ShowCurrentState("After Decompiling");
 }
 
-std::optional<ShaderCharacteristics> ScanFlow(const ProgramCode& program_code,
-                                              std::size_t program_size, u32 start_address) {
-    CFGRebuildState state{program_code, program_size, start_address};
-
+std::unique_ptr<ShaderCharacteristics> ScanFlow(const ProgramCode& program_code, u32 program_size,
+                                                u32 start_address, ASTManager& manager) {
+    CFGRebuildState state{program_code, program_size, start_address, manager};
     // Inspect Code and generate blocks
     state.labels.clear();
     state.labels.emplace(start_address);
@@ -503,12 +503,21 @@ std::optional<ShaderCharacteristics> ScanFlow(const ProgramCode& program_code,
               [](const BlockInfo& a, const BlockInfo& b) -> bool { return a.start < b.start; });
     if (decompiled) {
         DecompileShader(state);
+        decompiled = state.manager.IsFullyDecompiled();
+        if (!decompiled) {
+            LOG_CRITICAL(HW_GPU, "Failed to remove all the gotos!:");
+            state.manager.ShowCurrentState("Of Shader");
+            state.manager.Clear();
+        }
+    }
+    auto result_out = std::make_unique<ShaderCharacteristics>();
+    result_out->decompiled = decompiled;
+    result_out->start = start_address;
+    if (decompiled) {
+        result_out->end = state.block_info.back().end + 1;
+        return std::move(result_out);
     }
-    ShaderCharacteristics result_out{};
-    result_out.decompilable = decompiled;
-    result_out.start = start_address;
-    result_out.end = start_address;
-    for (const auto& block : state.block_info) {
+    for (auto& block : state.block_info) {
         ShaderBlock new_block{};
         new_block.start = block.start;
         new_block.end = block.end;
@@ -518,26 +527,20 @@ std::optional<ShaderCharacteristics> ScanFlow(const ProgramCode& program_code,
             new_block.branch.kills = block.branch.kill;
             new_block.branch.address = block.branch.address;
         }
-        result_out.end = std::max(result_out.end, block.end);
-        result_out.blocks.push_back(new_block);
-    }
-    if (result_out.decompilable) {
-        result_out.labels = std::move(state.labels);
-        return {std::move(result_out)};
+        result_out->end = std::max(result_out->end, block.end);
+        result_out->blocks.push_back(new_block);
     }
-
-    // If it's not decompilable, merge the unlabelled blocks together
-    auto back = result_out.blocks.begin();
+    auto back = result_out->blocks.begin();
     auto next = std::next(back);
-    while (next != result_out.blocks.end()) {
+    while (next != result_out->blocks.end()) {
         if (state.labels.count(next->start) == 0 && next->start == back->end + 1) {
             back->end = next->end;
-            next = result_out.blocks.erase(next);
+            next = result_out->blocks.erase(next);
             continue;
         }
         back = next;
         ++next;
     }
-    return {std::move(result_out)};
+    return std::move(result_out);
 }
 } // namespace VideoCommon::Shader
diff --git a/src/video_core/shader/control_flow.h b/src/video_core/shader/control_flow.h
index efd037f1a..2805d975c 100644
--- a/src/video_core/shader/control_flow.h
+++ b/src/video_core/shader/control_flow.h
@@ -10,6 +10,7 @@
 
 #include "video_core/engines/shader_bytecode.h"
 #include "video_core/shader/shader_ir.h"
+#include "video_core/shader/ast.h"
 
 namespace VideoCommon::Shader {
 
@@ -67,13 +68,12 @@ struct ShaderBlock {
 
 struct ShaderCharacteristics {
     std::list<ShaderBlock> blocks{};
-    bool decompilable{};
+    bool decompiled{};
     u32 start{};
     u32 end{};
-    std::set<u32> labels{};
 };
 
-std::optional<ShaderCharacteristics> ScanFlow(const ProgramCode& program_code,
-                                              std::size_t program_size, u32 start_address);
+std::unique_ptr<ShaderCharacteristics> ScanFlow(const ProgramCode& program_code, u32 program_size,
+                                              u32 start_address, ASTManager& manager);
 
 } // namespace VideoCommon::Shader
diff --git a/src/video_core/shader/decode.cpp b/src/video_core/shader/decode.cpp
index 47a9fd961..381e87415 100644
--- a/src/video_core/shader/decode.cpp
+++ b/src/video_core/shader/decode.cpp
@@ -39,36 +39,14 @@ void ShaderIR::Decode() {
     std::memcpy(&header, program_code.data(), sizeof(Tegra::Shader::Header));
 
     disable_flow_stack = false;
-    const auto info = ScanFlow(program_code, program_size, main_offset);
+    const auto info =
+        ScanFlow(program_code, program_size, main_offset, program_manager);
     if (info) {
         const auto& shader_info = *info;
         coverage_begin = shader_info.start;
         coverage_end = shader_info.end;
-        if (shader_info.decompilable) {
+        if (shader_info.decompiled) {
             disable_flow_stack = true;
-            const auto insert_block = [this](NodeBlock& nodes, u32 label) {
-                if (label == static_cast<u32>(exit_branch)) {
-                    return;
-                }
-                basic_blocks.insert({label, nodes});
-            };
-            const auto& blocks = shader_info.blocks;
-            NodeBlock current_block;
-            u32 current_label = static_cast<u32>(exit_branch);
-            for (auto& block : blocks) {
-                if (shader_info.labels.count(block.start) != 0) {
-                    insert_block(current_block, current_label);
-                    current_block.clear();
-                    current_label = block.start;
-                }
-                if (!block.ignore_branch) {
-                    DecodeRangeInner(current_block, block.start, block.end);
-                    InsertControlFlow(current_block, block);
-                } else {
-                    DecodeRangeInner(current_block, block.start, block.end + 1);
-                }
-            }
-            insert_block(current_block, current_label);
             return;
         }
         LOG_WARNING(HW_GPU, "Flow Stack Removing Failed! Falling back to old method");
diff --git a/src/video_core/shader/shader_ir.cpp b/src/video_core/shader/shader_ir.cpp
index 2c357f310..c79f80e04 100644
--- a/src/video_core/shader/shader_ir.cpp
+++ b/src/video_core/shader/shader_ir.cpp
@@ -23,7 +23,7 @@ using Tegra::Shader::PredOperation;
 using Tegra::Shader::Register;
 
 ShaderIR::ShaderIR(const ProgramCode& program_code, u32 main_offset, const std::size_t size)
-    : program_code{program_code}, main_offset{main_offset}, program_size{size} {
+    : program_code{program_code}, main_offset{main_offset}, program_size{size}, program_manager{} {
     Decode();
 }
 
diff --git a/src/video_core/shader/shader_ir.h b/src/video_core/shader/shader_ir.h
index 6f666ee30..a91cd7d67 100644
--- a/src/video_core/shader/shader_ir.h
+++ b/src/video_core/shader/shader_ir.h
@@ -16,6 +16,7 @@
 #include "video_core/engines/shader_bytecode.h"
 #include "video_core/engines/shader_header.h"
 #include "video_core/shader/node.h"
+#include "video_core/shader/ast.h"
 
 namespace VideoCommon::Shader {
 
@@ -364,6 +365,7 @@ private:
 
     std::map<u32, NodeBlock> basic_blocks;
     NodeBlock global_code;
+    ASTManager program_manager;
 
     std::set<u32> used_registers;
     std::set<Tegra::Shader::Pred> used_predicates;