parse_tree_ptr.hpp

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

using namespace recomp;

// Is lce_query.hpp line 482 necessary? will it ever even trigger? i dont understand why it should

class ParseTreePtr {
    public:
    static const int64_t LEFT = -1;
    static const int64_t RIGHT = -2;
    static const int64_t NONE = -3;

    struct PtrStackItem {
        int64_t sig;

        // may want to 1-index later?
        int64_t dir;
        int64_t repr_index;
    };

    using Stack = FStack<PtrStackItem>;
    using StackPtr = std::shared_ptr<Stack>;

    ParseTreeLevels &slp_w_levels;
    StackPtr stack;

    public:
    ParseTreePtr(ParseTreeLevels &slp_w_levels_) : slp_w_levels(slp_w_levels_), stack(root_stackptr()) {
        
    }

    ParseTreePtr(ParseTreeLevels &slp_w_levels_, const StackPtr &stack_) : slp_w_levels(slp_w_levels_), stack(stack_) {
        
    }

    // NEED TO TEST
    static ParseTreePtr at(ParseTreeLevels &slp_w_levels, size_t index) {
        auto &slp = slp_w_levels.slp();
        auto len = [&](int64_t sig) -> int64_t {
            if (sig < slp.terminals) return 1;
            return slp.non_terminals[sig - slp.terminals].len;
        };
        ParseTreePtr ptr(slp_w_levels);
        while (index > 0) {
            assert(ptr.sig() >= slp.terminals);
            assert(index < slp.non_terminals[ptr.sig()-slp.terminals].len);
            auto &non_term = slp.non_terminals[ptr.sig() - slp.terminals];
            if (slp.is_block(ptr.sig())) {
                size_t child_idx = index / len(non_term.first());
                ptr = ptr.child(child_idx);
                index -= child_idx * len(non_term.first());
            } else {
                if (index < len(non_term.first())) {
                    ptr = ptr.child(0);
                } else {
                    ptr = ptr.child(1);
                    index -= len(non_term.first());
                }
            }
        }
        while (!slp.is_terminal(ptr.sig())) {
            ptr = ptr.child(0);
        }
        return ptr;
    }

    ParseTreePtr& operator=(const ParseTreePtr &other) {
        // ONCE A PTR IS CREATED, IT CAN'T BE MOVED TO A DIFFERENT PARSE TREE
        stack = other.stack;
        return *this;
    }

    ParseTreePtr push(int64_t sig, int64_t dir) const {
        assert(dir != NONE);
        auto &slp = slp_w_levels.slp();
        auto [curr_sig, curr_dir, repr_index] = Stack::top(stack);

        assert(curr_sig >= slp.terminals);
        // can't do the asserts below because for left and right, we add after collapses would have
        // occurred, so we're not adding a direct child -- maybe add a direct child check?
        if (slp.is_block(curr_sig)) {
            // assert(slp.non_terminals[curr_sig-slp.terminals].first() == sig);
        } else {
            // assert(slp.non_terminals[curr_sig-slp.terminals].first() == sig || slp.non_terminals[curr_sig-slp.terminals].second() == sig);
        }

        auto len = [&](int64_t sig) -> int64_t {
            if (sig < slp.terminals) return 1;
            return slp.non_terminals[sig - slp.terminals].len;
        };

        // assert(slp_w_levels.level(sig) < slp_w_levels.level(curr_sig));
        // changed to LEQ because it's legal if we collapse after(?)
        assert(slp_w_levels.level(sig) <= slp_w_levels.level(curr_sig));

        if (dir == LEFT) {
            return ParseTreePtr(slp_w_levels, Stack::push(stack, {sig, dir, repr_index}));
        } else if (dir == RIGHT) {
            int64_t new_repr_index = repr_index + len(curr_sig) - len(sig);
            return ParseTreePtr(slp_w_levels, Stack::push(stack, {sig, dir, new_repr_index}));
        } else {
            assert(slp.is_block(curr_sig));

            // this check is ok here bc we can't do collapsing on non left-right paths
            assert(slp.non_terminals[curr_sig-slp.terminals].first() == sig);
            int64_t new_repr_index = repr_index + dir * len(sig);
            return ParseTreePtr(slp_w_levels, Stack::push(stack, {sig, dir, new_repr_index}));
        }
    }

    ParseTreePtr parent() const {
        // NEED TO IMPL
        auto &slp = slp_w_levels.slp();
        auto [sig, dir, repr_index] = Stack::top(stack);
        if (dir > 0) {
            return ParseTreePtr(slp_w_levels, Stack::pop(stack));
        } else if (dir == LEFT) {
            auto popped = Stack::pop(stack);
            auto [prev_sig, prev_dir, prev_repr_index] = Stack::top(popped);
            assert(prev_repr_index == repr_index);
            // level+1, paper was wrong?
            auto par_sig = (int64_t)slp_w_levels.first(prev_sig, slp_w_levels.level(sig)+1);
            assert(par_sig != -1);

            auto parent_ptr = ParseTreePtr(slp_w_levels, popped);
            // modified compared to paper, but i think this is necessary - otherwise we'll keep pushing copies of the root for example, when finding fourth parent of begin ptr of "banana"
            if (par_sig == prev_sig) return parent_ptr;
            return parent_ptr.push(par_sig, LEFT).collapse();
        } else if (dir == RIGHT) {
            auto popped = Stack::pop(stack);
            auto [prev_sig, prev_dir, prev_repr_index] = Stack::top(popped);
            // level+1, paper was wrong?
            auto par_sig = (int64_t)slp_w_levels.last(prev_sig, slp_w_levels.level(sig)+1);
            assert(par_sig != -1);

            auto parent_ptr = ParseTreePtr(slp_w_levels, popped);
            if (par_sig == prev_sig) return parent_ptr;
            return parent_ptr.push(par_sig, RIGHT).collapse();
        } else {
            std::cout << "ILLEGAL CASE" << std::endl;
            exit(0);
            // ???
        }
    }

    ParseTreePtr collapse() const {
        auto top_el = Stack::top(stack);
        auto [sig, dir, repr_index] = top_el;
        auto one_back = Stack::pop(stack);
        auto [prev_sig, prev_dir, prev_repr_index] = Stack::top(one_back);
        if (dir == prev_dir && (dir == LEFT || dir == RIGHT)) {
            return ParseTreePtr(slp_w_levels, Stack::push(Stack::pop(one_back), top_el));
        }
        return *this;
    }

    ParseTreePtr child(int64_t k) const {
        auto &slp = slp_w_levels.slp();
        auto [top_sig, top_dir, top_repr_index] = Stack::top(stack);
        assert(top_sig >= (int64_t)slp.terminals);

        auto &non_term = slp.non_terminals[top_sig - slp.terminals];
        if (slp.is_block(top_sig)) {
            assert(k >= 0 && k < non_term.second());
            if (k > 0 && k < non_term.second()-1) {
                // doesn't need a collapse(?)
                return push(non_term.first(), k);
            } else {
                auto dir = k == 0 ? LEFT : RIGHT;
                return push(non_term.first(), dir).collapse();
            }
        } else {
            assert(k == 0 || k == 1);
            auto dir = k == 0 ? LEFT : RIGHT;
            auto new_sig = dir == LEFT ? non_term.first() : non_term.second();
            return push(new_sig, dir).collapse();
        }
    }

    ParseTreePtr right(int64_t l) const {
        auto p_prime = Stack::top(stack).dir != RIGHT ? stack : Stack::pop(stack);
        auto q = ParseTreePtr(slp_w_levels, p_prime).parent();
        int64_t j = Stack::top(p_prime).dir == LEFT ? 1 : (Stack::top(p_prime).dir + 1);

        // in the paper, but what the hell is that? isn't it just getting a child of q?
        // auto &non_term_p_prime = slp.non_terminals[Stack::top(q.stack).sig - slp.terminals];
        // auto num_children_q = slp.is_block(Stack::top(q.stack).sig) ? non_term_p_prime.second() : 2;
        // int64_t y = j == (num_children_q-1) ? RIGHT : j;

        // auto child_sig = q.child(j).sig();
        // auto new_repr_index = Stack::top(q.stack).repr_index + (int64_t)slp.non_terminals[child_sig - slp.terminals].len;

        // auto r = ParseTreePtr(slp_w_levels, Stack::push(q.stack, {child_sig, y, })).collapse();

        auto r = q.child(j).collapse();
        auto b_level = slp_w_levels.level(r.sig());
        if (b_level <= l) return r;

        auto first_sig = (int64_t)(slp_w_levels.first(r.sig(), l));
        assert(first_sig != -1);
        auto ret = r.push(first_sig, LEFT);

        // not in the paper, but i think the paper just forgot abt this??? otherwise need a diff "first" array
        if (slp_w_levels.level(ret.sig()) <= l) return ret;
        return ret.child(0);
    }

    ParseTreePtr left(int64_t l) const {
        auto p_prime = Stack::top(stack).dir != LEFT ? stack : Stack::pop(stack);
        auto q = ParseTreePtr(slp_w_levels, p_prime).parent();
        int64_t j = Stack::top(p_prime).dir == RIGHT ? (q.degree()-2) : (Stack::top(p_prime).dir - 1);

        auto r = q.child(j).collapse();
        auto b_level = slp_w_levels.level(r.sig());
        if (b_level <= l) return r;

        // is returning 255 when we want it to return 97 -- wanted one lower level?
        auto last_sig = (int64_t)(slp_w_levels.last(r.sig(), l));
        assert(last_sig != -1);
        auto ret = r.push(last_sig, RIGHT);

        // not in the paper, but i think the paper just forgot abt this??? otherwise need a diff "last" array
        if (slp_w_levels.level(ret.sig()) <= l) return ret;
        return ret.child(ret.degree()-1);
    }

    StackPtr root_stackptr() const {
        auto &slp = slp_w_levels.slp();
        return Stack::push(Stack::empty(), {slp.root, NONE, 0});
    }

    ParseTreePtr root() const {
        return ParseTreePtr(slp_w_levels, root_stackptr());
    }

    ParseTreePtr begin() const {
        int64_t first_sig = slp_w_levels.first(sig(), 0);
        assert(first_sig != -1);
        return root().push(first_sig, LEFT);
    }

    ParseTreePtr end() const {
        int64_t last_sig = slp_w_levels.last(sig(), 0);
        assert(last_sig != -1);
        return root().push(last_sig, RIGHT);
    }

    int64_t sig() const {
        return Stack::top(stack).sig;
    }

    int64_t degree() const {
        auto &slp = slp_w_levels.slp();
        int64_t curr_sig = sig();
        if (slp.is_terminal(curr_sig)) {
            return 0;
        } else if (slp.is_block(curr_sig)) {
            return slp.non_terminals[curr_sig - slp.terminals].second();
        } else {
            return 2;
        }
    }

    // may want to 1-index later?
    int64_t index() const {
        auto dir = Stack::top(stack).dir;
        if (dir >= 0) return dir;

        // should I 1-index?
        if (dir == LEFT) return 0;

        // should I 1-index?
        if (dir == RIGHT) return parent().degree() - 1;

        assert(false);

        return -1;
    }

    int64_t level() const {
        return slp_w_levels.level(sig());
    }

    // closed, closed, ie [a,b]
    std::pair<int64_t, int64_t> repr() const {
        auto &slp = slp_w_levels.slp();
        auto [sig, dir, repr_index] = Stack::top(stack);
        int64_t len = sig < slp.terminals ? 1 : slp.non_terminals[sig - slp.terminals].len;
        return {repr_index, repr_index + len - 1};
    }

    bool operator==(const ParseTreePtr &other) const {
        return repr() == other.repr();
    }

    std::string to_string() const {
        if (stack == Stack::empty()) return "";
        auto &slp = slp_w_levels.slp();
        auto [sig, dir, repr_index] = Stack::top(stack);

        std::string rep = "";
        rep += ParseTreePtr(slp_w_levels, Stack::pop(stack)).to_string();
        std::string dir_s = dir == LEFT ? "L" : (dir == RIGHT ? "R" : std::to_string(dir));
        rep += std::string{} + " {" + std::to_string(sig) + "," + dir_s + "}";

        return rep;
    }
};