//===--- ClosureLifetimeFixup.cpp - Fixup the lifetime of closures --------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "closure-lifetime-fixup"

#include "swift/SIL/DebugUtils.h"
#include "swift/SIL/InstructionUtils.h"
#include "swift/SIL/SILArgument.h"
#include "swift/SIL/SILBuilder.h"
#include "swift/SIL/SILInstruction.h"
#include "swift/SILOptimizer/PassManager/Passes.h"
#include "swift/SILOptimizer/PassManager/Transforms.h"
#include "swift/SILOptimizer/Utils/CFG.h"
#include "swift/SILOptimizer/Utils/Local.h"

#include "llvm/Support/CommandLine.h"

llvm::cl::opt<bool> DisableConvertEscapeToNoEscapeSwitchEnumPeephole(
    "sil-disable-convert-escape-to-noescape-switch-peephole",
    llvm::cl::init(false),
    llvm::cl::desc(
        "Disable the convert_escape_to_noescape switch enum peephole. "),
    llvm::cl::Hidden);

using namespace swift;

static SILBasicBlock *getOptionalDiamondSuccessor(SwitchEnumInst *sei) {
  auto numSuccs = sei->getNumSuccessors();
  if (numSuccs != 2)
    return nullptr;
  auto *succSome = sei->getCase(0).second;
  auto *succNone = sei->getCase(1).second;
  if (succSome->args_size() != 1)
    std::swap(succSome, succNone);

  if (succSome->args_size() != 1 || succNone->args_size() != 0)
    return nullptr;

  auto *succ = succSome->getSingleSuccessorBlock();
  if (!succ)
    return nullptr;

  if (succNone == succ)
    return succ;

  succNone = succNone->getSingleSuccessorBlock();
  if (succNone == succ)
    return succ;

  if (succNone == nullptr)
    return nullptr;

  succNone = succNone->getSingleSuccessorBlock();
  if (succNone == succ)
    return succ;

  return nullptr;
}

/// Find a safe insertion point for closure destruction. We might create a
/// closure that captures self in deinit of self. In this situation it is not
/// safe to destroy the closure after we called super deinit. We have to place
/// the closure destruction before that call.
///
///  %deinit = objc_super_method %0 : $C, #A.deinit!deallocator.foreign
///  %super = upcast %0 : $C to $A
///  apply %deinit(%super) : $@convention(objc_method) (A) -> ()
///  end_lifetime %super : $A
static SILInstruction *getDeinitSafeClosureDestructionPoint(TermInst *Term) {
  for (auto It = Term->getParent()->rbegin(), E = Term->getParent()->rend();
       It != E; ++It) {
    if (auto *EndLifetime = dyn_cast<EndLifetimeInst>(&*It)) {
      auto *SuperInstance = EndLifetime->getOperand()->getDefiningInstruction();
      assert(SuperInstance && "Expected an instruction");
      return SuperInstance;
    }
  }
  return Term;
}

/// Extend the lifetime of the convert_escape_to_noescape's operand to the end
/// of the function.
static void extendLifetimeToEndOfFunction(SILFunction &Fn,
                                          ConvertEscapeToNoEscapeInst *Cvt) {
  auto EscapingClosure = Cvt->getOperand();
  auto EscapingClosureTy = EscapingClosure->getType();
  auto OptionalEscapingClosureTy = SILType::getOptionalType(EscapingClosureTy);
  auto loc = RegularLocation::getAutoGeneratedLocation();

  SILBuilderWithScope B(Cvt);
  auto NewCvt = B.createConvertEscapeToNoEscape(
      Cvt->getLoc(), Cvt->getOperand(), Cvt->getType(), true);
  Cvt->replaceAllUsesWith(NewCvt);
  Cvt->eraseFromParent();
  Cvt = NewCvt;

  // Create an alloc_stack Optional<() -> ()> at the beginning of the function.
  AllocStackInst *Slot;
  auto &Context = Cvt->getModule().getASTContext();
  {
    SILBuilderWithScope B(Fn.getEntryBlock()->begin());
    Slot = B.createAllocStack(loc, OptionalEscapingClosureTy);
    auto *NoneDecl = Context.getOptionalNoneDecl();
    // Store None to it.
    B.createStore(
        loc, B.createEnum(loc, SILValue(), NoneDecl, OptionalEscapingClosureTy),
        Slot, StoreOwnershipQualifier::Init);
  }
  // Insert a copy before the convert_escape_to_noescape and store it to the
  // alloc_stack location.
  {
    SILBuilderWithScope B(Cvt);
    auto *SomeDecl = Context.getOptionalSomeDecl();
    B.createDestroyAddr(loc, Slot);
    auto ClosureCopy = B.createCopyValue(loc, EscapingClosure);
    B.createStore(
        loc,
        B.createEnum(loc, ClosureCopy, SomeDecl, OptionalEscapingClosureTy),
        Slot, StoreOwnershipQualifier::Init);
  }
  // Insert destroys at the function exits.
  SmallVector<SILBasicBlock *, 4> ExitingBlocks;
  Fn.findExitingBlocks(ExitingBlocks);
  for (auto *Exit : ExitingBlocks) {
    auto *Term = Exit->getTerminator();
    auto *SafeClosureDestructionPt = getDeinitSafeClosureDestructionPoint(Term);
    SILBuilderWithScope B(SafeClosureDestructionPt);
    B.createDestroyAddr(loc, Slot);
    SILBuilderWithScope B2(Term);
    B2.createDeallocStack(loc, Slot);
  }
}

static SILInstruction *lookThroughRebastractionUsers(
    SILInstruction *Inst,
    llvm::DenseMap<SILInstruction *, SILInstruction *> &Memoized) {

  if (Inst == nullptr)
    return nullptr;

  // Try a cached lookup.
  auto Res = Memoized.find(Inst);
  if (Res != Memoized.end())
    return Res->second;

  // Cache recursive results.
  auto memoizeResult = [&] (SILInstruction *from, SILInstruction *toResult) {
    Memoized[from] = toResult;
    return toResult;
  };

  // If we have a convert_function, just look at its user.
  if (auto *Cvt = dyn_cast<ConvertFunctionInst>(Inst))
    return memoizeResult(Inst, lookThroughRebastractionUsers(
                                   getSingleNonDebugUser(Cvt), Memoized));
  if (auto *Cvt = dyn_cast<ConvertEscapeToNoEscapeInst>(Inst))
    return memoizeResult(Inst, lookThroughRebastractionUsers(
                                   getSingleNonDebugUser(Cvt), Memoized));

  // If we have a partial_apply user look at its single (non release) user.
  auto *PA = dyn_cast<PartialApplyInst>(Inst);
  if (!PA) return Inst;

  SILInstruction *SingleNonDebugNonRefCountUser = nullptr;
  for (auto *Usr : getNonDebugUses(PA)) {
    auto *I = Usr->getUser();
    if (onlyAffectsRefCount(I))
      continue;
    if (SingleNonDebugNonRefCountUser) {
      SingleNonDebugNonRefCountUser = nullptr;
      break;
    }
    SingleNonDebugNonRefCountUser = I;
  }

  return memoizeResult(Inst, lookThroughRebastractionUsers(
                                 SingleNonDebugNonRefCountUser, Memoized));
}

static bool tryExtendLifetimeToLastUse(
    ConvertEscapeToNoEscapeInst *Cvt,
    llvm::DenseMap<SILInstruction *, SILInstruction *> &Memoized) {
  // If there is a single user that is an apply this is simple: extend the
  // lifetime of the operand until after the apply.
  auto SingleUser = lookThroughRebastractionUsers(Cvt, Memoized);
  if (!SingleUser)
    return false;

  // Handle an apply.
  if (auto SingleApplyUser = FullApplySite::isa(SingleUser)) {
    // FIXME: Don't know how-to handle begin_apply/end_apply yet.
    if (isa<BeginApplyInst>(SingleApplyUser.getInstruction())) {
      return false;
    }

    auto loc = RegularLocation::getAutoGeneratedLocation();

    // Insert a copy at the convert_escape_to_noescape [not_guaranteed] and
    // change the instruction to the guaranteed form.
    auto EscapingClosure = Cvt->getOperand();
    {
      SILBuilderWithScope B(Cvt);
      auto NewCvt = B.createConvertEscapeToNoEscape(
          Cvt->getLoc(), Cvt->getOperand(), Cvt->getType(), true);
      Cvt->replaceAllUsesWith(NewCvt);
      Cvt->eraseFromParent();
      Cvt = NewCvt;
    }

    SILBuilderWithScope B2(Cvt);
    auto ClosureCopy = B2.createCopyValue(loc, EscapingClosure);

    // Insert a destroy after the apply.
    if (auto *Apply = dyn_cast<ApplyInst>(SingleApplyUser.getInstruction())) {
      auto InsertPt = std::next(SILBasicBlock::iterator(Apply));
      SILBuilderWithScope B3(InsertPt);
      B3.createDestroyValue(loc, ClosureCopy);

    } else if (auto *Try =
                   dyn_cast<TryApplyInst>(SingleApplyUser.getInstruction())) {
      for (auto *SuccBB : Try->getSuccessorBlocks()) {
        SILBuilderWithScope B3(SuccBB->begin());
        B3.createDestroyValue(loc, ClosureCopy);
      }
    } else {
      llvm_unreachable("Unknown FullApplySite instruction kind");
    }
    return true;
  }
  return false;
}

/// Ensure the lifetime of the closure accross an
///
///   optional<@escaping () -> ()> to
///   optional<@noescape @convention(block) () -> ()>
///
///   conversion and its use.
///
///   The pattern this is looking for
///                            switch_enum %closure
///                           /           \
///     convert_escape_to_noescape          nil
///                             switch_enum
///                           /           \
///                convertToBlock          nil
///                           \            /
///                     (%convertOptionalBlock :)
///   We will insert a copy_value of the original %closure before the two
///   diamonds. And a destroy of %closure at the last destroy of
///   %convertOptionalBlock.
static bool trySwitchEnumPeephole(ConvertEscapeToNoEscapeInst *Cvt) {
  auto *blockArg = dyn_cast<SILArgument>(Cvt->getOperand());
  if (!blockArg)
    return false;
  auto *PredBB = Cvt->getParent()->getSinglePredecessorBlock();
  if (!PredBB)
    return false;
  auto *ConvertSuccessorBlock = Cvt->getParent()->getSingleSuccessorBlock();
  if (!ConvertSuccessorBlock)
    return false;
  auto *SwitchEnum1 = dyn_cast<SwitchEnumInst>(PredBB->getTerminator());
  if (!SwitchEnum1)
    return false;
  auto *DiamondSucc = getOptionalDiamondSuccessor(SwitchEnum1);
  if (!DiamondSucc)
    return false;
  auto *SwitchEnum2 = dyn_cast<SwitchEnumInst>(DiamondSucc->getTerminator());
  if (!SwitchEnum2)
    return false;
  auto *DiamondSucc2 = getOptionalDiamondSuccessor(SwitchEnum2);
  if (!DiamondSucc2)
    return false;
  if (DiamondSucc2->getNumArguments() != 1)
    return false;

  // Look for the last and only destroy.
  SILInstruction *onlyDestroy = [&]() -> SILInstruction * {
    SILInstruction *lastDestroy = nullptr;
    for (auto *Use : DiamondSucc2->getArgument(0)->getUses()) {
      SILInstruction *Usr = Use->getUser();
      if (isa<ReleaseValueInst>(Usr) || isa<StrongReleaseInst>(Usr) ||
          isa<DestroyValueInst>(Usr)) {
        if (lastDestroy)
          return nullptr;
        lastDestroy = Usr;
      }
    }
    return lastDestroy;
  }();
  if (!onlyDestroy)
    return false;

  // Replace the convert_escape_to_noescape instruction.
  {
    SILBuilderWithScope B(Cvt);
    auto NewCvt = B.createConvertEscapeToNoEscape(
        Cvt->getLoc(), Cvt->getOperand(), Cvt->getType(), true);
    Cvt->replaceAllUsesWith(NewCvt);
    Cvt->eraseFromParent();
  }

  // Extend the lifetime.
  SILBuilderWithScope B(SwitchEnum1);
  auto loc = RegularLocation::getAutoGeneratedLocation();
  auto copy =
      B.createCopyValue(loc, SwitchEnum1->getOperand());
  B.setInsertionPoint(onlyDestroy);
  B.createDestroyValue(loc, copy);
  return true;
}

/// Look for a single destroy user and possibly unowned apply uses.
static SILInstruction *getOnlyDestroy(CopyBlockWithoutEscapingInst *CB) {
  SILInstruction *onlyDestroy = nullptr;

  for (auto *Use : getNonDebugUses(CB)) {
    SILInstruction *Inst = Use->getUser();

    // If this an apply use, only handle unowned parameters.
    if (auto Apply = FullApplySite::isa(Inst)) {
      SILArgumentConvention Conv = Apply.getArgumentConvention(*Use);
      if (Conv != SILArgumentConvention::Direct_Unowned)
        return nullptr;
      continue;
    }

    // We have already seen one destroy.
    if (onlyDestroy)
      return nullptr;

    if (isa<DestroyValueInst>(Inst) || isa<ReleaseValueInst>(Inst) ||
        isa<StrongReleaseInst>(Inst)) {
      onlyDestroy = Inst;
      continue;
    }

    // Some other instruction.
    return nullptr;
  }

  if (!onlyDestroy)
    return nullptr;

  // Now look at whether the dealloc_stack or the destroy postdominates and
  // return the post dominator.
  auto *BlockInit = dyn_cast<InitBlockStorageHeaderInst>(CB->getBlock());
  if (!BlockInit)
    return nullptr;

  auto *AS = dyn_cast<AllocStackInst>(BlockInit->getBlockStorage());
  if (!AS)
    return nullptr;
  auto *Dealloc = AS->getSingleDeallocStack();
  if (!Dealloc || Dealloc->getParent() != onlyDestroy->getParent())
    return nullptr;

  // Return the later instruction.
  for (auto It = SILBasicBlock::iterator(onlyDestroy),
            E = Dealloc->getParent()->end();
       It != E; ++It) {
    if (&*It == Dealloc)
      return Dealloc;
  }
  return onlyDestroy;
}

/// Lower a copy_block_without_escaping instruction.
///
///    This involves replacing:
///
///      %copy = copy_block_without_escaping %block withoutEscaping %closure
///
///      ...
///      destroy_value %copy
///
///    by (roughly) the instruction sequence:
///
///      %copy = copy_block %block
///
///      ...
///      destroy_value %copy
///      %e = is_escaping %closure
///      cond_fail %e
///      destroy_value %closure
static bool fixupCopyBlockWithoutEscaping(CopyBlockWithoutEscapingInst *CB) {
  SmallVector<SILInstruction *, 4> LifetimeEndPoints;

  // Find the end of the lifetime of the copy_block_without_escaping
  // instruction.
  auto &Fn  = *CB->getFunction();
  auto *SingleDestroy = getOnlyDestroy(CB);
  SmallVector<SILBasicBlock *, 4> ExitingBlocks;
  Fn.findExitingBlocks(ExitingBlocks);
  if (SingleDestroy) {
    LifetimeEndPoints.push_back(
        &*std::next(SILBasicBlock::iterator(SingleDestroy)));
  } else {
    // Otherwise, conservatively insert verification at the end of the function.
    for (auto *Exit : ExitingBlocks)
      LifetimeEndPoints.push_back(Exit->getTerminator());
  }

  auto SentinelClosure = CB->getClosure();
  auto Loc = CB->getLoc();

  SILBuilderWithScope B(CB);
  auto *NewCB = B.createCopyBlock(Loc, CB->getBlock());
  CB->replaceAllUsesWith(NewCB);
  CB->eraseFromParent();

  // Create an stack slot for the closure sentinel and store the sentinel (or
  // none on other paths.
  auto generatedLoc = RegularLocation::getAutoGeneratedLocation();
  AllocStackInst *Slot;
  auto &Context = NewCB->getModule().getASTContext();
  auto OptionalEscapingClosureTy =
      SILType::getOptionalType(SentinelClosure->getType());
  auto *NoneDecl = Context.getOptionalNoneDecl();
  {
    SILBuilderWithScope B(Fn.getEntryBlock()->begin());
    Slot = B.createAllocStack(generatedLoc, OptionalEscapingClosureTy);
    // Store None to it.
    B.createStore(generatedLoc,
                  B.createEnum(generatedLoc, SILValue(), NoneDecl,
                               OptionalEscapingClosureTy),
                  Slot, StoreOwnershipQualifier::Init);
  }
  {
    SILBuilderWithScope B(NewCB);
    // Store the closure sentinel (the copy_block_without_escaping closure
    // operand consumed at +1, so we don't need a copy) to it.
    B.createDestroyAddr(generatedLoc, Slot); // We could be in a loop.
    auto *SomeDecl = Context.getOptionalSomeDecl();
    B.createStore(generatedLoc,
                  B.createEnum(generatedLoc, SentinelClosure, SomeDecl,
                               OptionalEscapingClosureTy),
                  Slot, StoreOwnershipQualifier::Init);
  }

  for (auto LifetimeEndPoint : LifetimeEndPoints) {
    SILBuilderWithScope B(LifetimeEndPoint);
    SILValue isEscaping;
    B.emitScopedBorrowOperation(generatedLoc, Slot, [&](SILValue value) {
      isEscaping = B.createIsEscapingClosure(
          Loc, value, IsEscapingClosureInst::ObjCEscaping);
    });
    B.createCondFail(Loc, isEscaping);
    B.createDestroyAddr(generatedLoc, Slot);
    // Store None to it.
    B.createStore(generatedLoc,
                  B.createEnum(generatedLoc, SILValue(), NoneDecl,
                               OptionalEscapingClosureTy),
                  Slot, StoreOwnershipQualifier::Init);
  }

  // Insert the dealloc_stack in all exiting blocks.
  for (auto *ExitBlock: ExitingBlocks) {
    auto *Terminator = ExitBlock->getTerminator();
    SILBuilderWithScope B(Terminator);
    B.createDeallocStack(generatedLoc, Slot);
  }

  return true;
}

static bool fixupClosureLifetimes(SILFunction &Fn) {
  bool Changed = false;

  // tryExtendLifetimeToLastUse uses a cache of recursive instruction use
  // queries.
  llvm::DenseMap<SILInstruction *, SILInstruction *> MemoizedQueries;

  for (auto &BB : Fn) {
    auto I = BB.begin();
    while (I != BB.end()) {
      SILInstruction *Inst = &*I;
      ++I;

      // Handle, copy_block_without_escaping instructions.
      if (auto *CB = dyn_cast<CopyBlockWithoutEscapingInst>(Inst)) {
        Changed |= fixupCopyBlockWithoutEscaping(CB);
        continue;
      }

      // Otherwise, look at convert_escape_to_noescape [not_guaranteed]
      // instructions.
      auto *Cvt = dyn_cast<ConvertEscapeToNoEscapeInst>(Inst);
      if (!Cvt || Cvt->isLifetimeGuaranteed())
        continue;

      // First try to peephole a known pattern.
      if (!DisableConvertEscapeToNoEscapeSwitchEnumPeephole &&
          trySwitchEnumPeephole(Cvt)) {
        Changed |= true;
        continue;
      }

      if (tryExtendLifetimeToLastUse(Cvt, MemoizedQueries)) {
        Changed |= true;
        continue;
      }

      // Otherwise, extend the lifetime of the operand to the end of the
      // function.
      extendLifetimeToEndOfFunction(Fn, Cvt);
      Changed |= true;
    }
  }
  return Changed;
}

/// Fix-up the lifetime of the escaping closure argument of
/// convert_escape_to_noescape [not_guaranteed] instructions.
///
/// convert_escape_to_noescape [not_guaranteed] assume that someone guarantees
/// the lifetime of the operand for the duration of the trivial closure result.
/// SILGen does not guarantee this for '[not_guaranteed]' instructions so we
/// ensure it here.
namespace {
class ClosureLifetimeFixup : public SILFunctionTransform {

  /// The entry point to the transformation.
  void run() override {
    // Don't rerun diagnostics on deserialized functions.
    if (getFunction()->wasDeserializedCanonical())
      return;

    // Fixup convert_escape_to_noescape [not_guaranteed] and
    // copy_block_without_escaping instructions.
    if (fixupClosureLifetimes(*getFunction()))
      invalidateAnalysis(SILAnalysis::InvalidationKind::FunctionBody);
    LLVM_DEBUG(getFunction()->verify());

  }

};
} // end anonymous namespace

SILTransform *swift::createClosureLifetimeFixup() {
  return new ClosureLifetimeFixup();
}