/// Emit code to release/destroy temporaries.
void PartialApplyCombiner::releaseTemporaries() {
    // Insert releases and destroy_addrs as early as possible,
    // because we don't want to keep objects alive longer than
    // its really needed.
    for (auto Op : Tmps) {
        auto TmpType = Op->getType().getObjectType();
        if (TmpType.isTrivial(PAI->getModule()))
            continue;
        for (auto *EndPoint : PAFrontier) {
            Builder.setInsertionPoint(EndPoint);
            if (!TmpType.isAddressOnly(PAI->getModule())) {
                auto *Load = Builder.createLoad(PAI->getLoc(), Op);
                Builder.createReleaseValue(PAI->getLoc(), Load, Atomicity::Atomic);
            } else {
                Builder.createDestroyAddr(PAI->getLoc(), Op);
            }
        }
    }
}
/// Process an apply instruction which uses a partial_apply
/// as its callee.
/// Returns true on success.
bool PartialApplyCombiner::processSingleApply(FullApplySite AI) {
    Builder.setInsertionPoint(AI.getInstruction());
    Builder.setCurrentDebugScope(AI.getDebugScope());

    // Prepare the args.
    SmallVector<SILValue, 8> Args;
    // First the ApplyInst args.
    for (auto Op : AI.getArguments())
        Args.push_back(Op);

    SILInstruction *InsertionPoint = &*Builder.getInsertionPoint();
    // Next, the partial apply args.

    // Pre-process partial_apply arguments only once, lazily.
    if (isFirstTime) {
        isFirstTime = false;
        if (!allocateTemporaries())
            return false;
    }

    // Now, copy over the partial apply args.
    for (auto Op : PAI->getArguments()) {
        auto Arg = Op;
        // If there is new temporary for this argument, use it instead.
        if (isa<AllocStackInst>(Arg)) {
            if (ArgToTmp.count(Arg)) {
                Op = ArgToTmp.lookup(Arg);
            }
        }
        Args.push_back(Op);
    }

    Builder.setInsertionPoint(InsertionPoint);
    Builder.setCurrentDebugScope(AI.getDebugScope());

    // The thunk that implements the partial apply calls the closure function
    // that expects all arguments to be consumed by the function. However, the
    // captured arguments are not arguments of *this* apply, so they are not
    // pre-incremented. When we combine the partial_apply and this apply into
    // a new apply we need to retain all of the closure non-address type
    // arguments.
    auto ParamInfo = PAI->getSubstCalleeType()->getParameters();
    auto PartialApplyArgs = PAI->getArguments();
    // Set of arguments that need to be released after each invocation.
    SmallVector<SILValue, 8> ToBeReleasedArgs;
    for (unsigned i = 0, e = PartialApplyArgs.size(); i < e; ++i) {
        SILValue Arg = PartialApplyArgs[i];
        if (!Arg->getType().isAddress()) {
            // Retain the argument as the callee may consume it.
            Builder.emitRetainValueOperation(PAI->getLoc(), Arg);
            // For non consumed parameters (e.g. guaranteed), we also need to
            // insert releases after each apply instruction that we create.
            if (!ParamInfo[ParamInfo.size() - PartialApplyArgs.size() + i].
                    isConsumed())
                ToBeReleasedArgs.push_back(Arg);
        }
    }

    auto *F = FRI->getReferencedFunction();
    SILType FnType = F->getLoweredType();
    SILType ResultTy = F->getLoweredFunctionType()->getSILResult();
    ArrayRef<Substitution> Subs = PAI->getSubstitutions();
    if (!Subs.empty()) {
        FnType = FnType.substGenericArgs(PAI->getModule(), Subs);
        ResultTy = FnType.getAs<SILFunctionType>()->getSILResult();
    }

    FullApplySite NAI;
    if (auto *TAI = dyn_cast<TryApplyInst>(AI))
        NAI =
            Builder.createTryApply(AI.getLoc(), FRI, FnType, Subs, Args,
                                   TAI->getNormalBB(), TAI->getErrorBB());
    else
        NAI =
            Builder.createApply(AI.getLoc(), FRI, FnType, ResultTy, Subs, Args,
                                cast<ApplyInst>(AI)->isNonThrowing());

    // We also need to release the partial_apply instruction itself because it
    // is consumed by the apply_instruction.
    if (auto *TAI = dyn_cast<TryApplyInst>(AI)) {
        Builder.setInsertionPoint(TAI->getNormalBB()->begin());
        for (auto Arg : ToBeReleasedArgs) {
            Builder.emitReleaseValueOperation(PAI->getLoc(), Arg);
        }
        Builder.createStrongRelease(AI.getLoc(), PAI, Atomicity::Atomic);
        Builder.setInsertionPoint(TAI->getErrorBB()->begin());
        // Release the non-consumed parameters.
        for (auto Arg : ToBeReleasedArgs) {
            Builder.emitReleaseValueOperation(PAI->getLoc(), Arg);
        }
        Builder.createStrongRelease(AI.getLoc(), PAI, Atomicity::Atomic);
        Builder.setInsertionPoint(AI.getInstruction());
    } else {
        // Release the non-consumed parameters.
        for (auto Arg : ToBeReleasedArgs) {
            Builder.emitReleaseValueOperation(PAI->getLoc(), Arg);
        }
        Builder.createStrongRelease(AI.getLoc(), PAI, Atomicity::Atomic);
    }

    SilCombiner->replaceInstUsesWith(*AI.getInstruction(), NAI.getInstruction());
    SilCombiner->eraseInstFromFunction(*AI.getInstruction());
    return true;
}
/// Returns true on success.
bool PartialApplyCombiner::allocateTemporaries() {
    // Copy the original arguments of the partial_apply into
    // newly created temporaries and use these temporaries instead of
    // the original arguments afterwards.
    // This is done to "extend" the life-time of original partial_apply
    // arguments, as they may be destroyed/deallocated before the last
    // use by one of the apply instructions.
    // TODO:
    // Copy arguments of the partial_apply into new temporaries
    // only if the lifetime of arguments ends before their uses
    // by apply instructions.
    bool needsReleases = false;
    CanSILFunctionType PAITy =
        dyn_cast<SILFunctionType>(PAI->getCallee()->getType().getSwiftType());

    // Emit a destroy value for each captured closure argument.
    ArrayRef<SILParameterInfo> Params = PAITy->getParameters();
    auto Args = PAI->getArguments();
    unsigned Delta = Params.size() - Args.size();

    llvm::SmallVector<std::pair<SILValue, unsigned>, 8> ArgsToHandle;
    for (unsigned AI = 0, AE = Args.size(); AI != AE; ++AI) {
        SILValue Arg = Args[AI];
        SILParameterInfo Param = Params[AI + Delta];
        if (Param.isIndirectMutating())
            continue;
        // Create a temporary and copy the argument into it, if:
        // - the argument stems from an alloc_stack
        // - the argument is consumed by the callee and is indirect
        //   (e.g. it is an @in argument)
        if (isa<AllocStackInst>(Arg) ||
                (Param.isConsumed() && Param.isIndirect())) {
            // If the temporary is non-trivial, we need to release it later.
            if (!Arg->getType().isTrivial(PAI->getModule()))
                needsReleases = true;
            ArgsToHandle.push_back(std::make_pair(Arg, AI));
        }
    }

    if (needsReleases) {
        // Compute the set of endpoints, which will be used to insert releases of
        // temporaries. This may fail if the frontier is located on a critical edge
        // which we may not split (no CFG changes in SILCombine).
        ValueLifetimeAnalysis VLA(PAI);
        if (!VLA.computeFrontier(PAFrontier, ValueLifetimeAnalysis::DontModifyCFG))
            return false;
    }

    for (auto ArgWithIdx : ArgsToHandle) {
        SILValue Arg = ArgWithIdx.first;
        Builder.setInsertionPoint(PAI->getFunction()->begin()->begin());
        // Create a new temporary at the beginning of a function.
        auto *Tmp = Builder.createAllocStack(PAI->getLoc(), Arg->getType(),
        {/*Constant*/ true, ArgWithIdx.second});
        Builder.setInsertionPoint(PAI);
        // Copy argument into this temporary.
        Builder.createCopyAddr(PAI->getLoc(), Arg, Tmp,
                               IsTake_t::IsNotTake,
                               IsInitialization_t::IsInitialization);

        Tmps.push_back(Tmp);
        ArgToTmp.insert(std::make_pair(Arg, Tmp));
    }
    return true;
}