// See the file "COPYING" in the main distribution directory for copyright.

#include "zeek/script_opt/UseDefs.h"

#include "zeek/Desc.h"
#include "zeek/Reporter.h"
#include "zeek/Stmt.h"
#include "zeek/script_opt/Reduce.h"
#include "zeek/script_opt/ScriptOpt.h"

namespace zeek::detail {

void UseDefSet::Dump() const {
    for ( const auto& u : IterateOver() )
        printf(" %s", u->Name());
}

UseDefs::UseDefs(StmtPtr _body, std::shared_ptr<Reducer> _rc, FuncTypePtr _ft) {
    body = std::move(_body);
    rc = std::move(_rc);
    ft = std::move(_ft);
}

void UseDefs::Analyze() {
    // Start afresh.
    use_defs_map.clear();
    UDs_are_copies.clear();
    stmts.clear();
    successor.clear();
    successor2.clear();

    (void)PropagateUDs(body, nullptr, nullptr, false);
}

StmtPtr UseDefs::RemoveUnused() {
    int iter = 0;
    while ( RemoveUnused(++iter) ) {
        body = rc->Reduce(body);
        Analyze();

        if ( reporter->Errors() > 0 )
            break;
    }

    return body;
}

void UseDefs::Dump() {
    for ( int i = stmts.size(); --i >= 0; ) {
        const auto& s = stmts[i];
        auto uds = FindUsage(s);
        auto are_copies = (UDs_are_copies.find(s) != UDs_are_copies.end());

        printf("UDs (%s) for %s:\n", are_copies ? "copy" : "orig", obj_desc(s).c_str());

        if ( uds )
            uds->Dump();
        else
            printf(" <none>");

        printf("\n\n");
    }
}

bool UseDefs::RemoveUnused(int iter) {
    rc->ResetAlteredStmts();

    bool did_omission = false;

    for ( const auto& s : stmts ) {
        if ( s->Tag() == STMT_INIT ) {
            auto init = s->AsInitStmt();
            const auto& inits = init->Inits();
            std::vector<IDPtr> used_ids;

            for ( const auto& id : inits )
                if ( is_atomic_type(id->GetType()) || ! CheckIfUnused(s, id.get(), false) )
                    used_ids.emplace_back(id);

            if ( used_ids.empty() ) { // There aren't any ID's to keep.
                rc->AddStmtToOmit(s);
                continue;
            }

            if ( used_ids.size() < inits.size() ) {
                // Need to replace the current Init statement
                // with one that only includes the actually
                // used identifiers.

                auto new_init = with_location_of(make_intrusive<InitStmt>(used_ids), s);
                rc->AddStmtToReplace(s, std::move(new_init));
            }

            continue;
        }

        // The only other statements we might revise or remove
        // are assignments.

        if ( s->Tag() != STMT_EXPR )
            continue;

        auto s_e = s->AsExprStmt();
        auto e = s_e->StmtExpr();

        if ( e->Tag() != EXPR_ASSIGN )
            continue;

        auto a = e->AsAssignExpr();
        auto r = a->GetOp1();

        // Because we're dealing with reduced statements, the
        // assignment expression should be to a simple variable.
        if ( r->Tag() != EXPR_REF )
            reporter->InternalError("lhs ref inconsistency in UseDefs::RemoveUnused");

        auto n = r->AsRefExprPtr()->GetOp1();
        if ( n->Tag() != EXPR_NAME )
            reporter->InternalError("lhs name inconsistency in UseDefs::RemoveUnused");

        auto id = n->AsNameExpr()->Id();

        auto rhs = a->GetOp2();
        auto rt = rhs->Tag();

        if ( rt == EXPR_CALL && ! rhs->IsPure() )
            // Need to do the call for the side effects.
            // Could prune out the assignment and just
            // keep the call, but not clear that that's
            // worth the complexity.
            continue;

        if ( rt == EXPR_EVENT || rt == EXPR_SCHEDULE )
            // These always have side effects.
            continue;

        // Check for degenerate assignment "x = x".
        bool degen = rt == EXPR_NAME && id == rhs->AsNameExpr()->Id();

        if ( CheckIfUnused(s, id, iter == 1) || degen ) {
            rc->AddStmtToOmit(s);
            did_omission = true;
        }
    }

    return did_omission;
}

bool UseDefs::CheckIfUnused(const Stmt* s, const ID* id, bool report) {
    if ( id->IsGlobal() )
        return false;

    if ( auto& captures = ft->GetCaptures() ) {
        for ( auto& c : *captures )
            if ( c.Id() == id )
                return false;
    }

    auto uds = FindSuccUsage(s);
    if ( ! uds || ! uds->HasID(id) ) {
        if ( report && analysis_options.usage_issues > 0 && ! rc->IsTemporary(id) && ! rc->IsConstantVar(id) &&
             ! rc->IsNewLocal(id) && ! id->GetAttr(ATTR_IS_USED) )
            reporter->Warning("%s assignment unused: %s", id->Name(), obj_desc(s).c_str());

        return true;
    }

    return false;
}

UDs UseDefs::PropagateUDs(const Stmt* s, UDs succ_UDs, const Stmt* succ_stmt, bool second_pass) {
    if ( ! second_pass ) {
        om.AddObj(s);
        stmts.push_back(s);
    }

    switch ( s->Tag() ) {
        case STMT_LIST: {
            auto sl = s->AsStmtList();
            const auto& stmts = sl->Stmts();

            for ( int i = stmts.size(); --i >= 0; ) {
                auto s_i = stmts[i].get();

                const Stmt* succ;

                if ( i == int(stmts.size()) - 1 ) { // Very last statement.
                    succ = succ_stmt;
                    if ( successor2.find(s) != successor2.end() ) {
                        om.AddObj(s_i);
                        successor2[s_i] = successor2[s];
                    }
                }
                else
                    succ = stmts[i + 1].get();

                succ_UDs = PropagateUDs(s_i, succ_UDs, succ, second_pass);
            }

            return UseUDs(s, succ_UDs);
        }

        case STMT_CATCH_RETURN: {
            auto cr = s->AsCatchReturnStmt();
            auto block = cr->Block();

            auto uds = PropagateUDs(block.get(), succ_UDs, succ_stmt, second_pass);

            return UseUDs(s, uds);
        }

        case STMT_NULL:
        case STMT_NEXT:
        case STMT_BREAK:
        case STMT_FALLTHROUGH:
            // When we back up to one of these, its successor isn't
            // actually succ_stmt (other than for STMT_NULL).  However,
            // in the contexts in which these can occur, it doesn't
            // actually do any harm to use the successor anyway.
            return UseUDs(s, succ_UDs);

        case STMT_PRINT: return CreateExprUDs(s, s->AsPrintStmt()->ExprList(), succ_UDs);

        case STMT_EVENT:
        case STMT_CHECK_ANY_LEN:
        case STMT_ADD:
        case STMT_DELETE:
        case STMT_RETURN: {
            auto e = static_cast<const ExprStmt*>(s)->StmtExpr();

            if ( e )
                return CreateExprUDs(s, e, succ_UDs);
            else
                return UseUDs(s, succ_UDs);
        }

        case STMT_EXPR: {
            auto e = s->AsExprStmt()->StmtExpr();

            if ( e->Tag() != EXPR_ASSIGN )
                return CreateExprUDs(s, e, succ_UDs);

            // Change in use-defs as here we have a definition.
            auto a = e->AsAssignExpr();
            auto lhs_ref = a->GetOp1();

            if ( lhs_ref->Tag() != EXPR_REF )
                // Since we're working on reduced form ...
                reporter->InternalError("lhs inconsistency in UseDefs::ExprUDs");

            auto lhs_var = lhs_ref->GetOp1();
            auto lhs_id = lhs_var->AsNameExpr()->Id();
            auto lhs_UDs = RemoveID(lhs_id, succ_UDs);
            auto rhs_UDs = ExprUDs(a->GetOp2().get());
            auto uds = UD_Union(lhs_UDs, rhs_UDs);

            if ( ! second_pass )
                successor[s] = succ_stmt;

            return CreateUDs(s, uds);
        }

        case STMT_IF: {
            auto i = s->AsIfStmt();
            auto cond = i->StmtExpr();

            auto cond_UDs = ExprUDs(cond);
            auto true_UDs = PropagateUDs(i->TrueBranch(), succ_UDs, succ_stmt, second_pass);
            auto false_UDs = PropagateUDs(i->FalseBranch(), succ_UDs, succ_stmt, second_pass);

            return CreateUDs(s, UD_Union(cond_UDs, true_UDs, false_UDs));
        }

        case STMT_INIT:
            if ( ! second_pass )
                successor[s] = succ_stmt;

            return UseUDs(s, succ_UDs);

        case STMT_WHEN: {
            auto w = s->AsWhenStmt();
            auto wi = w->Info();
            auto uds = UD_Union(succ_UDs, ExprUDs(wi->Lambda().get()));

            auto timeout = wi->TimeoutExpr();
            if ( timeout )
                uds = UD_Union(uds, ExprUDs(timeout.get()));

            return CreateUDs(s, uds);
        }

        case STMT_SWITCH: {
            auto sw_UDs = std::make_shared<UseDefSet>();

            auto sw = s->AsSwitchStmt();
            auto cases = sw->Cases();

            for ( const auto& c : *cases ) {
                auto body = c->Body();
                auto uds = PropagateUDs(body, succ_UDs, succ_stmt, second_pass);

                auto exprs = c->ExprCases();
                if ( exprs ) {
                    auto e_UDs = ExprUDs(exprs);
                    uds = UD_Union(uds, e_UDs);
                }

                auto type_ids = c->TypeCases();
                if ( type_ids )
                    for ( const auto& id : *type_ids )
                        uds = RemoveID(id, uds);

                FoldInUDs(sw_UDs, uds);
            }

            auto e_UDs = ExprUDs(sw->StmtExpr());

            if ( sw->HasDefault() )
                FoldInUDs(sw_UDs, e_UDs);
            else
                // keep successor definitions in the mix
                FoldInUDs(sw_UDs, succ_UDs, e_UDs);

            return CreateUDs(s, sw_UDs);
        }

        case STMT_FOR: {
            auto f = s->AsForStmt();
            auto body = f->LoopBody();

            // The loop body has two potential successors, itself
            // and the successor of the entire "for" statement.
            om.AddObj(body);
            successor2[body] = succ_stmt;
            auto body_UDs = PropagateUDs(body, succ_UDs, body, second_pass);

            auto e = f->LoopExpr();
            auto f_UDs = ExprUDs(e);
            FoldInUDs(f_UDs, body_UDs);

            // Confluence: loop the top UDs back around to the bottom.
            auto bottom_UDs = UD_Union(f_UDs, succ_UDs);
            (void)PropagateUDs(body, bottom_UDs, body, true);

            auto ids = f->LoopVars();
            for ( const auto& id : *ids )
                RemoveUDFrom(f_UDs, id);

            auto val_var = f->ValueVar();
            if ( val_var )
                RemoveUDFrom(f_UDs, val_var.get());

            // The loop might not execute at all.
            FoldInUDs(f_UDs, succ_UDs);

            return CreateUDs(s, f_UDs);
        }

        case STMT_WHILE: {
            auto w = s->AsWhileStmt();
            auto body = w->Body();
            auto cond_stmt = w->CondPredStmt();

            // See note above for STMT_FOR regarding propagating
            // around the loop.
            auto succ = cond_stmt ? cond_stmt : body;
            om.AddObj(body.get());
            successor2[body.get()] = succ_stmt;
            auto body_UDs = PropagateUDs(body.get(), succ_UDs, succ.get(), second_pass);

            const auto& cond = w->Condition();
            auto w_UDs = UD_Union(ExprUDs(cond.get()), body_UDs);
            FoldInUDs(w_UDs, body_UDs);

            if ( cond_stmt ) {
                // Create a successor for the cond_stmt
                // that has the correct UDs associated with it.
                const auto& c_as_s = w->ConditionAsStmt();
                auto c_as_s_UDs = std::make_shared<UseDefSet>(w_UDs);
                CreateUDs(c_as_s.get(), c_as_s_UDs);

                w_UDs = PropagateUDs(cond_stmt, w_UDs, c_as_s, second_pass);
            }

            // Confluence: loop the top UDs back around to the bottom.
            auto bottom_UDs = UD_Union(w_UDs, succ_UDs);
            (void)PropagateUDs(body, bottom_UDs, succ, true);

            // The loop might not execute at all.
            FoldInUDs(w_UDs, succ_UDs);

            return CreateUDs(s, w_UDs);
        }

        default: reporter->InternalError("non-reduced statement in use-def analysis");
    }
}

UDs UseDefs::FindUsage(const Stmt* s) const {
    auto s_map = use_defs_map.find(s);

    if ( s_map == use_defs_map.end() )
        reporter->InternalError("missing use-defs");

    return s_map->second;
}

UDs UseDefs::FindSuccUsage(const Stmt* s) const {
    auto succ = successor.find(s);
    auto no_succ = (succ == successor.end() || ! succ->second);
    auto uds = no_succ ? nullptr : FindUsage(succ->second);

    auto succ2 = successor2.find(s);
    auto no_succ2 = (succ2 == successor2.end() || ! succ2->second);
    auto uds2 = no_succ2 ? nullptr : FindUsage(succ2->second);

    if ( uds && uds2 )
        return UD_Union(uds, uds2);
    else if ( uds )
        return uds;
    else
        return uds2;
}

UDs UseDefs::ExprUDs(const Expr* e) {
    auto uds = std::make_shared<UseDefSet>();

    switch ( e->Tag() ) {
        case EXPR_NAME:
        case EXPR_LAMBDA: AddInExprUDs(uds, e); break;

        case EXPR_FIELD_LHS_ASSIGN: {
            AddInExprUDs(uds, e->GetOp1().get());
            auto rhs_UDs = ExprUDs(e->GetOp2().get());
            uds = UD_Union(uds, rhs_UDs);
            break;
        }

        case EXPR_INCR:
        case EXPR_DECR: AddInExprUDs(uds, e->GetOp1()->AsRefExprPtr()->GetOp1().get()); break;

        case EXPR_ADD_TO:
        case EXPR_REMOVE_FROM: {
            AddInExprUDs(uds, e->GetOp1().get());
            auto rhs_UDs = ExprUDs(e->GetOp2().get());
            uds = UD_Union(uds, rhs_UDs);
            break;
        }

        case EXPR_RECORD_CONSTRUCTOR: {
            auto r = static_cast<const RecordConstructorExpr*>(e);
            AddInExprUDs(uds, r->Op().get());
            break;
        }

        case EXPR_TABLE_CONSTRUCTOR: {
            auto t = static_cast<const TableConstructorExpr*>(e);
            AddInExprUDs(uds, t->GetOp1().get());

            auto& t_attrs = t->GetAttrs();
            auto def_attr = t_attrs ? t_attrs->Find(ATTR_DEFAULT) : nullptr;
            auto& def_expr = def_attr ? def_attr->GetExpr() : nullptr;
            if ( def_expr && def_expr->Tag() == EXPR_LAMBDA )
                uds = ExprUDs(def_expr.get());

            break;
        }

        case EXPR_CONST: break;

        case EXPR_CALL: {
            auto c = e->AsCallExpr();
            AddInExprUDs(uds, c->Func());
            AddInExprUDs(uds, c->Args());
            break;
        }

        case EXPR_LIST: {
            auto l = e->AsListExpr();
            for ( const auto& l_e : l->Exprs() )
                AddInExprUDs(uds, l_e);

            break;
        }

        default:
            auto op1 = e->GetOp1();
            auto op2 = e->GetOp2();
            auto op3 = e->GetOp3();

            if ( ! op1 )
                reporter->InternalError("expression inconsistency in UseDefs::ExprUDs");

            AddInExprUDs(uds, op1.get());
            if ( op2 )
                AddInExprUDs(uds, op2.get());
            if ( op3 )
                AddInExprUDs(uds, op3.get());

            break;
    }

    return uds;
}

void UseDefs::AddInExprUDs(UDs uds, const Expr* e) {
    switch ( e->Tag() ) {
        case EXPR_REF: AddInExprUDs(uds, e->GetOp1().get()); break;

        case EXPR_NAME: AddID(uds, e->AsNameExpr()->Id()); break;

        case EXPR_LIST: {
            auto l = e->AsListExpr();
            for ( const auto& l_e : l->Exprs() )
                AddInExprUDs(uds, l_e);
        } break;

        case EXPR_EVENT: AddInExprUDs(uds, e->GetOp1().get()); break;

        case EXPR_INCR:
        case EXPR_DECR: AddInExprUDs(uds, e->GetOp1()->AsRefExprPtr()->GetOp1().get()); break;

        case EXPR_ASSIGN: // can occur inside a table constructor
        case EXPR_ADD_TO:
        case EXPR_REMOVE_FROM:
            AddInExprUDs(uds, e->GetOp1().get());
            AddInExprUDs(uds, e->GetOp2().get());
            break;

        case EXPR_FIELD_ASSIGN: AddInExprUDs(uds, static_cast<const FieldAssignExpr*>(e)->Op()); break;

        case EXPR_FIELD:
            // This happens for append-to-field.
            AddInExprUDs(uds, e->AsFieldExpr()->Op());
            break;

        case EXPR_LAMBDA: {
            auto outer_ids = e->AsLambdaExpr()->OuterIDs();
            for ( auto& i : outer_ids )
                AddID(uds, i);
            break;
        }

        case EXPR_CONST:
            // Nothing to do.
            break;

        default: reporter->InternalError("bad tag in UseDefs::AddInExprUDs"); break;
    }
}

void UseDefs::AddID(UDs uds, const ID* id) const { uds->Add(id); }

UDs UseDefs::RemoveID(const ID* id, const UDs& uds) {
    if ( ! uds )
        return nullptr;

    UDs new_uds = std::make_shared<UseDefSet>();

    new_uds->Replicate(uds);
    new_uds->Remove(id);

    return new_uds;
}

void UseDefs::RemoveUDFrom(UDs uds, const ID* id) {
    if ( uds )
        uds->Remove(id);
}

void UseDefs::FoldInUDs(UDs& main_UDs, const UDs& u1, const UDs& u2) {
    auto old_main = main_UDs;
    main_UDs = std::make_shared<UseDefSet>();

    if ( old_main )
        main_UDs->Replicate(old_main);

    if ( u1 )
        for ( auto ud : u1->IterateOver() )
            main_UDs->Add(ud);

    if ( u2 )
        for ( auto ud : u2->IterateOver() )
            main_UDs->Add(ud);
}

void UseDefs::UpdateUDs(const Stmt* s, const UDs& uds) {
    auto curr_uds = FindUsage(s);

    if ( ! curr_uds || UDs_are_copies.find(s) != UDs_are_copies.end() ) {
        // Copy-on-write.
        auto new_uds = std::make_shared<UseDefSet>();

        if ( curr_uds )
            new_uds->Replicate(curr_uds);

        CreateUDs(s, new_uds);

        curr_uds = new_uds;
    }

    if ( uds ) {
        for ( auto u : uds->IterateOver() )
            curr_uds->Add(u);
    }
}

UDs UseDefs::UD_Union(const UDs& u1, const UDs& u2, const UDs& u3) const {
    auto new_uds = std::make_shared<UseDefSet>();

    if ( u1 )
        new_uds->Replicate(u1);

    if ( u2 )
        for ( auto& u : u2->IterateOver() )
            AddID(new_uds, u);

    if ( u3 )
        for ( auto& u : u3->IterateOver() )
            AddID(new_uds, u);

    return new_uds;
}

UDs UseDefs::UseUDs(const Stmt* s, UDs uds) {
    om.AddObj(s);
    use_defs_map[s] = uds;
    UDs_are_copies.insert(s);
    return uds;
}

UDs UseDefs::CreateExprUDs(const Stmt* s, const Expr* e, const UDs& uds) {
    auto e_UDs = ExprUDs(e);
    auto new_UDs = UD_Union(uds, e_UDs);

    return CreateUDs(s, new_UDs);
}

UDs UseDefs::CreateUDs(const Stmt* s, UDs uds) {
    om.AddObj(s);
    use_defs_map[s] = uds;
    UDs_are_copies.erase(s);
    return uds;
}

} // namespace zeek::detail