From c17953978b16f82a3b2049f8b961275020c73dd0 Mon Sep 17 00:00:00 2001
From: Fernando Sahmkow <fsahmkow27@gmail.com>
Date: Thu, 27 Jun 2019 00:39:40 -0400
Subject: [PATCH] shader_ir: Initial Decompile Setup

---
 src/video_core/CMakeLists.txt          |   3 +
 src/video_core/shader/ast.cpp          | 180 ++++++++++++++++++++++++
 src/video_core/shader/ast.h            | 184 +++++++++++++++++++++++++
 src/video_core/shader/control_flow.cpp |  58 +++++++-
 src/video_core/shader/control_flow.h   |   4 +-
 src/video_core/shader/expr.h           |  86 ++++++++++++
 6 files changed, 510 insertions(+), 5 deletions(-)
 create mode 100644 src/video_core/shader/ast.cpp
 create mode 100644 src/video_core/shader/ast.h
 create mode 100644 src/video_core/shader/expr.h

diff --git a/src/video_core/CMakeLists.txt b/src/video_core/CMakeLists.txt
index e2f85c5f1..32049a2e7 100644
--- a/src/video_core/CMakeLists.txt
+++ b/src/video_core/CMakeLists.txt
@@ -105,9 +105,12 @@ add_library(video_core STATIC
     shader/decode/warp.cpp
     shader/decode/xmad.cpp
     shader/decode/other.cpp
+    shader/ast.cpp
+    shader/ast.h
     shader/control_flow.cpp
     shader/control_flow.h
     shader/decode.cpp
+    shader/expr.h
     shader/node_helper.cpp
     shader/node_helper.h
     shader/node.h
diff --git a/src/video_core/shader/ast.cpp b/src/video_core/shader/ast.cpp
new file mode 100644
index 000000000..5d0e85f42
--- /dev/null
+++ b/src/video_core/shader/ast.cpp
@@ -0,0 +1,180 @@
+// Copyright 2019 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#include <string>
+
+#include "common/assert.h"
+#include "common/common_types.h"
+#include "video_core/shader/ast.h"
+#include "video_core/shader/expr.h"
+
+namespace VideoCommon::Shader {
+
+class ExprPrinter final {
+public:
+    ExprPrinter() = default;
+
+    void operator()(ExprAnd const& expr) {
+        inner += "( ";
+        std::visit(*this, *expr.operand1);
+        inner += " && ";
+        std::visit(*this, *expr.operand2);
+        inner += ')';
+    }
+
+    void operator()(ExprOr const& expr) {
+        inner += "( ";
+        std::visit(*this, *expr.operand1);
+        inner += " || ";
+        std::visit(*this, *expr.operand2);
+        inner += ')';
+    }
+
+    void operator()(ExprNot const& expr) {
+        inner += "!";
+        std::visit(*this, *expr.operand1);
+    }
+
+    void operator()(ExprPredicate const& expr) {
+        u32 pred = static_cast<u32>(expr.predicate);
+        if (pred > 7) {
+            inner += "!";
+            pred -= 8;
+        }
+        inner += "P" + std::to_string(pred);
+    }
+
+    void operator()(ExprCondCode const& expr) {
+        u32 cc = static_cast<u32>(expr.cc);
+        inner += "CC" + std::to_string(cc);
+    }
+
+    void operator()(ExprVar const& expr) {
+        inner += "V" + std::to_string(expr.var_index);
+    }
+
+    void operator()(ExprBoolean const& expr) {
+        inner += expr.value ? "true" : "false";
+    }
+
+    std::string& GetResult() {
+        return inner;
+    }
+
+    std::string inner{};
+};
+
+class ASTPrinter {
+public:
+    ASTPrinter() = default;
+
+    void operator()(ASTProgram& ast) {
+        scope++;
+        inner += "program {\n";
+        for (ASTNode& node : ast.nodes) {
+            Visit(node);
+        }
+        inner += "}\n";
+        scope--;
+    }
+
+    void operator()(ASTIf& ast) {
+        ExprPrinter expr_parser{};
+        std::visit(expr_parser, *ast.condition);
+        inner += Ident() + "if (" + expr_parser.GetResult() + ") {\n";
+        scope++;
+        for (auto& node : ast.then_nodes) {
+            Visit(node);
+        }
+        scope--;
+        if (ast.else_nodes.size() > 0) {
+            inner += Ident() + "} else {\n";
+            scope++;
+            for (auto& node : ast.else_nodes) {
+                Visit(node);
+            }
+            scope--;
+        } else {
+            inner += Ident() + "}\n";
+        }
+    }
+
+    void operator()(ASTBlockEncoded& ast) {
+        inner += Ident() + "Block(" + std::to_string(ast.start) + ", " + std::to_string(ast.end) +
+                 ");\n";
+    }
+
+    void operator()(ASTVarSet& ast) {
+        ExprPrinter expr_parser{};
+        std::visit(expr_parser, *ast.condition);
+        inner +=
+            Ident() + "V" + std::to_string(ast.index) + " := " + expr_parser.GetResult() + ";\n";
+    }
+
+    void operator()(ASTLabel& ast) {
+        inner += "Label_" + std::to_string(ast.index) + ":\n";
+    }
+
+    void operator()(ASTGoto& ast) {
+        ExprPrinter expr_parser{};
+        std::visit(expr_parser, *ast.condition);
+        inner += Ident() + "(" + expr_parser.GetResult() + ") -> goto Label_" +
+                 std::to_string(ast.label) + ";\n";
+    }
+
+    void operator()(ASTDoWhile& ast) {
+        ExprPrinter expr_parser{};
+        std::visit(expr_parser, *ast.condition);
+        inner += Ident() + "do {\n";
+        scope++;
+        for (auto& node : ast.loop_nodes) {
+            Visit(node);
+        }
+        scope--;
+        inner += Ident() + "} while (" + expr_parser.GetResult() + ")\n";
+    }
+
+    void operator()(ASTReturn& ast) {
+        ExprPrinter expr_parser{};
+        std::visit(expr_parser, *ast.condition);
+        inner += Ident() + "(" + expr_parser.GetResult() + ") -> " +
+                 (ast.kills ? "discard" : "exit") + ";\n";
+    }
+
+    std::string& Ident() {
+        if (memo_scope == scope) {
+            return tabs_memo;
+        }
+        tabs_memo = tabs.substr(0, scope * 2);
+        memo_scope = scope;
+        return tabs_memo;
+    }
+
+    void Visit(ASTNode& node) {
+        std::visit(*this, *node->GetInnerData());
+    }
+
+    std::string& GetResult() {
+        return inner;
+    }
+
+private:
+    std::string inner{};
+    u32 scope{};
+
+    std::string tabs_memo{};
+    u32 memo_scope{};
+
+    static std::string tabs;
+};
+
+std::string ASTPrinter::tabs = "                                    ";
+
+std::string ASTManager::Print() {
+    ASTPrinter printer{};
+    printer.Visit(main_node);
+    return printer.GetResult();
+}
+
+} // namespace VideoCommon::Shader
diff --git a/src/video_core/shader/ast.h b/src/video_core/shader/ast.h
new file mode 100644
index 000000000..ca71543fb
--- /dev/null
+++ b/src/video_core/shader/ast.h
@@ -0,0 +1,184 @@
+// Copyright 2019 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include <list>
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "video_core/shader/expr.h"
+#include "video_core/shader/node.h"
+
+namespace VideoCommon::Shader {
+
+class ASTBase;
+class ASTProgram;
+class ASTIf;
+class ASTBlockEncoded;
+class ASTVarSet;
+class ASTGoto;
+class ASTLabel;
+class ASTDoWhile;
+class ASTReturn;
+
+using ASTData = std::variant<ASTProgram, ASTIf, ASTBlockEncoded, ASTVarSet, ASTGoto, ASTLabel,
+                             ASTDoWhile, ASTReturn>;
+
+using ASTNode = std::shared_ptr<ASTBase>;
+
+class ASTProgram {
+public:
+    ASTProgram() = default;
+    std::list<ASTNode> nodes;
+};
+
+class ASTIf {
+public:
+    ASTIf(Expr condition, std::list<ASTNode> then_nodes, std::list<ASTNode> else_nodes)
+        : condition(condition), then_nodes{then_nodes}, else_nodes{then_nodes} {}
+    Expr condition;
+    std::list<ASTNode> then_nodes;
+    std::list<ASTNode> else_nodes;
+};
+
+class ASTBlockEncoded {
+public:
+    ASTBlockEncoded(u32 start, u32 end) : start{start}, end{end} {}
+    u32 start;
+    u32 end;
+};
+
+class ASTVarSet {
+public:
+    ASTVarSet(u32 index, Expr condition) : index{index}, condition{condition} {}
+    u32 index;
+    Expr condition;
+};
+
+class ASTLabel {
+public:
+    ASTLabel(u32 index) : index{index} {}
+    u32 index;
+};
+
+class ASTGoto {
+public:
+    ASTGoto(Expr condition, u32 label) : condition{condition}, label{label} {}
+    Expr condition;
+    u32 label;
+};
+
+class ASTDoWhile {
+public:
+    ASTDoWhile(Expr condition, std::list<ASTNode> loop_nodes)
+        : condition(condition), loop_nodes{loop_nodes} {}
+    Expr condition;
+    std::list<ASTNode> loop_nodes;
+};
+
+class ASTReturn {
+public:
+    ASTReturn(Expr condition, bool kills) : condition{condition}, kills{kills} {}
+    Expr condition;
+    bool kills;
+};
+
+class ASTBase {
+public:
+    explicit ASTBase(ASTNode parent, ASTData data) : parent{parent}, data{data} {}
+
+    template <class U, class... Args>
+    static ASTNode Make(ASTNode parent, Args&&... args) {
+        return std::make_shared<ASTBase>(parent, ASTData(U(std::forward<Args>(args)...)));
+    }
+
+    void SetParent(ASTNode new_parent) {
+        parent = new_parent;
+    }
+
+    ASTNode& GetParent() {
+        return parent;
+    }
+
+    const ASTNode& GetParent() const {
+        return parent;
+    }
+
+    u32 GetLevel() const {
+        u32 level = 0;
+        auto next = parent;
+        while (next) {
+            next = next->GetParent();
+            level++;
+        }
+        return level;
+    }
+
+    ASTData* GetInnerData() {
+        return &data;
+    }
+
+private:
+    ASTData data;
+    ASTNode parent;
+};
+
+class ASTManager final {
+public:
+    explicit ASTManager() {
+        main_node = ASTBase::Make<ASTProgram>(nullptr);
+        program = std::get_if<ASTProgram>(main_node->GetInnerData());
+    }
+
+    void DeclareLabel(u32 address) {
+        const auto pair = labels_map.emplace(address, labels_count);
+        if (pair.second) {
+            labels_count++;
+            labels.resize(labels_count);
+        }
+    }
+
+    void InsertLabel(u32 address) {
+        u32 index = labels_map[address];
+        ASTNode label = ASTBase::Make<ASTLabel>(main_node, index);
+        labels[index] = label;
+        program->nodes.push_back(label);
+    }
+
+    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.push_back(goto_node);
+    }
+
+    void InsertBlock(u32 start_address, u32 end_address) {
+        ASTNode block = ASTBase::Make<ASTBlockEncoded>(main_node, start_address, end_address);
+        program->nodes.push_back(block);
+    }
+
+    void InsertReturn(Expr condition, bool kills) {
+        ASTNode node = ASTBase::Make<ASTReturn>(main_node, condition, kills);
+        program->nodes.push_back(node);
+    }
+
+    std::string Print();
+
+    void Decompile() {}
+
+private:
+    std::unordered_map<u32, u32> labels_map{};
+    u32 labels_count{};
+    std::vector<ASTNode> labels{};
+    std::list<ASTNode> gotos{};
+    u32 variables{};
+    ASTProgram* program;
+    ASTNode main_node;
+};
+
+} // namespace VideoCommon::Shader
diff --git a/src/video_core/shader/control_flow.cpp b/src/video_core/shader/control_flow.cpp
index ec3a76690..bea7f767c 100644
--- a/src/video_core/shader/control_flow.cpp
+++ b/src/video_core/shader/control_flow.cpp
@@ -4,13 +4,14 @@
 
 #include <list>
 #include <map>
+#include <set>
 #include <stack>
 #include <unordered_map>
-#include <unordered_set>
 #include <vector>
 
 #include "common/assert.h"
 #include "common/common_types.h"
+#include "video_core/shader/ast.h"
 #include "video_core/shader/control_flow.h"
 #include "video_core/shader/shader_ir.h"
 
@@ -64,7 +65,7 @@ struct CFGRebuildState {
     std::list<u32> inspect_queries{};
     std::list<Query> queries{};
     std::unordered_map<u32, u32> registered{};
-    std::unordered_set<u32> labels{};
+    std::set<u32> labels{};
     std::map<u32, u32> ssy_labels{};
     std::map<u32, u32> pbk_labels{};
     std::unordered_map<u32, BlockStack> stacks{};
@@ -415,6 +416,54 @@ bool TryQuery(CFGRebuildState& state) {
 }
 } // Anonymous namespace
 
+void InsertBranch(ASTManager& mm, const BlockBranchInfo& branch) {
+    const auto get_expr = ([&](const Condition& cond) -> Expr {
+        Expr result{};
+        if (cond.cc != ConditionCode::T) {
+            result = MakeExpr<ExprCondCode>(cond.cc);
+        }
+        if (cond.predicate != Pred::UnusedIndex) {
+            Expr extra = MakeExpr<ExprPredicate>(cond.predicate);
+            if (result) {
+                return MakeExpr<ExprAnd>(extra, result);
+            }
+            return extra;
+        }
+        if (result) {
+            return result;
+        }
+        return MakeExpr<ExprBoolean>(true);
+    });
+    if (branch.address < 0) {
+        if (branch.kill) {
+            mm.InsertReturn(get_expr(branch.condition), true);
+            return;
+        }
+        mm.InsertReturn(get_expr(branch.condition), false);
+        return;
+    }
+    mm.InsertGoto(get_expr(branch.condition), branch.address);
+}
+
+void DecompileShader(CFGRebuildState& state) {
+    ASTManager manager{};
+    for (auto label : state.labels) {
+        manager.DeclareLabel(label);
+    }
+    for (auto& block : state.block_info) {
+        if (state.labels.count(block.start) != 0) {
+            manager.InsertLabel(block.start);
+        }
+        u32 end = block.branch.ignore ? block.end + 1 : block.end;
+        manager.InsertBlock(block.start, end);
+        if (!block.branch.ignore) {
+            InsertBranch(manager, block.branch);
+        }
+    }
+    manager.Decompile();
+    LOG_CRITICAL(HW_GPU, "Decompiled Shader:\n{} \n", manager.Print());
+}
+
 std::optional<ShaderCharacteristics> ScanFlow(const ProgramCode& program_code,
                                               std::size_t program_size, u32 start_address) {
     CFGRebuildState state{program_code, program_size, start_address};
@@ -441,7 +490,10 @@ std::optional<ShaderCharacteristics> ScanFlow(const ProgramCode& program_code,
 
     // Sort and organize results
     std::sort(state.block_info.begin(), state.block_info.end(),
-              [](const BlockInfo& a, const BlockInfo& b) { return a.start < b.start; });
+              [](const BlockInfo& a, const BlockInfo& b) -> bool { return a.start < b.start; });
+    if (decompiled) {
+        DecompileShader(state);
+    }
     ShaderCharacteristics result_out{};
     result_out.decompilable = decompiled;
     result_out.start = start_address;
diff --git a/src/video_core/shader/control_flow.h b/src/video_core/shader/control_flow.h
index b0a5e4f8c..efd037f1a 100644
--- a/src/video_core/shader/control_flow.h
+++ b/src/video_core/shader/control_flow.h
@@ -6,7 +6,7 @@
 
 #include <list>
 #include <optional>
-#include <unordered_set>
+#include <set>
 
 #include "video_core/engines/shader_bytecode.h"
 #include "video_core/shader/shader_ir.h"
@@ -70,7 +70,7 @@ struct ShaderCharacteristics {
     bool decompilable{};
     u32 start{};
     u32 end{};
-    std::unordered_set<u32> labels{};
+    std::set<u32> labels{};
 };
 
 std::optional<ShaderCharacteristics> ScanFlow(const ProgramCode& program_code,
diff --git a/src/video_core/shader/expr.h b/src/video_core/shader/expr.h
new file mode 100644
index 000000000..94678f09a
--- /dev/null
+++ b/src/video_core/shader/expr.h
@@ -0,0 +1,86 @@
+// Copyright 2019 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include <variant>
+#include <memory>
+
+#include "video_core/engines/shader_bytecode.h"
+
+namespace VideoCommon::Shader {
+
+using Tegra::Shader::ConditionCode;
+using Tegra::Shader::Pred;
+
+class ExprAnd;
+class ExprOr;
+class ExprNot;
+class ExprPredicate;
+class ExprCondCode;
+class ExprVar;
+class ExprBoolean;
+
+using ExprData =
+    std::variant<ExprVar, ExprCondCode, ExprPredicate, ExprNot, ExprOr, ExprAnd, ExprBoolean>;
+using Expr = std::shared_ptr<ExprData>;
+
+class ExprAnd final {
+public:
+    ExprAnd(Expr a, Expr b) : operand1{a}, operand2{b} {}
+
+    Expr operand1;
+    Expr operand2;
+};
+
+class ExprOr final {
+public:
+    ExprOr(Expr a, Expr b) : operand1{a}, operand2{b} {}
+
+    Expr operand1;
+    Expr operand2;
+};
+
+class ExprNot final {
+public:
+    ExprNot(Expr a) : operand1{a} {}
+
+    Expr operand1;
+};
+
+class ExprVar final {
+public:
+    ExprVar(u32 index) : var_index{index} {}
+
+    u32 var_index;
+};
+
+class ExprPredicate final {
+public:
+    ExprPredicate(Pred predicate) : predicate{predicate} {}
+
+    Pred predicate;
+};
+
+class ExprCondCode final {
+public:
+    ExprCondCode(ConditionCode cc) : cc{cc} {}
+
+    ConditionCode cc;
+};
+
+class ExprBoolean final {
+public:
+    ExprBoolean(bool val) : value{val} {}
+
+    bool value;
+};
+
+template <typename T, typename... Args>
+Expr MakeExpr(Args&&... args) {
+    static_assert(std::is_convertible_v<T, ExprData>);
+    return std::make_shared<ExprData>(T(std::forward<Args>(args)...));
+}
+
+} // namespace VideoCommon::Shader