/*

 * Copyright (c) 2022-2024 Huawei Device Co., Ltd.

 * 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.

 */



#include "ecmascript/compiler/async_function_lowering.h"



#include "ecmascript/js_generator_object.h"



namespace panda::ecmascript::kungfu {

void AsyncFunctionLowering::ProcessAll()

{

    ProcessJumpTable();



    if (IsLogEnabled()) {

        LOG_COMPILER(INFO) << "";

        LOG_COMPILER(INFO) << "\033[34m"

                           << "===================="

                           << " After async function lowering "

                           << "[" << GetMethodName() << "]"

                           << "===================="

                           << "\033[0m";

        circuit_->PrintAllGatesWithBytecode();

        LOG_COMPILER(INFO) << "\033[34m" << "========================= End ==========================" << "\033[0m";

    }

}



void AsyncFunctionLowering::ProcessJumpTable()

{

    GateRef newTarget = argAccessor_.GetCommonArgGate(CommonArgIdx::NEW_TARGET);

    GateRef isEqual = builder_.Equal(newTarget, builder_.Undefined());

    auto firstUse = accessor_.ConstUses(stateEntry_).begin();

    GateRef ifBranchCondition = builder_.Branch(stateEntry_, isEqual, 1, 1, "checkNewTarget");

    GateRef ifTrueCondition = builder_.IfTrue(ifBranchCondition);

    GateRef ifFalseCondition = builder_.IfFalse(ifBranchCondition);

    while (accessor_.GetOpCode(*firstUse) == OpCode::STATE_SPLIT) {

        firstUse++;

    }

    accessor_.ReplaceStateIn(*firstUse, ifTrueCondition);



    GateRef contextOffset = builder_.IntPtr(JSGeneratorObject::GENERATOR_CONTEXT_OFFSET);

    GateRef val = builder_.PtrAdd(newTarget, contextOffset);

    GateRef dependStart = builder_.DependRelay(ifFalseCondition, dependEntry_);

    auto bit = LoadStoreAccessor::ToValue(MemoryAttribute::Default());

    GateRef contextGate = circuit_->NewGate(circuit_->Load(bit), MachineType::I64, {dependStart, glue_, val},

                                            GateType::TaggedPointer());

    GateRef bcOffset = builder_.IntPtr(GeneratorContext::GENERATOR_BC_OFFSET_OFFSET);

    val = builder_.PtrAdd(contextGate, bcOffset);

    GateRef restoreOffsetGate = circuit_->NewGate(circuit_->Load(bit), MachineType::I32, {contextGate, glue_, val},

                                                  GateType::NJSValue());

    GateRef firstState = Circuit::NullGate();

    const auto &suspendAndResumeGates = bcBuilder_->GetAsyncRelatedGates();

    for (const auto &gate : suspendAndResumeGates) {

        EcmaOpcode ecmaOpcode = accessor_.GetByteCodeOpcode(gate);

        if (ecmaOpcode == EcmaOpcode::RESUMEGENERATOR) {

            RebuildGeneratorCfg(gate, restoreOffsetGate, ifFalseCondition, newTarget, firstState);

        }

    }

}



void AsyncFunctionLowering::RebuildGeneratorCfg(GateRef resumeGate, GateRef restoreOffsetGate, GateRef ifFalseCondition,

                                                GateRef newTarget, GateRef &firstState)

{

    GateRef stateGate = accessor_.GetState(resumeGate);

    GateRef suspendGate = stateGate;

    if (accessor_.GetOpCode(suspendGate) == OpCode::IF_SUCCESS) {

        suspendGate = accessor_.GetState(suspendGate);

    }

    GateRef offsetConstantGate = accessor_.GetValueIn(suspendGate);

    offsetConstantGate = builder_.TruncInt64ToInt32(offsetConstantGate);

    auto stateInGate = accessor_.GetState(resumeGate);

    bool flag = true;

    GateRef prevLoopBeginGate = Circuit::NullGate();

    GateRef loopBeginStateIn = Circuit::NullGate();

    GateRef prevBcOffsetPhiGate = Circuit::NullGate();

    while (true) {

        if (stateInGate == GetEntryBBStateOut()) {  // from state entry

            GateRef condition = builder_.Equal(offsetConstantGate, restoreOffsetGate);

            GateRef ifBranch = circuit_->NewGate(circuit_->IfBranch(0), { ifFalseCondition, condition });

            GateRef ifTrue = circuit_->NewGate(circuit_->IfTrue(), {ifBranch});

            GateRef ifFalse = circuit_->NewGate(circuit_->IfFalse(), {ifBranch});

            GateRef ifTrueDepend = builder_.DependRelay(ifTrue, restoreOffsetGate);

            GateRef ifFalseDepend = builder_.DependRelay(ifFalse, restoreOffsetGate);

            if (flag) {

                accessor_.ReplaceStateIn(resumeGate, ifTrue);

                accessor_.ReplaceValueIn(resumeGate, newTarget);

                accessor_.ReplaceDependIn(resumeGate, ifTrueDepend);

                circuit_->NewGate(circuit_->Return(), MachineType::NOVALUE,

                    { stateGate, suspendGate, suspendGate, circuit_->GetReturnRoot() },

                    GateType::AnyType());

            } else {

                loopBeginStateIn = ifTrue;

            }

            accessor_.ReplaceStateIn(ifBranch, ifFalseCondition);

            if (firstState != Circuit::NullGate()) {

                accessor_.ReplaceStateIn(firstState, ifFalse);

            } else {

                auto constant = builder_.UndefineConstant();

                circuit_->NewGate(circuit_->Return(), MachineType::NOVALUE,

                    { ifFalse, ifFalseDepend, constant, circuit_->GetReturnRoot() },

                    GateType::AnyType());

            }

            firstState = ifBranch;

        }

        auto opcode = accessor_.GetOpCode(stateInGate);

        if (opcode == OpCode::LOOP_BEGIN) {

            bool resumeInLoopBody = false;

            CheckResumeInLoopBody(stateInGate, resumeInLoopBody);

            if (resumeInLoopBody) {

                // This constant gate must be created by the NewGate method to distinguish whether the while

                // loop needs to modify the phi node or not.

                GateRef emptyOffsetGate = circuit_->NewGate(circuit_->GetMetaBuilder()->Constant(-1),

                                                            MachineType::I32, GateType::NJSValue());



                auto numIn = accessor_.GetNumIns(stateInGate);

                std::vector<GateRef> inList(numIn + 1, emptyOffsetGate);

                inList[0] = stateInGate; // 0 : state in

                inList[1] = restoreOffsetGate; // 1 : outloop value in

                GateRef bcOffsetPhiGate = circuit_->NewGate(circuit_->ValueSelector(numIn), MachineType::I32,

                                                            inList, GateType::NJSValue());



                GateRef condition = builder_.Equal(offsetConstantGate, bcOffsetPhiGate);

                GateRef ifBranch = circuit_->NewGate(circuit_->IfBranch(0), {stateInGate, condition});

                GateRef ifTrue = circuit_->NewGate(circuit_->IfTrue(), {ifBranch});

                GateRef ifFalse = circuit_->NewGate(circuit_->IfFalse(), {ifBranch});



                GateRef resumeStateGate = accessor_.GetState(resumeGate);

                if (accessor_.GetOpCode(resumeStateGate) != OpCode::IF_TRUE) {

                    accessor_.ReplaceStateIn(resumeGate, ifTrue);

                    accessor_.ReplaceValueIn(resumeGate, newTarget);

                    accessor_.ReplaceDependIn(resumeGate, GetDependPhiFromLoopBegin(stateInGate));

                    circuit_->NewGate(circuit_->Return(), MachineType::NOVALUE,

                        { stateGate, suspendGate, suspendGate, circuit_->GetReturnRoot() },

                        GateType::AnyType());

                } else {

                    // Handling multi-layer for loops

                    // When in a multi-layer loop, the value-selector node of the prev-loop

                    // should be used directly instead of generating a new node

                    UpdateValueSelector(prevLoopBeginGate, ifTrue, prevBcOffsetPhiGate, false);

                    accessor_.ReplaceValueIn(prevBcOffsetPhiGate, bcOffsetPhiGate);

                }

                accessor_.ReplaceStateIn(ifBranch, stateInGate);

                ModifyStateInput(stateInGate, ifBranch, ifFalse);



                prevLoopBeginGate = stateInGate;

                prevBcOffsetPhiGate = bcOffsetPhiGate;

                stateInGate = accessor_.GetState(stateInGate);

                flag = false;

                continue;

            }

        }

        if (loopBeginStateIn != Circuit::NullGate()) {

            UpdateValueSelector(prevLoopBeginGate, loopBeginStateIn, prevBcOffsetPhiGate);

            break;

        }

        if (stateInGate == GetEntryBBStateOut()) {

            break;

        }

        stateInGate = accessor_.GetState(stateInGate);

    }

}



void AsyncFunctionLowering::UpdateValueSelector(GateRef prevLoopBeginGate,

                                                GateRef controlStateGate,

                                                GateRef prevBcOffsetPhiGate,

                                                bool genNewValuePhiGate)

{

    GateRef loopBeginFirstState = accessor_.GetState(prevLoopBeginGate);

    // 2: statesIn

    GateRef newGate = circuit_->NewGate(circuit_->Merge(2),

                                        {controlStateGate, loopBeginFirstState});



    if (genNewValuePhiGate) {

        GateRef emptyOffsetGate =

            circuit_->NewGate(circuit_->GetMetaBuilder()->Constant(-1), // -1: distinguish bcoffset

            MachineType::I32, GateType::NJSValue());

        GateRef restoreOffset = accessor_.GetValueIn(prevBcOffsetPhiGate);

        // this value selector is compatible with await in the loop body

        GateRef valueSelector = circuit_->NewGate(circuit_->ValueSelector(2), MachineType::I32, // 2: num of valueIn

                                                  {newGate, restoreOffset, emptyOffsetGate},

                                                  GateType::NJSValue());

        accessor_.ReplaceValueIn(prevBcOffsetPhiGate, valueSelector);

    }

    accessor_.ReplaceStateIn(prevLoopBeginGate, newGate);

    auto loopBeginUses = accessor_.Uses(prevLoopBeginGate);

    for (auto use : loopBeginUses) {

        if (accessor_.GetOpCode(use) == OpCode::VALUE_SELECTOR && use != prevBcOffsetPhiGate) {

            auto machineType = accessor_.GetMachineType(use);

            auto gateType = accessor_.GetGateType(use);

            GateRef undefinedGate = Circuit::NullGate();

            if (gateType.IsNumberType()) {

                undefinedGate =

                circuit_->NewGate(circuit_->GetMetaBuilder()->Constant(JSTaggedValue::VALUE_ZERO),

                                  machineType, GateType::IntType());

            } else {

                undefinedGate =

                circuit_->NewGate(circuit_->GetMetaBuilder()->Constant(JSTaggedValue::VALUE_UNDEFINED),

                                  machineType, gateType);

            }

            auto firstValueGate = accessor_.GetValueIn(use, 0);

            auto newValueSelector = circuit_->NewGate(circuit_->ValueSelector(2), machineType, // 2: valuesIn

                                                      {newGate, undefinedGate, firstValueGate},

                                                      gateType);

            accessor_.ReplaceValueIn(use, newValueSelector);

        } else if (accessor_.GetOpCode(use) == OpCode::DEPEND_SELECTOR) {

            // if there is a dependSelector in the use node of the loop-begin, a new dependSelector node needs

            // to be generated. This node is bound to the merge node (newGate) before the loop-begin, and its

            // input corresponds to the 'dependEntry' (not the frist time enter the function) and

            // 'dependGate' (the first time enter the function) nodes.

            auto dependGate = accessor_.GetDep(use);

            auto newDependSelector = circuit_->NewGate(circuit_->DependSelector(2), // 2: num of dependIn

                                                       {newGate, circuit_->GetDependRoot(), dependGate});

            accessor_.ReplaceDependIn(use, newDependSelector);

        }

    }

}



bool AsyncFunctionLowering::IsAsyncRelated() const

{

    return bcBuilder_->GetAsyncRelatedGates().size() > 0;

}



void AsyncFunctionLowering::ModifyStateInput(GateRef stateInGate, GateRef ifBranch, GateRef ifFalse)

{

    // Find the node with LOOP_BEGIN as State input and modify its

    // state input to the newly created IF_FALSE node.

    auto uses = accessor_.Uses(stateInGate);

    for (auto useIt = uses.begin(); useIt != uses.end();) {

        GateRef use = *useIt;

        if (accessor_.IsState(use) && use != ifBranch) {

            useIt = accessor_.ReplaceIn(useIt, ifFalse);

        } else {

            useIt++;

        }

    }

}



void AsyncFunctionLowering::CheckResumeInLoopBody(GateRef stateInGate, bool &resumeInLoopBody)

{

    ASSERT(accessor_.GetOpCode(stateInGate) == OpCode::LOOP_BEGIN);

    ChunkQueue<GateRef> resumeList(circuit_->chunk());

    ChunkVector<VisitState> visited(circuit_->GetMaxGateId() + 1, VisitState::UNVISITED, circuit_->chunk());

    for (size_t i = 0; i < accessor_.GetNumIns(stateInGate); i++) {

        GateRef inGate = accessor_.GetIn(stateInGate, i);

        if (accessor_.GetOpCode(inGate) == OpCode::LOOP_BACK) {

            resumeList.push(inGate);

            visited[accessor_.GetId(inGate)] = VisitState::VISITED;

        }

    }

    auto loopBeginId = accessor_.GetId(stateInGate);

    visited[loopBeginId] = VisitState::VISITED;

    while (!resumeList.empty()) {

        GateRef curGate = resumeList.front();

        if (accessor_.GetOpCode(curGate) == OpCode::JS_BYTECODE &&

            accessor_.GetByteCodeOpcode(curGate) == EcmaOpcode::RESUMEGENERATOR) {

            resumeInLoopBody = true;

            break;

        }

        resumeList.pop();

        size_t stateStart = 0;

        size_t stateEnd = accessor_.GetStateCount(curGate);

        for (size_t idx = stateStart; idx < stateEnd; idx++) {

            GateRef gate = accessor_.GetState(curGate, idx);

            auto id = accessor_.GetId(gate);

            if (visited[id] == VisitState::UNVISITED) {

                visited[id] = VisitState::VISITED;

                resumeList.push(gate);

            }

        }

    }

}



GateRef AsyncFunctionLowering::GetDependPhiFromLoopBegin(GateRef gate) const

{

    auto loopBeginUses = accessor_.ConstUses(gate);

    for (auto use : loopBeginUses) {

        if (accessor_.GetOpCode(use) == OpCode::DEPEND_SELECTOR) {

            return use;

        }

    }

    LOG_COMPILER(FATAL) << "Can not find depend-selector from loopbegin";

    return Circuit::NullGate();

}



GateRef AsyncFunctionLowering::GetEntryBBStateOut() const

{

    auto& bb = bcBuilder_->GetBasicBlockById(0);   // 0 : Entry Block Id

    // state may CheckSafePointAndStackOver

    auto state = bb.dependCache;

    if (state == Circuit::NullGate()) {

        return circuit_->GetStateRoot();

    } else {

        return state;

    }

}



GateRef AsyncFunctionLowering::GetEntryBBDependOut() const

{

    auto& bb = bcBuilder_->GetBasicBlockById(0);   // 0 : Entry Block Id

    auto depend = bb.dependCache;

    if (depend == Circuit::NullGate()) {

        return circuit_->GetDependRoot();

    } else {

        return depend;

    }

}

}  // panda::ecmascript::kungfu