From fbdb1a3d61868c97f03ebb03f31e9deed18caaea Mon Sep 17 00:00:00 2001
From: Dan Klishch <danilklishch@gmail.com>
Date: Sun, 1 Oct 2023 23:01:10 -0400
Subject: [PATCH] JSSpecCompiler: Add SSA form building pass

---
 .../JSSpecCompiler/CMakeLists.txt             |   1 +
 .../Compiler/Passes/SSABuildingPass.cpp       | 435 ++++++++++++++++++
 .../Compiler/Passes/SSABuildingPass.h         |  91 ++++
 .../Tests/simple.cpp.expectation              |  69 +++
 .../CodeGenerators/JSSpecCompiler/main.cpp    |   2 +
 5 files changed, 598 insertions(+)
 create mode 100644 Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/Passes/SSABuildingPass.cpp
 create mode 100644 Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/Passes/SSABuildingPass.h

diff --git a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/CMakeLists.txt b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/CMakeLists.txt
index 431bd1ada1..3206048ed1 100644
--- a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/CMakeLists.txt
+++ b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/CMakeLists.txt
@@ -9,6 +9,7 @@ set(SOURCES
     Compiler/Passes/FunctionCallCanonicalizationPass.cpp
     Compiler/Passes/IfBranchMergingPass.cpp
     Compiler/Passes/ReferenceResolvingPass.cpp
+    Compiler/Passes/SSABuildingPass.cpp
     Parser/CppASTConverter.cpp
     Parser/Lexer.cpp
     Parser/ParseError.cpp
diff --git a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/Passes/SSABuildingPass.cpp b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/Passes/SSABuildingPass.cpp
new file mode 100644
index 0000000000..d10ba7cf07
--- /dev/null
+++ b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/Passes/SSABuildingPass.cpp
@@ -0,0 +1,435 @@
+/*
+ * Copyright (c) 2023, Dan Klishch <danilklishch@gmail.com>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include <AK/Queue.h>
+
+#include "AST/AST.h"
+#include "Compiler/GenericASTPass.h"
+#include "Compiler/Passes/SSABuildingPass.h"
+#include "Function.h"
+
+namespace JSSpecCompiler {
+
+void SSABuildingPass::process_function()
+{
+    m_dtree_timer = 0;
+    m_order.clear();
+    m_mark_version = 1;
+    m_def_stack.clear();
+    m_next_id.clear();
+    m_undo_vector.clear();
+    m_graph = m_function->m_cfg;
+
+    with_graph(m_graph->blocks_count(), [&] {
+        compute_dominator_tree();
+        compute_dominance_frontiers();
+        place_phi_nodes();
+        rename_variables();
+    });
+}
+
+// ===== compute_dominator_tree =====
+namespace {
+class DSU {
+    struct NodeData {
+        size_t sdom;
+        size_t parent;
+    };
+
+public:
+    DSU(size_t n)
+        : n(n)
+    {
+        m_nodes.resize(n);
+        for (size_t i = 0; i < n; ++i)
+            m_nodes[i] = { i, i };
+    }
+
+    NodeData get(size_t u)
+    {
+        if (m_nodes[u].parent == u)
+            return { n, u };
+        auto [sdom, root] = get(m_nodes[u].parent);
+        sdom = min(sdom, m_nodes[u].sdom);
+        return m_nodes[u] = { sdom, root };
+    }
+
+    void merge(size_t u, size_t v, size_t v_sdom)
+    {
+        m_nodes[v] = { v_sdom, u };
+    }
+
+private:
+    size_t n;
+    Vector<NodeData> m_nodes;
+};
+}
+
+void SSABuildingPass::compute_order(BasicBlockRef u, Vertex parent)
+{
+    if (m_nodes[u->m_index].is_used)
+        return;
+    m_nodes[u->m_index].is_used = true;
+
+    Vertex reordered_u = m_order.size();
+    m_order.append(RefPtr<BasicBlock>(u).release_nonnull());
+    reordered_u->parent = parent;
+
+    for (auto* v : u->m_continuation->references())
+        compute_order(*v, reordered_u);
+}
+
+void SSABuildingPass::compute_dominator_tree()
+{
+    size_t n = m_graph->blocks_count();
+    m_nodes.resize(n);
+
+    // Algorithm is from https://tanujkhattar.wordpress.com/2016/01/11/dominator-tree-of-a-directed-graph/ ,
+    // an author writes awful CP-style write-only code, but the explanation is pretty good.
+
+    // Step 1
+    compute_order(m_graph->start_block);
+    VERIFY(m_order.size() == n);
+    for (size_t i = 0; i < n; ++i)
+        m_order[i]->m_index = i;
+    m_graph->blocks = m_order;
+
+    for (size_t i = 0; i < n; ++i) {
+        Vertex u = i;
+
+        for (auto* reference : u.block()->m_continuation->references()) {
+            Vertex v { *reference };
+
+            v->incoming_edges.append(u);
+            u->outgoing_edges.append(v);
+        }
+    }
+
+    // Steps 2 & 3
+    DSU dsu(n);
+
+    for (size_t i = n - 1; i > 0; --i) {
+        Vertex u = i;
+
+        Vertex& current_sdom = u->semi_dominator;
+        current_sdom = n;
+
+        for (Vertex v : u->incoming_edges) {
+            if (v < u)
+                current_sdom = min(current_sdom, v);
+            else
+                current_sdom = min(current_sdom, dsu.get(v).sdom);
+        }
+
+        current_sdom->buckets.append(u);
+        for (Vertex w : u->buckets) {
+            Vertex v = dsu.get(w).sdom;
+
+            if (v->semi_dominator == w->semi_dominator)
+                w->immediate_dominator = v->semi_dominator;
+            else
+                w->immediate_dominator = v;
+        }
+        dsu.merge(u->parent, u, current_sdom);
+    }
+
+    m_nodes[0].immediate_dominator = invalid_node;
+    for (size_t i = 1; i < n; ++i) {
+        Vertex u = i;
+
+        if (u->immediate_dominator.is_invalid())
+            u->immediate_dominator = 0;
+        else if (u->immediate_dominator != u->semi_dominator)
+            u->immediate_dominator = u->immediate_dominator->immediate_dominator;
+    }
+
+    // Populate dtree_children & BasicBlock::immediate_dominator
+    for (size_t i = 0; i < n; ++i) {
+        Vertex u = i;
+
+        if (i != 0) {
+            u.block()->m_immediate_dominator = u->immediate_dominator.block();
+            u->immediate_dominator->dtree_children.append(u);
+        } else {
+            u.block()->m_immediate_dominator = nullptr;
+        }
+    }
+}
+
+// ===== compute_dominance_frontiers =====
+template<typename... Args>
+Vector<SSABuildingPass::Vertex> SSABuildingPass::unique(Args const&... args)
+{
+    ++m_mark_version;
+
+    Vector<Vertex> result;
+    (([&](auto const& list) {
+        for (Vertex u : list) {
+            if (u->mark != m_mark_version) {
+                u->mark = m_mark_version;
+                result.append(u);
+            }
+        }
+    })(args),
+        ...);
+    return result;
+}
+
+void SSABuildingPass::compute_dtree_tin_tout(Vertex u)
+{
+    u->tin = m_dtree_timer++;
+    for (Vertex v : u->dtree_children)
+        compute_dtree_tin_tout(v);
+    u->tout = m_dtree_timer++;
+}
+
+bool SSABuildingPass::is_strictly_dominating(Vertex u, Vertex v)
+{
+    return u != v && u->tin <= v->tin && v->tout <= u->tout;
+}
+
+void SSABuildingPass::compute_dominance_frontiers()
+{
+    compute_dtree_tin_tout(0);
+
+    // Algorithm from https://en.wikipedia.org/wiki/Static_single-assignment_form#Converting%20to%20SSA:~:text=their%20paper%20titled-,A%20Simple%2C%20Fast%20Dominance%20Algorithm,-%3A%5B13%5D .
+    // DF(u) = {w : !(u sdom w) /\ (\exists v \in incoming_edges(v) : u dom v)}
+    for (size_t wi = 0; wi < m_nodes.size(); ++wi) {
+        Vertex w = wi;
+
+        for (Vertex v : w->incoming_edges) {
+            Vertex u = v;
+            while (u != invalid_node && !is_strictly_dominating(u, w)) {
+                u->d_frontier.append(w);
+                u = u->immediate_dominator;
+            }
+        }
+    }
+
+    for (size_t i = 0; i < m_nodes.size(); ++i) {
+        Vertex u = i;
+
+        u->d_frontier = unique(u->d_frontier);
+    }
+}
+
+// ===== place_phi_nodes =====
+namespace {
+class VariableAssignmentCollector : private RecursiveASTVisitor {
+public:
+    VariableAssignmentCollector(OrderedHashMap<NamedVariableDeclarationRef, Vector<BasicBlockRef>>& declarations)
+        : m_declarations(declarations)
+    {
+    }
+
+    void run(BasicBlockRef block)
+    {
+        m_current_block = block;
+
+        for (auto& expression : block->m_expressions)
+            run_in_subtree(expression);
+        run_in_const_subtree(block->m_continuation);
+    }
+
+protected:
+    RecursionDecision on_entry(Tree tree) override
+    {
+        if (tree->is_statement())
+            TODO();
+        return RecursionDecision::Recurse;
+    }
+
+    void on_leave(Tree tree) override
+    {
+        if (auto binary_operation = as<BinaryOperation>(tree); binary_operation) {
+            if (binary_operation->m_operation != BinaryOperator::Assignment)
+                return;
+
+            if (auto variable = as<Variable>(binary_operation->m_left); variable) {
+                auto& vector = m_declarations.get(variable->m_name).value();
+                if (vector.is_empty() || vector.last() != m_current_block)
+                    vector.append(m_current_block);
+            }
+        }
+    }
+
+private:
+    BasicBlockRef m_current_block;
+    OrderedHashMap<NamedVariableDeclarationRef, Vector<BasicBlockRef>>& m_declarations;
+};
+}
+
+void SSABuildingPass::add_phi_node(BasicBlockRef block, NamedVariableDeclarationRef decl)
+{
+    BasicBlock::PhiNode node { .var = make_ref_counted<Variable>(decl) };
+    for (Vertex incoming : Vertex(block)->incoming_edges) {
+        BasicBlockRef incoming_block = incoming.block();
+        auto value = make_ref_counted<Variable>(decl);
+        node.branches.append({ .block = incoming_block, .value = value });
+    }
+    block->m_phi_nodes.append(move(node));
+}
+
+void SSABuildingPass::place_phi_nodes()
+{
+    // Entry block has implicit declarations of all variables.
+    OrderedHashMap<NamedVariableDeclarationRef, Vector<BasicBlockRef>> m_declarations;
+    for (auto const& [name, var_decl] : m_function->m_local_variables)
+        m_declarations.set(var_decl, { m_order[0] });
+    m_declarations.set(m_function->m_return_value, { m_order[0] });
+
+    VariableAssignmentCollector collector(m_declarations);
+    for (auto const& block : m_order)
+        collector.run(block);
+
+    for (auto const& [decl, blocks] : m_declarations) {
+        ++m_mark_version;
+
+        Queue<BasicBlockRef> queue;
+        for (auto const& block : blocks)
+            queue.enqueue(block);
+
+        while (!queue.is_empty()) {
+            Vertex u(queue.dequeue());
+
+            for (Vertex frontier : u->d_frontier) {
+                if (frontier->mark == m_mark_version)
+                    continue;
+                frontier->mark = m_mark_version;
+                add_phi_node(frontier.block(), decl);
+            }
+        }
+    }
+}
+
+// ===== rename_variables =====
+namespace {
+template<typename CreateSSAVariableFunc, typename RenameVariableFunc>
+class VariableRenamer : private RecursiveASTVisitor {
+public:
+    VariableRenamer(CreateSSAVariableFunc create, RenameVariableFunc rename)
+        : m_create(create)
+        , m_rename(rename)
+    {
+    }
+
+    void run(BasicBlockRef block)
+    {
+        for (auto& expression : block->m_expressions)
+            run_in_subtree(expression);
+        run_in_const_subtree(block->m_continuation);
+    }
+
+protected:
+    RecursionDecision on_entry(Tree tree) override
+    {
+        if (tree->is_statement())
+            TODO();
+
+        auto binary_operation = as<BinaryOperation>(tree);
+        if (binary_operation && binary_operation->m_operation == BinaryOperator::Assignment) {
+            run_in_subtree(binary_operation->m_right);
+            if (auto variable = as<Variable>(binary_operation->m_left); variable) {
+                m_create(variable->m_name);
+                m_rename(variable.release_nonnull());
+            } else {
+                run_in_subtree(binary_operation->m_left);
+            }
+            return RecursionDecision::Continue;
+        }
+
+        if (auto variable = as<Variable>(tree); variable) {
+            m_rename(variable.release_nonnull());
+            return RecursionDecision::Continue;
+        }
+
+        return RecursionDecision::Recurse;
+    }
+
+private:
+    CreateSSAVariableFunc m_create;
+    RenameVariableFunc m_rename;
+};
+}
+
+void SSABuildingPass::make_new_ssa_variable_for(NamedVariableDeclarationRef var)
+{
+    m_undo_vector.append(var);
+
+    u64 id = 0;
+    if (auto it = m_next_id.find(var); it == m_next_id.end())
+        m_next_id.set(var, 1);
+    else
+        id = it->value++;
+    auto ssa_decl = make_ref_counted<SSAVariableDeclaration>(id);
+
+    m_function->m_local_ssa_variables.append(ssa_decl);
+
+    if (auto it = m_def_stack.find(var); it == m_def_stack.end())
+        m_def_stack.set(var, { ssa_decl });
+    else
+        it->value.append(ssa_decl);
+}
+
+void SSABuildingPass::rename_variable(VariableRef var)
+{
+    var->m_ssa = m_def_stack.get(var->m_name).value().last();
+}
+
+void SSABuildingPass::rename_variables(Vertex u, Vertex from)
+{
+    size_t rollback_point = m_undo_vector.size();
+
+    for (auto& phi_node : u.block()->m_phi_nodes) {
+        // TODO: Find the right branch index without iterating through all of the branches.
+        bool found = false;
+        for (auto& branch : phi_node.branches) {
+            if (branch.block->m_index == from) {
+                rename_variable(branch.value);
+                found = true;
+                break;
+            }
+        }
+        VERIFY(found);
+    }
+
+    if (u->mark == m_mark_version)
+        return;
+    u->mark = m_mark_version;
+
+    for (auto& phi_node : u.block()->m_phi_nodes) {
+        make_new_ssa_variable_for(phi_node.var->m_name);
+        rename_variable(phi_node.var);
+    }
+
+    VariableRenamer renamer(
+        [&](NamedVariableDeclarationRef decl) {
+            make_new_ssa_variable_for(move(decl));
+        },
+        [&](VariableRef var) {
+            rename_variable(move(var));
+        });
+    renamer.run(u.block());
+
+    for (size_t j : u->outgoing_edges)
+        rename_variables(j, u);
+
+    while (m_undo_vector.size() > rollback_point)
+        (void)m_def_stack.get(m_undo_vector.take_last()).value().take_last();
+}
+
+void SSABuildingPass::rename_variables()
+{
+    for (auto const& [name, var_decl] : m_function->m_local_variables)
+        make_new_ssa_variable_for(var_decl);
+    make_new_ssa_variable_for(m_function->m_return_value);
+
+    ++m_mark_version;
+    rename_variables(0);
+    m_function->reindex_ssa_variables();
+}
+
+}
diff --git a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/Passes/SSABuildingPass.h b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/Passes/SSABuildingPass.h
new file mode 100644
index 0000000000..be644c18d5
--- /dev/null
+++ b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Compiler/Passes/SSABuildingPass.h
@@ -0,0 +1,91 @@
+/*
+ * Copyright (c) 2023, Dan Klishch <danilklishch@gmail.com>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <AK/HashMap.h>
+
+#include "Compiler/CompilerPass.h"
+#include "Compiler/ControlFlowGraph.h"
+#include "Compiler/EnableGraphPointers.h"
+
+namespace JSSpecCompiler {
+
+// TODO: Add a LOT of unit tests.
+
+class SSABuildingPass
+    : public IntraproceduralCompilerPass
+    , private EnableGraphPointers<SSABuildingPass, BasicBlockRef> {
+public:
+    inline static constexpr StringView name = "ssa-building"sv;
+
+    using IntraproceduralCompilerPass::IntraproceduralCompilerPass;
+
+protected:
+    void process_function() override;
+
+private:
+    friend EnableGraphPointers;
+
+    class Vertex : public VertexBase {
+    public:
+        using VertexBase::VertexBase;
+
+        BasicBlockRef block() const { return m_instance->m_order[m_index]; }
+    };
+
+    void compute_order(BasicBlockRef u, Vertex parent = invalid_node);
+    void compute_dominator_tree();
+
+    template<typename... Args>
+    Vector<Vertex> unique(Args const&... args);
+    void compute_dtree_tin_tout(Vertex u);
+    bool is_strictly_dominating(Vertex u, Vertex v);
+    void compute_dominance_frontiers();
+
+    void add_phi_node(BasicBlockRef block, NamedVariableDeclarationRef decl);
+    void place_phi_nodes();
+
+    void make_new_ssa_variable_for(NamedVariableDeclarationRef var);
+    void rename_variable(VariableRef var);
+    void rename_variables(Vertex u, Vertex from = invalid_node);
+    void rename_variables();
+
+    struct NodeData {
+        Vector<Vertex> incoming_edges;
+        Vector<Vertex> outgoing_edges;
+
+        Vector<Vertex> buckets;
+
+        bool is_used = false;
+        Vertex parent;
+        Vertex semi_dominator;
+        Vertex immediate_dominator;
+
+        Vector<Vertex> dtree_children;
+        u64 tin, tout;
+
+        Vector<Vertex> d_frontier;
+
+        HashMap<NamedVariableDeclarationRef, Vertex> phi_nodes;
+
+        u64 mark = 0;
+    };
+
+    u64 m_dtree_timer;
+    Vector<NodeData> m_nodes;
+    Vector<NonnullRefPtr<BasicBlock>> m_order;
+
+    u64 m_mark_version;
+
+    HashMap<NamedVariableDeclarationRef, Vector<SSAVariableDeclarationRef>> m_def_stack;
+    HashMap<NamedVariableDeclarationRef, u64> m_next_id;
+    Vector<NamedVariableDeclarationRef> m_undo_vector;
+
+    ControlFlowGraph* m_graph;
+};
+
+}
diff --git a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Tests/simple.cpp.expectation b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Tests/simple.cpp.expectation
index 25917762db..2ecb952867 100644
--- a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Tests/simple.cpp.expectation
+++ b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/Tests/simple.cpp.expectation
@@ -249,3 +249,72 @@ BinaryOperation Assignment
   MathematicalConstant 4
 ControlFlowJump jump=2
 
+===== AST after ssa-building =====
+f():
+TreeList
+  IfElseIfChain
+    UnresolvedReference cond1
+    TreeList
+      BinaryOperation Assignment
+        Var a@1
+        MathematicalConstant 1
+      IfElseIfChain
+        UnresolvedReference cond2
+        TreeList
+          BinaryOperation Assignment
+            Var b@1
+            Var a@1
+        TreeList
+          BinaryOperation Assignment
+            Var b@3
+            MathematicalConstant 3
+    TreeList
+      BinaryOperation Assignment
+        Var b@4
+        MathematicalConstant 4
+  ReturnNode
+    Var b@2
+
+===== CFG after ssa-building =====
+f():
+0:
+ControlFlowBranch true=1 false=6
+  UnresolvedReference cond1
+
+1:
+BinaryOperation Assignment
+  Var a@1
+  MathematicalConstant 1
+ControlFlowBranch true=2 false=5
+  UnresolvedReference cond2
+
+2:
+BinaryOperation Assignment
+  Var b@1
+  Var a@1
+ControlFlowJump jump=3
+
+3:
+a@2 = phi(2: a@1, 5: a@1, 6: a@0)
+b@2 = phi(2: b@1, 5: b@3, 6: b@4)
+BinaryOperation Assignment
+  Var $return@1
+  Var b@2
+ControlFlowJump jump=4
+
+4:
+ControlFlowFunctionReturn
+  Var $return@1
+
+5:
+BinaryOperation Assignment
+  Var b@3
+  MathematicalConstant 3
+ControlFlowJump jump=3
+
+6:
+BinaryOperation Assignment
+  Var b@4
+  MathematicalConstant 4
+ControlFlowJump jump=3
+
diff --git a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/main.cpp b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/main.cpp
index 3b8b097781..b8cbf3bb95 100644
--- a/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/main.cpp
+++ b/Meta/Lagom/Tools/CodeGenerators/JSSpecCompiler/main.cpp
@@ -13,6 +13,7 @@
 #include "Compiler/Passes/FunctionCallCanonicalizationPass.h"
 #include "Compiler/Passes/IfBranchMergingPass.h"
 #include "Compiler/Passes/ReferenceResolvingPass.h"
+#include "Compiler/Passes/SSABuildingPass.h"
 #include "Function.h"
 #include "Parser/CppASTConverter.h"
 #include "Parser/SpecParser.h"
@@ -108,6 +109,7 @@ ErrorOr<int> serenity_main(Main::Arguments arguments)
     pipeline.add_compilation_pass<ReferenceResolvingPass>();
     pipeline.add_compilation_pass<CFGBuildingPass>();
     pipeline.add_compilation_pass<CFGSimplificationPass>();
+    pipeline.add_compilation_pass<SSABuildingPass>();
 
     pipeline.for_each_step_in(passes_to_dump_ast, [](CompilationStepWithDumpOptions& step) {
         step.dump_ast = true;