#include "dvm.h"
#include "IR_domTree.h"

namespace SAPFOR {
    void DominatorFinder::DFS(BasicBlock* v, int parent_num) {
        dfs_num[v] = n;
        vertex[n] = n;
        semi[n] = n;
        label[n] = n;
        ancestor[n] = -1;
        parent[n] = parent_num;
        vertices[n++] = v;

        for (const auto& w : v->getNext()) {
            if (dfs_num[w] == -1)
                DFS(w, dfs_num[v]);
        }
    }

    void DominatorFinder::Compress(int v) {
        if (ancestor[ancestor[v]] != -1) {
            Compress(ancestor[v]);

            if (semi[label[ancestor[v]]] < semi[label[v]])
                label[v] = label[ancestor[v]];
            ancestor[v] = ancestor[ancestor[v]];
        }
    }

    int DominatorFinder::Eval(int v) {
        if (ancestor[v] == -1) 
            return v;

        Compress(v);
        return label[v];
    }

    void DominatorFinder::Link(int v, int w) {
        ancestor[w] = v;
    }

    DominatorFinder::DominatorFinder(std::vector<BasicBlock*>& blocks) {
        if (blocks.empty()) 
            return;

        entry = blocks[0];
        n = 0;

        for (auto block : blocks) 
            dfs_num[block] = -1;

        int max_size = blocks.size();
        vertices.resize(max_size);
        parent.assign(max_size, -1);
        semi.assign(max_size, -1);
        vertex.assign(max_size, -1);
        ancestor.assign(max_size, -1);
        label.assign(max_size, -1);
        bucket.resize(max_size);

        DFS(entry, -1);

        for (int i = n - 1; i > 0; --i) {
            int w = vertex[i];

            for (BasicBlock* v : vertices[w]->getPrev()) {
                if (dfs_num[v] == -1)
                    continue;
                int u = Eval(dfs_num[v]);

                if (semi[u] < semi[w])
                    semi[w] = semi[u];
            }

            bucket[vertex[semi[w]]].push_back(w);
            Link(parent[w], w);

            for (int v : bucket[parent[w]])
            {
                int u = Eval(v);
                if (semi[u] < semi[v])
                    vertices[v]->setDom(vertices[u]);
                else
                    vertices[v]->setDom(vertices[parent[w]]);
            }
            bucket[parent[w]].clear();
        }

        for (int i = 1; i < n; ++i) {
            int w = vertex[i];

            if (vertices[w]->getDom() != vertices[vertex[semi[w]]])
                vertices[w]->setDom(vertices[w]->getDom()->getDom());
        }

        entry->setDom(nullptr);
    }

    void buildDominatorTree(std::vector<BasicBlock*>& blocks) {
        DominatorFinder finder(blocks);
    }
}