// Copyright 2014 Wouter van Oortmerssen. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

namespace lobster {

struct Optimizer {
    Parser &parser;
    SymbolTable &st;
    TypeChecker &tc;
    size_t total_changes = 0;
    vector<SubFunction *> sfstack;
    bool functions_removed = false;
    int runtime_checks;
    size_t always_inline = 16;
    size_t never_inline = 256;

    Optimizer(Parser &_p, SymbolTable &_st, TypeChecker &_tc, int runtime_checks)
        : parser(_p), st(_st), tc(_tc), runtime_checks(runtime_checks) {
        if (runtime_checks >= RUNTIME_DEBUG) {
            // User wants to see useful stack-traces, only inline the tiniest of functions.
            always_inline = 2;
            never_inline = 4;
        }
        // We don't optimize parser.root, it only contains a single call.
        for (auto f : parser.st.functiontable) {
            again:
            for (auto ov : f->overloads) {
                auto sf = ov->sf;
                if (sf && sf->typechecked) {
                    for (; sf; sf = sf->next) {
                        functions_removed = false;
                        OptimizeFunction(*sf);
                        if (functions_removed) goto again;
                    }
                }
            }
        }
        LOG_INFO("optimizer: ", total_changes, " optimizations");
    }

    void OptimizeFunction(SubFunction &sf) {
        if (sf.optimized) return;
        sf.optimized = true;
        if (!sf.sbody) return;
        if (!sf.typechecked) {
            delete sf.sbody;
            sf.sbody = nullptr;
            return;
        }
        sfstack.push_back(&sf);
        auto nb = sf.sbody->Optimize(*this);
        sf.node_count = sf.sbody->Count();
        assert(nb == sf.sbody);
        (void)nb;
        sfstack.pop_back();
    }

    void Changed() { total_changes++; }

    Node *Typed(TypeRef type, Lifetime lt, Node *n) {
        n->exptype = type;
        n->lt = lt;
        return n;
    }
};

Node *Node::Optimize(Optimizer &opt) {
    assert(exptype->t != V_UNDEFINED);
    for (size_t i = 0; i < Arity(); i++) {
        Children()[i] = Children()[i]->Optimize(opt);
    }
    Value cval = NilVal();
    auto t = ConstVal(&opt.tc, cval);
    if (t == V_VOID) return this;
    Node *r;
    switch (t) {
        case V_INT:   r = new IntConstant(line, cval.ival()); break;
        case V_FLOAT: r = new FloatConstant(line, cval.fval()); break;
        case V_NIL:   r = new Nil(line, { exptype }); break;
        default:      assert(false); return this;
    }
    r = opt.Typed(exptype, LT_ANY, r);
    delete this;
    opt.Changed();
    return r->Optimize(opt);
}

Node *VectorConstructor::Optimize(Optimizer &opt) {
    return Node::Optimize(opt);
}

Node *ObjectConstructor::Optimize(Optimizer &opt) {
    return Node::Optimize(opt);
}

Node *Nil::Optimize(Optimizer &) {
    return this;
}

Node *IntConstant::Optimize(Optimizer &) {
    return this;
}

Node *FloatConstant::Optimize(Optimizer &) {
    return this;
}

Node *Call::Optimize(Optimizer &opt) {
    Node::Optimize(opt);
    assert(sf->numcallers > 0);
    auto parent = opt.sfstack.back();
    // FIXME: Reduce these requirements where possible.
    bool is_inlinable =
        !sf->isrecursivelycalled &&
        sf->num_returns <= 1 &&
        // Have to double-check this, since last return may be a return from.
        AssertIs<Return>(sf->sbody->children.back())->sf == sf &&
        // This may happen even if num_returns==1 when body of sf is a non-local return also,
        // See e.g. exception_handler
        // FIXME: if the function that caused this to be !=0 gets inlined, this needs to be
        // decremented so it can be inlined after all.
        sf->num_returns_non_local == 0 &&
        vtable_idx < 0 &&
        sf->returntype->NumValues() <= 1 &&
        // Because we inline so aggressively, it is possible to generate huuge functions,
        // which may cause a problem for our stack, or that of e.g. V8 in Wasm.
        parent->locals.size() < 1024;
    // Attempt to optimize function we're calling first, that way if it shrinks (or grows) it's
    // more or less likely to be inlined.
    if (is_inlinable) opt.OptimizeFunction(*sf);
    // Check if we should inline this call.
    if (!is_inlinable ||
        // Inline small functions even if called multiple times.
        (sf->numcallers > 1 && sf->node_count >= opt.always_inline) ||
        // Don't inline really gigantic functions, this helps with not flattening the call-graph too
        // much for stack traces, profiling and such, and may also make them easier to reg-alloc etc.
        (sf->numcallers <= 1 && sf->node_count >= opt.never_inline) ||
        // Don't inline functions that are being profiled.
        (LOBSTER_FRAME_PROFILER && sf->attributes.find("profile") != sf->attributes.end())) {
        return this;
    }
    auto AddToLocals = [&](const vector<Arg> &av) {
        for (auto &arg : av) {
            // We have to check if the sid already exists, since inlining the same function
            // multiple times in the same parent can cause this. This variable is shared
            // between the copies in the parent, second use overwrites the first etc.
            // We generally have to keep using this sid rather than creating a new one, since
            // this function may call others that may refer to this sid, etc.
            for (auto &loc : parent->locals) if (loc.sid == arg.sid) goto already;
            parent->locals.push_back(arg);
            arg.sid->sf_def = parent;
            already:;
        }
    };
    AddToLocals(sf->args);
    AddToLocals(sf->locals);
    int ai = 0;
    auto list = new Block(line);
    for (auto c : children) {
        auto &arg = sf->args[ai];
        // NOTE: this introduces locals which potentially borrow, which the typechecker so far
        // never introduces.
        // We can't just force these to LT_KEEP, since the sids maybe shared with non-inlined
        // instances (e.g. in dynamic dispatch).
        // Borrowing could be problematic if 2 copies of the same function get inlined, since
        // that creates an overwite of a borrowed variable, but the codegen ensures the overwrite
        // does not decref.
        auto def = new Define(line, c);
        def->tsids.push_back({ arg.sid, { arg.spec_type } });
        list->Add(opt.Typed(type_void, LT_ANY, def));
        ai++;
    }
    // TODO: triple-check this similar in semantics to what happens in CloneFunction() in the
    // typechecker.
    if (sf->numcallers == 1) {
        list->children.append(sf->sbody->children.data(), sf->sbody->children.size());
        sf->sbody->children.clear();
        delete sf->sbody;
        sf->sbody = nullptr;
        opt.functions_removed = sf->parent->RemoveSubFunction(sf);
        assert(opt.functions_removed);
    } else {
        for (auto c : sf->sbody->children) {
            auto nc = c->Clone(false);
            list->children.push_back(nc);
            nc->Iterate([](Node *i) {
                if (auto call = Is<Call>(i)) call->sf->numcallers++;
                if (auto dcall = Is<DynCall>(i)) dcall->sf->numcallers++;
            });
        }
    }
    sf->numcallers--;
    // Remove single return statement pointing to function that is now gone.
    auto ret = AssertIs<Return>(list->children.back());
    if (ret->sf == sf) {
        assert(ret->child->exptype->NumValues() <= 1);
        assert(sf->num_returns <= 1);
        assert(sf->num_returns_non_local == 0);
        // This is not great: having to undo the optimization in Return::TypeCheck where this
        // flag was set.
        // Since the caller generally expects to keep the return value of the now inlined
        // function, and since the variables of this function are not dead yet (they can be called
        // multiple times if inside a loop body or inlined multiple times, which may cause the
        // var to be decreffed when overwritten on second use), we have to incref.
        // TODO: investigate if setting them to null at scope exit would be an alternative?
        if (sf->consumes_vars_on_return) {
            AssertIs<IdentRef>(ret->child);
            opt.tc.MakeLifetime(ret->child, LT_KEEP, 1, 0);
        }
        list->children.back() = ret->child;
        ret->child = nullptr;
        delete ret;
    }
    auto r = opt.Typed(exptype, sf->ltret, list);
    children.clear();
    r = r->Optimize(opt);
    delete this;  // Do this after, since Optimize may touch this same call.
    opt.Changed();
    return r;
}

}  // namespace lobster


