#include "GVN.h"

#include "BasicBlock.h"
#include "Constant.h"
#include "DeadCode.h"
#include "FuncInfo.h"
#include "Function.h"
#include "Instruction.h"
#include "logging.hpp"

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <fstream>
#include <memory>
#include <sstream>
#include <tuple>
#include <utility>
#include <vector>

using namespace GVNExpression;
using std::string_literals::operator""s;
using std::shared_ptr;

static auto get_const_int_value = [](Value *v) {
    return dynamic_cast<ConstantInt *>(v)->get_value();
};
static auto get_const_fp_value = [](Value *v) {
    return dynamic_cast<ConstantFP *>(v)->get_value();
};

// Constant Propagation helper, folders are done for you
Constant *
ConstFolder::compute(Instruction *instr, Constant *value1, Constant *value2) {
    auto op = instr->get_instr_type();
    switch (op) {
        case Instruction::add:
            return ConstantInt::get(get_const_int_value(value1) +
                                        get_const_int_value(value2),
                                    module_);
        case Instruction::sub:
            return ConstantInt::get(get_const_int_value(value1) -
                                        get_const_int_value(value2),
                                    module_);
        case Instruction::mul:
            return ConstantInt::get(get_const_int_value(value1) *
                                        get_const_int_value(value2),
                                    module_);
        case Instruction::sdiv:
            return ConstantInt::get(get_const_int_value(value1) /
                                        get_const_int_value(value2),
                                    module_);
        case Instruction::fadd:
            return ConstantFP::get(get_const_fp_value(value1) +
                                       get_const_fp_value(value2),
                                   module_);
        case Instruction::fsub:
            return ConstantFP::get(get_const_fp_value(value1) -
                                       get_const_fp_value(value2),
                                   module_);
        case Instruction::fmul:
            return ConstantFP::get(get_const_fp_value(value1) *
                                       get_const_fp_value(value2),
                                   module_);
        case Instruction::fdiv:
            return ConstantFP::get(get_const_fp_value(value1) /
                                       get_const_fp_value(value2),
                                   module_);

        case Instruction::cmp:
            switch (dynamic_cast<CmpInst *>(instr)->get_cmp_op()) {
                case CmpInst::EQ:
                    return ConstantInt::get(get_const_int_value(value1) ==
                                                get_const_int_value(value2),
                                            module_);
                case CmpInst::NE:
                    return ConstantInt::get(get_const_int_value(value1) !=
                                                get_const_int_value(value2),
                                            module_);
                case CmpInst::GT:
                    return ConstantInt::get(get_const_int_value(value1) >
                                                get_const_int_value(value2),
                                            module_);
                case CmpInst::GE:
                    return ConstantInt::get(get_const_int_value(value1) >=
                                                get_const_int_value(value2),
                                            module_);
                case CmpInst::LT:
                    return ConstantInt::get(get_const_int_value(value1) <
                                                get_const_int_value(value2),
                                            module_);
                case CmpInst::LE:
                    return ConstantInt::get(get_const_int_value(value1) <=
                                                get_const_int_value(value2),
                                            module_);
            }
        case Instruction::fcmp:
            switch (dynamic_cast<FCmpInst *>(instr)->get_cmp_op()) {
                case FCmpInst::EQ:
                    return ConstantInt::get(get_const_fp_value(value1) ==
                                                get_const_fp_value(value2),
                                            module_);
                case FCmpInst::NE:
                    return ConstantInt::get(get_const_fp_value(value1) !=
                                                get_const_fp_value(value2),
                                            module_);
                case FCmpInst::GT:
                    return ConstantInt::get(get_const_fp_value(value1) >
                                                get_const_fp_value(value2),
                                            module_);
                case FCmpInst::GE:
                    return ConstantInt::get(get_const_fp_value(value1) >=
                                                get_const_fp_value(value2),
                                            module_);
                case FCmpInst::LT:
                    return ConstantInt::get(get_const_fp_value(value1) <
                                                get_const_fp_value(value2),
                                            module_);
                case FCmpInst::LE:
                    return ConstantInt::get(get_const_fp_value(value1) <=
                                                get_const_fp_value(value2),
                                            module_);
            }
        default:
            return nullptr;
    }
}

Constant *
ConstFolder::compute(Instruction *instr, Constant *value1) {
    auto op = instr->get_instr_type();
    switch (op) {
        case Instruction::sitofp:
            return ConstantFP::get((float)get_const_int_value(value1), module_);
        case Instruction::fptosi:
            return ConstantInt::get((int)get_const_fp_value(value1), module_);
        case Instruction::zext:
            return ConstantInt::get((int)get_const_int_value(value1), module_);
        default:
            return nullptr;
    }
}

namespace utils {
static std::string
print_congruence_class(const CongruenceClass &cc) {
    std::stringstream ss;
    if (cc.index_ == 0) {
        ss << "top class\n";
        return ss.str();
    }
    ss << "\nindex: " << cc.index_ << "\nleader: " << cc.leader_->print()
       << "\nvalue phi: "
       << (cc.value_phi_ ? cc.value_phi_->print() : "nullptr"s)
       << "\nvalue expr: "
       << (cc.value_expr_ ? cc.value_expr_->print() : "nullptr"s)
       << "\nmembers: {";
    for (auto &member : cc.members_)
        ss << member->print() << "; ";
    ss << "}\n";
    return ss.str();
}

static std::string
dump_cc_json(const CongruenceClass &cc) {
    std::string json;
    json += "[";
    for (auto member : cc.members_) {
        if (auto c = dynamic_cast<Constant *>(member))
            json += member->print() + ", ";
        else
            json += "\"%" + member->get_name() + "\", ";
    }
    json += "]";
    return json;
}

static std::string
dump_partition_json(const GVN::partitions &p) {
    std::string json;
    json += "[";
    for (auto cc : p)
        json += dump_cc_json(*cc) + ", ";
    json += "]";
    return json;
}

static std::string
dump_bb2partition(const std::map<BasicBlock *, GVN::partitions> &map) {
    std::string json;
    json += "{";
    for (auto [bb, p] : map)
        json += "\"" + bb->get_name() + "\": " + dump_partition_json(p) + ",";
    json += "}";
    return json;
}

// logging utility for you
static void
print_partitions(const GVN::partitions &p) {
    if (p.empty()) {
        // to del
        // std::cout << "empty partitions\n";
        LOG_DEBUG << "empty partitions\n";
        return;
    }
    std::string log;
    for (auto &cc : p)
        log += print_congruence_class(*cc);
    LOG_DEBUG << log; // please don't use std::cout
    // to del
    // std::cout << log;
}
} // namespace utils

GVN::partitions
GVN::join_helper(BasicBlock *pre1, BasicBlock *pre2) {
    assert(not _TOP[pre1] or not _TOP[pre2] && "should flow here, not jump");
    if (_TOP[pre1])
        return pout_[pre2];
    else if (_TOP[pre2])
        return pout_[pre1];
    return join(pout_[pre1], pout_[pre2]);
}

GVN::partitions
GVN::join(const partitions &P1, const partitions &P2) {
    // TODO: do intersection pair-wise
    partitions P{};
    for (auto c1 : P1)
        for (auto c2 : P2) {
            auto c = intersect(c1, c2);
            if (c->members_.empty())
                continue;
            P.insert(c);
        }

    return P;
}

shared_ptr<CongruenceClass>
GVN::intersect(shared_ptr<CongruenceClass> ci, shared_ptr<CongruenceClass> cj) {
    // TODO
    auto c = createCongruenceClass();
    std::set<Value *> intersection;

    std::set_intersection(ci->members_.begin(),
                          ci->members_.end(),
                          cj->members_.begin(),
                          cj->members_.end(),
                          std::inserter(intersection, intersection.begin()));
    c->members_ = intersection;
    if (ci->index_ == cj->index_)
        c->index_ = ci->index_;
    if (ci->value_expr_ == cj->value_expr_)
        c->value_expr_ = ci->value_expr_;
    if (ci->value_phi_ and cj->value_phi_ and
        *ci->value_phi_ == *cj->value_phi_)
        c->value_phi_ = ci->value_phi_;

    // ??
    // What if the ve is nullptr?
    if (c->members_.size()) // not empty
    {
        if (c->index_ == 0) {
            // we should use the `index` int the previous block
            auto instr = static_cast<Instruction *>(*c->members_.begin());
            auto instr_phi = dynamic_cast<PhiInst *>(instr);
            int exact_idx;
            // trick here: use the exact value number
            if (instr_phi) {
                auto exact_pre_bb =
                    *(instr_phi->get_parent()->get_pre_basic_blocks().begin());
                auto e = instr_phi->get_operand(
                    pretend_copy_stmt(instr_phi, exact_pre_bb));
                exact_idx = start_idx_.at(std::make_pair(instr_phi, e));
            } else {
                exact_idx = start_idx_.at(std::make_pair(instr, nullptr));
            }
            c->index_ = exact_idx;
            c->value_expr_ = c->value_phi_ =
                PhiExpression::create(ci->value_expr_, cj->value_expr_);
        }
        // ??
        c->leader_ = *c->members_.begin();
    }

    return c;
}

// assign start index for each instruction, including copy statement
// ther logic here:
// - use the same traver order as main run
// - assign an incremental number for each instruction
void
GVN::assign_start_idx_() {
    int res;
    for (auto &bb : func_->get_basic_blocks()) {
        for (auto &instr : bb.get_instructions()) {
            if (not instr.is_phi() and not instr.is_void())
                start_idx_[std::make_pair(&instr, nullptr)] = new_number();
        }
        // and the phi instruction in all the successors
        for (auto succ : bb.get_succ_basic_blocks()) {
            for (auto &instr : succ->get_instructions()) {
                if (instr.is_phi()) {
                    if ((res = pretend_copy_stmt(&instr, &bb)) == -1)
                        continue;
                    start_idx_[std::make_pair(&instr, instr.get_operand(res))] =
                        new_number();
                }
            }
        }
    }
    next_value_number_ = 1;
}

void
GVN::detectEquivalences() {
    int times = 0;
    bool changed;
    std::cout << "all the instruction address:" << std::endl;
    for (auto &bb : func_->get_basic_blocks()) {
        for (auto &instr : bb.get_instructions())
            std::cout << &instr << "\t" << instr.print() << std::endl;
    }
    assign_start_idx_();
    // initialize pout with top
    for (auto &bb : func_->get_basic_blocks()) {
        _TOP[&bb] = true;
    }

    auto Entry = func_->get_entry_block();
    _TOP[Entry] = false;

    pin_[Entry].clear();
    pout_[Entry] = transferFunction(Entry);

    // iterate until converge
    do {
        changed = false;
        std::cout << ++times << "th iteration" << std::endl;
        for (auto &_bb : func_->get_basic_blocks()) {
            auto bb = &_bb;
            if (bb == Entry)
                continue;
            // get PIN of bb from predecessor(s)
            auto pre_bbs_ = bb->get_pre_basic_blocks();
            if (bb != Entry) {
                // only update PIN for blocks that are not Entry
                // that is: the PIN for entry is always empty
                switch (pre_bbs_.size()) {
                    case 2: {
                        auto pre_1 = *pre_bbs_.begin();
                        auto pre_2 = *(++pre_bbs_.begin());
                        pin_[bb] = clone(join_helper(pre_1, pre_2));
                        break;
                    }
                    case 1: {
                        auto pre = *(pre_bbs_.begin());
                        pin_[bb] = clone(pout_[pre]);
                        break;
                    }
                    default:
                        LOG_ERROR << "block has count of predecessors: "
                                  << pre_bbs_.size();
                        abort();
                }
            }
            auto part = transferFunction(bb);

            // check changes in pout
            changed |= not(part == pout_[bb]);
            pout_[bb] = clone(part);
            _TOP[bb] = false;

            /* std::cout << "//-------\n"
             *           << "//after transferFunction(basic-block="
             *           << bb->get_name() << "), all pout:" << std::endl; */
            // dump_tmp(*func_);
        }
        // reset value number
    } while (changed);
}

// if v is already in one congruence set, return that value_expr_
// or return nullptr
shared_ptr<Expression>
GVN::search_ve(Value *v, const GVN::partitions &part) {
    for (auto c : part) {
        if (std::find(c->members_.begin(), c->members_.end(), v) !=
            c->members_.end()) {
            return c->value_expr_;
        }
    }
    return nullptr;
}

// try constant fold for `count` operands
// the related instr is `instr`
Constant *
GVN::constFold_core(const size_t count,
                    const partitions &part,
                    Instruction *instr,
                    const std::vector<Value *> &operands) {
    assert(count == 1 or count == 2);
    Constant *res = nullptr;

    // the first operand
    auto op0_const_value = dynamic_cast<Constant *>(operands[0]);
    auto op0_search_ = search_ve(operands[0], part);
    auto op0_const_expr =
        op0_search_ ? dynamic_cast<ConstantExpression *>(op0_search_.get())
                    : nullptr;
    if ((op0_const_value or op0_const_expr)) {
        auto op0_const =
            op0_const_value ? op0_const_value : op0_const_expr->get_val();
        // by now: the type cast instruction can do constant fold
        if (count == 1)
            res = folder_->compute(instr, op0_const);
        if (count == 2) {
            // the second operand
            auto op1_const_value = dynamic_cast<Constant *>(operands[1]);
            auto op1_search_ = search_ve(operands[1], part);
            auto op1_const_expr =
                op1_search_
                    ? dynamic_cast<ConstantExpression *>(op1_search_.get())
                    : nullptr;
            if (op1_const_value or op1_const_expr) {
                auto op1_const = op1_const_value ? op1_const_value
                                                 : op1_const_expr->get_val();
                res = folder_->compute(instr, op0_const, op1_const);
            }
        }
    }
    return res;
}

// for each op, try to find the
std::vector<shared_ptr<Expression>>
GVN::valueExpr_core_(Instruction *instr,
                     const partitions &part,
                     const size_t count,
                     bool fold_) {
    assert(not(fold_ and count > 2));

    Value *v;
    Constant *v_const;
    auto operands = instr->get_operands();
    std::vector<shared_ptr<Expression>> ret;

    // if able to fold, then fold
    if (fold_) {
        Constant *res = nullptr;
        if ((res = constFold_core(count, part, instr, operands))) {
            ret.push_back(ConstantExpression::create(res));
            return ret;
        }
    }

    // normal case:
    // - try to find expression that already exists
    // - take care of constant
    for (int i = 0; i != count; i++) {
        v = operands[i];
        v_const = dynamic_cast<Constant *>(v);
        if (v_const) {
            ret.push_back(ConstantExpression::create(v_const));
        } else {
            auto res = search_ve(v, part);
            assert(res);
            ret.push_back(res);
        }
    }
    return ret;
}

shared_ptr<Expression>
GVN::valueExpr(Instruction *instr, const partitions &part) {
    // TODO
    // ?? should use part?
    std::string err{"Undefined"};
    std::vector<shared_ptr<Expression>> res;

    // first check if there is already one congruence class inside
    auto tmp = search_ve(instr, part);
    if (tmp)
        return tmp;

    if (instr->isBinary() or instr->is_cmp() or instr->is_fcmp()) {
        res = valueExpr_core_(instr, part, 2);
        if (res.size() == 1) // constant fold
            return res[0];
        else
            return BinaryExpression::create(
                instr->get_instr_type(), res[0], res[1]);
    } else if (instr->is_phi()) {
        err = "phi";
    } else if (instr->is_fp2si() or instr->is_si2fp() or instr->is_zext()) {
        res = valueExpr_core_(instr, part, 1);
        if (res[0]->get_expr_type() == Expression::e_constant)
            return res[0];
        Type *dest_type;

        switch (instr->get_instr_type()) {
            case Instruction::fptosi:
                dest_type = static_cast<FpToSiInst *>(instr)->get_dest_type();
                break;
            case Instruction::sitofp:
                dest_type = static_cast<SiToFpInst *>(instr)->get_dest_type();
                break;
            case Instruction::zext:
                dest_type = static_cast<ZextInst *>(instr)->get_dest_type();
                break;
            default:
                dest_type = nullptr;
                err = "cast";
        }

        if (dest_type) {
            return CastExpression::create(
                instr->get_instr_type(), res[0], dest_type);
        }
    } else if (instr->is_gep()) {
        res = valueExpr_core_(instr, part, instr->get_operands().size(), false);
        auto ptr = res[0];
        res.erase(res.begin());
        return GEPExpression::create(ptr, res);
    } else if (instr->is_load() or instr->is_alloca() or instr->is_call()) {
        auto ret = search_ve(instr, part);
        if (ret)
            return ret;
        else
            return UniqueExpression::create(instr, next_value_number_);
    }

    std::cerr << "Undefined case: " << err << std::endl;
    abort();
}

// instruction of the form `x = e`, mostly x is just e (SSA),
// but for copy stmt x is a phi instruction in the successor.
// Phi values (not copy stmt) should be handled in detectEquiv
//
// assert the x is an instruction that can generate a new value
//
/// \param bb basic block in which the transfer function is called
GVN::partitions
GVN::transferFunction(Instruction *instr, Value *e, partitions pin) {
    next_value_number_ = start_idx_[std::make_pair(instr, e)];

    partitions pout = clone(pin);
    // TODO: deal with copy-stmt case
    // ?? deal with copy statement
    auto e_instr = dynamic_cast<Instruction *>(e);
    auto e_const = dynamic_cast<Constant *>(e);
    assert((not e or e_instr or e_const) &&
           "A value must be from an instruction or constant");
    // erase the old record for x
    std::set<Value *>::iterator it;
    for (auto c : pout)
        if ((it = std::find(c->members_.begin(), c->members_.end(), instr)) !=
            c->members_.end()) {
            c->members_.erase(it);
            if (c->members_.empty())
                pout.erase(c);
            break;
        }

    // TODO: get different ValueExpr by Instruction::OpID, modify pout
    // get ve and vpf
    shared_ptr<Expression> ve;
    if (e) {
        if (e_const)
            ve = ConstantExpression::create(e_const);
        else
            ve = valueExpr(e_instr, pout);
    } else
        ve = valueExpr(instr, pout);
    auto vpf = valuePhiFunc(ve, curr_bb, instr);

    // TODO: set leader
    for (auto c : pout) {
        if (ve == c->value_expr_ or
            (vpf and c->value_phi_ and *vpf == *c->value_phi_)) {
            // c->value_expr_ = ve;
            c->members_.insert(instr);
            return pout;
        }
    }
    auto c = createCongruenceClass(new_number());
    c->members_.insert(instr);
    c->value_expr_ = ve;
    c->value_phi_ = vpf;
    if (c->value_expr_->get_expr_type() == Expression::e_constant)
        c->leader_ =
            static_cast<ConstantExpression *>(c->value_expr_.get())->get_val();
    else
        c->leader_ = instr;
    pout.insert(c);
    return pout;
}

/*
 * read the pin for the block and then execute transferFunction() for all
 * instructions inside.
 */
GVN::partitions
GVN::transferFunction(BasicBlock *bb) {
    curr_bb = bb;
    int res;
    auto part = pin_[bb];
    /* LOG_INFO << "transferFunction(bb=" << bb->get_name() << ")\n";
     * LOG_INFO << "pin:\n";
     * utils::print_partitions(pin_[bb]);
     * LOG_INFO << "pout before:\n";
     * utils::print_partitions(pout_[bb]);  */

    // iterate through all instructions in the block
    for (auto &instr : bb->get_instructions()) {
        // ?? what about orther instructions? Are they all ok?
        if (not instr.is_phi() and not instr.is_void())
            part = transferFunction(&instr, nullptr, part);
    }
    // and the phi instruction in all the successors
    for (auto succ : bb->get_succ_basic_blocks()) {
        for (auto &instr : succ->get_instructions()) {
            if (instr.is_phi()) {
                if ((res = pretend_copy_stmt(&instr, bb)) == -1)
                    continue;
                part = transferFunction(&instr, instr.get_operand(res), part);
            }
        }
    }
    std::cout << "-------\n"
              << "for basic block " << bb->get_name() << ", pout:" << std::endl;
    utils::print_partitions(part);
    return part;
}

// NOTE: only instruction op matter
shared_ptr<PhiExpression>
GVN::valuePhiFunc(shared_ptr<Expression> ve,
                  BasicBlock *bb,
                  Instruction *instr) {
    // TODO
    // get 2 predecessors
    auto pre_bbs_ = bb->get_pre_basic_blocks();
    if (pre_bbs_.size() == 1)
        return nullptr;
    auto pre_1 = *pre_bbs_.begin();
    auto pre_2 = *(++pre_bbs_.begin());

    // check expression form
    if (ve->get_expr_type() != Expression::e_bin)
        return nullptr;
    auto ve_bin = static_cast<BinaryExpression *>(ve.get());
    if (ve_bin->lhs_->get_expr_type() != Expression::e_phi and
        ve_bin->rhs_->get_expr_type() != Expression::e_phi)
        return nullptr;

    shared_ptr<Expression> vi{}, vj{};
    shared_ptr<BinaryExpression> vl_merge{}, vr_merge{};
    shared_ptr<Expression> vl_merge_E{}, vr_merge_E{};
    // get 2 phi expressions
    auto lhs_phi = dynamic_cast<PhiExpression *>(ve_bin->lhs_.get());
    auto rhs_phi = dynamic_cast<PhiExpression *>(ve_bin->rhs_.get());

    // set vl_merge and vr_merge
    if (ve_bin->lhs_->get_expr_type() == Expression::e_phi and
        ve_bin->rhs_->get_expr_type() == Expression::e_phi) {
        // try to get the merged value expression
        vl_merge =
            BinaryExpression::create(ve_bin->op_, lhs_phi->lhs_, rhs_phi->lhs_);
        vr_merge =
            BinaryExpression::create(ve_bin->op_, lhs_phi->rhs_, rhs_phi->rhs_);
    } else if (ve_bin->lhs_->get_expr_type() == Expression::e_phi) {
        vl_merge =
            BinaryExpression::create(ve_bin->op_, lhs_phi->lhs_, ve_bin->rhs_);
        vr_merge =
            BinaryExpression::create(ve_bin->op_, lhs_phi->rhs_, ve_bin->rhs_);
    } else {
        vl_merge =
            BinaryExpression::create(ve_bin->op_, ve_bin->lhs_, rhs_phi->lhs_);
        vr_merge =
            BinaryExpression::create(ve_bin->op_, ve_bin->lhs_, rhs_phi->rhs_);
    }

    // constant fold
    if (vl_merge->lhs_->get_expr_type() == Expression::e_constant and
        vl_merge->rhs_->get_expr_type() == Expression::e_constant) {
        auto vl_merge_l =
            dynamic_cast<ConstantExpression *>(vl_merge->lhs_.get());
        auto vl_merge_r =
            dynamic_cast<ConstantExpression *>(vl_merge->rhs_.get());
        vl_merge_E = ConstantExpression::create(folder_->compute(
            instr, vl_merge_l->get_val(), vl_merge_r->get_val()));
    } else
        vl_merge_E = vl_merge;
    if (vr_merge->lhs_->get_expr_type() == Expression::e_constant and
        vr_merge->rhs_->get_expr_type() == Expression::e_constant) {
        auto vr_merge_l =
            dynamic_cast<ConstantExpression *>(vr_merge->lhs_.get());
        auto vr_merge_r =
            dynamic_cast<ConstantExpression *>(vr_merge->rhs_.get());
        vr_merge_E = ConstantExpression::create(folder_->compute(
            instr, vr_merge_l->get_val(), vr_merge_r->get_val()));
    } else
        vr_merge_E = vr_merge;

    vi = getVN(pout_[pre_1], vl_merge_E);
    vj = getVN(pout_[pre_2], vr_merge_E);
    if (vi == nullptr)
        vi = valuePhiFunc(vl_merge_E, pre_1, instr);
    if (vj == nullptr)
        vj = valuePhiFunc(vr_merge_E, pre_2, instr);

    if (vi and vj)
        return PhiExpression::create(vi, vj);
    else
        return nullptr;
}

shared_ptr<Expression>
GVN::getVN(const partitions &pout, shared_ptr<Expression> ve) {
    // TODO: return what?
    /* for (auto c : pout) {
     *     if (c->value_expr_ == ve)
     *         return ve;
     * } */
    for (auto it = pout.begin(); it != pout.end(); it++)
        if ((*it)->value_expr_ and *(*it)->value_expr_ == *ve)
            return ve;
    return nullptr;
}

void
GVN::initPerFunction() {
    next_value_number_ = 1;
    pin_.clear();
    pout_.clear();
}

void
GVN::replace_cc_members() {
    for (auto &[_bb, part] : pout_) {
        auto bb = _bb; // workaround: structured bindings can't be captured
                       // in C++17
        for (auto &cc : part) {
            if (cc->index_ == 0)
                continue;
            // if you are planning to do constant propagation, leaders
            // should be set to constant at some point
            for (auto &member : cc->members_) {
                bool member_is_phi = dynamic_cast<PhiInst *>(member);
                bool value_phi = cc->value_phi_ != nullptr;
                if (member != cc->leader_ and (value_phi or !member_is_phi)) {
                    // only replace the members if users are in the same
                    // block as bb
                    member->replace_use_with_when(
                        cc->leader_, [bb](User *user) {
                            if (auto instr =
                                    dynamic_cast<Instruction *>(user)) {
                                auto parent = instr->get_parent();
                                auto &bb_pre = parent->get_pre_basic_blocks();
                                if (instr->is_phi()) // as copy stmt, the phi
                                    // belongs to this block
                                    return std::find(bb_pre.begin(),
                                                     bb_pre.end(),
                                                     bb) != bb_pre.end();
                                else
                                    return parent == bb;
                            }
                            return false;
                        });
                }
            }
        }
    }
    return;
}

// top-level function, done for you
void
GVN::run() {
    std::ofstream gvn_json;
    if (dump_json_) {
        gvn_json.open("gvn.json", std::ios::out);
        gvn_json << "[";
    }

    folder_ = std::make_unique<ConstFolder>(m_);
    func_info_ = std::make_unique<FuncInfo>(m_);
    func_info_->run();
    dce_ = std::make_unique<DeadCode>(m_);
    dce_->run(); // let dce take care of some dead phis with undef

    for (auto &f : m_->get_functions()) {
        if (f.get_basic_blocks().empty())
            continue;
        func_ = &f;
        initPerFunction();
        LOG_INFO << "Processing " << f.get_name();
        detectEquivalences();
        LOG_INFO << "===============pin=========================\n";
        for (auto &[bb, part] : pin_) {
            LOG_INFO << "\n===============bb: " << bb->get_name()
                     << "=========================\npartitionIn: ";
            for (auto &cc : part)
                LOG_INFO << utils::print_congruence_class(*cc);
        }
        LOG_INFO << "\n===============pout=========================\n";
        for (auto &[bb, part] : pout_) {
            LOG_INFO << "\n=====bb: " << bb->get_name()
                     << "=====\npartitionOut: ";
            for (auto &cc : part)
                LOG_INFO << utils::print_congruence_class(*cc);
        }
        if (dump_json_) {
            gvn_json << "{\n\"function\": ";
            gvn_json << "\"" << f.get_name() << "\", ";
            gvn_json << "\n\"pout\": " << utils::dump_bb2partition(pout_);
            gvn_json << "},";
        }
        replace_cc_members(); // don't delete instructions, just replace
                              // them
    }
    dce_->run(); // let dce do that for us
    if (dump_json_)
        gvn_json << "]";
}

void
GVN::dump_tmp(Function &f) {
    std::string gvn_json;
    if (dump_json_) {
        gvn_json += "{\n\"function\": ";
        gvn_json += "\"" + f.get_name() + "\", ";
        gvn_json += "\n\"pout\": " + utils::dump_bb2partition(pout_);
        gvn_json += "},";
    }
    gvn_json += "]";
    std::cout << gvn_json << std::endl;
}

template <typename T>
static bool
equiv_as(const Expression &lhs, const Expression &rhs) {
    // we use static_cast because we are very sure that both operands are
    // actually T, not other types.
    return static_cast<const T *>(&lhs)->equiv(static_cast<const T *>(&rhs));
}

bool
GVNExpression::operator==(const Expression &lhs, const Expression &rhs) {
    if (lhs.get_expr_type() != rhs.get_expr_type())
        return false;
    switch (lhs.get_expr_type()) {
        case Expression::e_constant:
            return equiv_as<ConstantExpression>(lhs, rhs);
        case Expression::e_bin:
            return equiv_as<BinaryExpression>(lhs, rhs);
        case Expression::e_phi:
            return equiv_as<PhiExpression>(lhs, rhs);
        case Expression::e_cast:
            return equiv_as<CastExpression>(lhs, rhs);
        case Expression::e_gep:
            return equiv_as<GEPExpression>(lhs, rhs);
        case Expression::e_unique:
            return equiv_as<UniqueExpression>(lhs, rhs);
    }
}

bool
GVNExpression::operator==(const shared_ptr<Expression> &lhs,
                          const shared_ptr<Expression> &rhs) {
    if (lhs == nullptr and
        rhs == nullptr) // is the nullptr check necessary here?
        return true;
    return lhs and rhs and *lhs == *rhs;
}

GVN::partitions
GVN::clone(const partitions &p) {
    partitions data;
    for (auto &cc : p) {
        assert(not cc->members_.empty());
        data.insert(std::make_shared<CongruenceClass>(*cc));
    }
    return data;
}

bool
operator==(const GVN::partitions &p1, const GVN::partitions &p2) {
    // TODO: how to compare partitions?
    // cannot direct compare???
    if (p1.size() != p2.size())
        return false;
    auto it1 = p1.begin();
    auto it2 = p2.begin();
    for (; it1 != p1.end(); ++it1, ++it2)
        if (not(**it1 == **it2))
            return false;
    return true;
}

// only compare members
bool
CongruenceClass::operator==(const CongruenceClass &other) const {
    // TODO: which fields need to be compared?
    if (members_.size() != other.members_.size())
        return false;
    return members_ == other.members_;
}

int
GVN::pretend_copy_stmt(Instruction *instr, BasicBlock *bb) {
    auto phi = static_cast<PhiInst *>(instr);
    // res = phi [op1, name1], [op2, name2]
    //            ^0   ^1       ^2   ^3
    if (phi->get_operand(1)->get_name() == bb->get_name()) {
        // pretend copy statement:
        // `res = op1`
        return 0;
    } else if (phi->get_operand(3)->get_name() == bb->get_name()) {
        // `res = op2`
        return 2;
    }
    return -1;
}