parse_tree_levels.hpp

#pragma once
#include <memory>
#include <vector>
#include <deque>
#include "recompression.hpp"
#include "fstack.hpp"

using namespace recomp;

class ParseTreeLevels {
    rlslp<var_t> &slp_;
    // for the terminals and the non-terminals
    std::vector<size_t> levels;
    size_t num_levels;

    std::vector<std::vector<int64_t>> firsts, lasts;

    public:
    ParseTreeLevels(rlslp<var_t> &slp) : slp_(slp), levels(slp.non_terminals.size() + slp.terminals) {
        std::vector<size_t> level_order;
        level_order.reserve(slp.size() + slp.terminals);

        std::vector<size_t> outdeg(slp_.non_terminals.size() + slp_.terminals);
        for (size_t i = 0; i < slp_.non_terminals.size(); i++) {
            size_t sig = i + slp_.terminals;
            outdeg[sig] = slp_.is_block(sig) ? 1 : 2;
        }
        std::vector<std::vector<size_t>> rev_adj(slp_.non_terminals.size() + slp_.terminals);
        for (size_t i = 0; i < slp_.non_terminals.size(); i++) {
            size_t curr_idx = i + slp_.terminals;
            auto &non_term = slp_.non_terminals[i];
            if (slp_.is_block(curr_idx)) {
                rev_adj[non_term.first()].push_back(curr_idx);
            } else {
                rev_adj[non_term.first()].push_back(curr_idx);
                rev_adj[non_term.second()].push_back(curr_idx);
            }
        }
        
        size_t curr_level = 0;
        std::vector<size_t> next_bcomped;
        std::vector<size_t> next_pcomped(slp_.terminals);
        std::iota(next_pcomped.begin(), next_pcomped.end(), 0);

        auto add_sig = [&](auto sig){
            if (slp_.is_block(sig)) {
                next_bcomped.push_back(sig);
            } else {
                next_pcomped.push_back(sig);
            }
        };

        while (!next_bcomped.empty() || !next_pcomped.empty()) {
            std::vector<size_t> pcomped;
            pcomped.swap(next_pcomped);

            for (auto sig : pcomped) {
                level_order.push_back(sig);
                levels[sig] = curr_level;
                for (auto parent_sig : rev_adj[sig]) {
                    outdeg[parent_sig]--;
                    if (outdeg[parent_sig] == 0) {
                        add_sig(parent_sig);
                    }
                }
            }

            curr_level++;

            std::vector<size_t> bcomped;
            bcomped.swap(next_bcomped);

            for (auto sig : bcomped) {
                level_order.push_back(sig);
                levels[sig] = curr_level;
                for (auto parent_sig : rev_adj[sig]) {
                    outdeg[parent_sig]--;
                    if (outdeg[parent_sig] == 0) {
                        add_sig(parent_sig);
                    }
                }
            }

            curr_level++;
        }

        num_levels = curr_level;

        firsts.assign(slp.non_terminals.size() + slp.terminals, std::vector<int64_t>(num_levels, -1));
        lasts.assign(slp.non_terminals.size() + slp.terminals, std::vector<int64_t>(num_levels, -1));

        for (size_t curr_level = 0; curr_level < num_levels; curr_level++) {
            for (auto sig : level_order) {
                if (levels[sig] < curr_level) {
                    continue;
                }
                if (slp.is_terminal(sig)) {
                    firsts[sig][curr_level] = sig;
                    lasts[sig][curr_level] = sig;
                } else {
                    auto &non_term = slp.non_terminals[sig - slp.terminals];
                    if (slp.is_block(sig)) {
                        if (firsts[non_term.first()][curr_level] != -1) {
                            firsts[sig][curr_level] = firsts[non_term.first()][curr_level];
                        } else {
                            firsts[sig][curr_level] = sig;
                        }

                        if (lasts[non_term.first()][curr_level] != -1) {
                            lasts[sig][curr_level] = lasts[non_term.first()][curr_level];
                        } else {
                            lasts[sig][curr_level] = sig;
                        }
                    } else {
                        if (firsts[non_term.first()][curr_level] != -1) {
                            firsts[sig][curr_level] = firsts[non_term.first()][curr_level];
                        }
                        else {
                            firsts[sig][curr_level] = sig;
                        }

                        if (lasts[non_term.second()][curr_level] != -1) {
                            lasts[sig][curr_level] = lasts[non_term.second()][curr_level];
                        }
                        else {
                            lasts[sig][curr_level] = sig;
                        }
                    }
                }
            }
        }

        // plot();
    }

    void plot() {
        auto &slp = slp_;
        std::deque<size_t> q;
        std::vector<bool> seen(slp.non_terminals.size() + slp.terminals);
        q.push_back(slp.root);

        size_t curr_id = 0;
        std::map<size_t, size_t> nodes_to_ids;
        std::vector<std::pair<size_t, size_t>> edges;
        while (!q.empty()) {
            size_t node_idx = q.front();
            q.pop_front();

            if (seen[node_idx]) {
                continue;
            }
            seen[node_idx] = true;
            nodes_to_ids[node_idx] = curr_id++;

            if (node_idx < slp.terminals) {
                continue;
            } else {
                auto &non_term = slp.non_terminals[node_idx - slp.terminals];
                auto is_block = slp.is_block(node_idx);
                if (is_block) {
                    for (size_t i = 0; i < non_term.second(); i++) {
                        q.push_back(non_term.first());
                        edges.push_back({node_idx, non_term.first()});
                    }
                } else {
                    q.push_back(non_term.first());
                    edges.push_back({node_idx, non_term.first()});

                    q.push_back(non_term.second());
                    edges.push_back({node_idx, non_term.second()});
                }
            }
        }

        std::ofstream out("graph.dot");
        out << "digraph G {\n";
        for (auto &node : nodes_to_ids) {
            std::string label = node.first < slp.terminals ? (std::string{(char)(node.first)}) : std::to_string(node.first - slp.terminals + slp.terminals);
            out << node.second << " [label=\"" << label << "|" << levels[node.first] << "\"];\n";
        }
        int i = 0;
        for (auto &edge : edges) {
            out << nodes_to_ids[edge.first] << " -> " << nodes_to_ids[edge.second] << " [label=\"" << i << "\"];\n";
            i++;
        }
        out << "}\n";
    }

    rlslp<var_t> &slp() {
        return slp_;
    }

    size_t level(int64_t sig) {
        return levels[sig];
    }

    size_t first(int64_t sig, int64_t level) {
        return firsts[sig][level];
    }

    size_t last(int64_t sig, int64_t level) {
        return lasts[sig][level];
    }

};