//===--- Existential.cpp - Functions analyzing existentials.  -------------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "swift/SILOptimizer/Utils/Existential.h"
#include "swift/AST/Module.h"
#include "swift/AST/ProtocolConformance.h"
#include "swift/SIL/BasicBlockUtils.h"
#include "swift/SIL/InstructionUtils.h"
#include "swift/SILOptimizer/Utils/CFGOptUtils.h"
#include "swift/SILOptimizer/Utils/InstOptUtils.h"
#include "llvm/ADT/SmallPtrSet.h"

using namespace swift;

/// Determine InitExistential from global_addr.
/// %3 = global_addr @$P : $*SomeP
/// %4 = init_existential_addr %3 : $*SomeP, $SomeC
/// %5 = alloc_ref $SomeC
/// store %5 to %4 : $*SomeC
/// %8 = alloc_stack $SomeP
/// copy_addr %3 to [initialization] %8 : $*SomeP
/// %10 = apply %9(%3) : $@convention(thin) (@in_guaranteed SomeP)
/// Assumptions: Insn is a direct user of GAI (e.g., copy_addr or 
/// apply pattern shown above) and that a valid init_existential_addr 
/// value is returned only if it can prove that the value it 
/// initializes is the same value at the use point.
static InitExistentialAddrInst *
findInitExistentialFromGlobalAddr(GlobalAddrInst *GAI, SILInstruction *Insn) {
  /// Check for a single InitExistential usage for GAI and
  /// a simple dominance check: both InitExistential and Insn are in
  /// the same basic block and only one InitExistential
  /// occurs between GAI and Insn.
  llvm::SmallPtrSet<SILInstruction *, 8> IEUses;
  for (auto *Use : GAI->getUses()) {
    if (auto *InitExistential =
            dyn_cast<InitExistentialAddrInst>(Use->getUser())) {
      IEUses.insert(InitExistential);
    }
  }

  /// No InitExistential found in the basic block.
  if (IEUses.empty())
    return nullptr;

  /// Walk backwards from Insn instruction till the begining of the basic block
  /// looking for an InitExistential.
  InitExistentialAddrInst *SingleIE = nullptr;
  for (auto II = Insn->getIterator().getReverse(),
            IE = Insn->getParent()->rend();
       II != IE; ++II) {
    if (!IEUses.count(&*II))
      continue;
    if (SingleIE)
      return nullptr;

    SingleIE = cast<InitExistentialAddrInst>(&*II);
  }
  return SingleIE;
}

/// Returns the instruction that initializes the given stack address. This is
/// currently either a init_existential_addr, unconditional_checked_cast_addr,
/// or copy_addr (if the instruction initializing the source of the copy cannot
/// be determined). Returns nullptr if the initializer does not dominate the
/// alloc_stack user \p ASIUser.  If the value is copied from another stack
/// location, \p isCopied is set to true.
///
/// allocStackAddr may either itself be an AllocStackInst or an
/// InitEnumDataAddrInst that projects the value of an AllocStackInst.
static SILInstruction *getStackInitInst(SILValue allocStackAddr,
                                        SILInstruction *ASIUser,
                                        bool &isCopied) {
  SILInstruction *SingleWrite = nullptr;
  // Check that this alloc_stack is initialized only once.
  for (auto Use : allocStackAddr->getUses()) {
    auto *User = Use->getUser();

    // Ignore instructions which don't write to the stack location.
    // Also ignore ASIUser (only kicks in if ASIUser is the original apply).
    if (isa<DeallocStackInst>(User) || isa<DebugValueAddrInst>(User) ||
        isa<DestroyAddrInst>(User) || isa<WitnessMethodInst>(User) ||
        isa<DeinitExistentialAddrInst>(User) ||
        isa<OpenExistentialAddrInst>(User) || User == ASIUser) {
      continue;
    }
    if (auto *CAI = dyn_cast<CopyAddrInst>(User)) {
      if (CAI->getDest() == allocStackAddr) {
        if (SingleWrite)
          return nullptr;
        SingleWrite = CAI;
        isCopied = true;
      }
      continue;
    }
    // An unconditional_checked_cast_addr also copies a value into this addr.
    if (auto *UCCAI = dyn_cast<UnconditionalCheckedCastAddrInst>(User)) {
      if (UCCAI->getDest() == allocStackAddr) {
        if (SingleWrite)
          return nullptr;
        SingleWrite = UCCAI;
        isCopied = true;
      }
      continue;
    }
    if (isa<InitExistentialAddrInst>(User)) {
      if (SingleWrite)
        return nullptr;
      SingleWrite = User;
      continue;
    }
    if (isa<ApplyInst>(User) || isa<TryApplyInst>(User)) {
      // Ignore function calls which do not write to the stack location.
      auto Conv = FullApplySite(User).getArgumentConvention(*Use);
      if (Conv != SILArgumentConvention::Indirect_In &&
          Conv != SILArgumentConvention::Indirect_In_Guaranteed)
        return nullptr;
      continue;
    }
    // Bail if there is any unknown (and potentially writing) instruction.
    return nullptr;
  }
  if (!SingleWrite)
    return nullptr;

  // A very simple dominance check. As ASI is an operand of ASIUser,
  // SingleWrite dominates ASIUser if it is in the same block as ASI or
  // ASIUser. (SingleWrite can't occur after ASIUser because the address would
  // be uninitialized on use).
  //
  // If allocStack holds an Optional, then ASI is an InitEnumDataAddrInst
  // projection and not strictly an operand of ASIUser. We rely on the guarantee
  // that this InitEnumDataAddrInst must occur before the InjectEnumAddrInst
  // that was the source of the existential address.
  SILBasicBlock *BB = SingleWrite->getParent();
  if (BB != allocStackAddr->getParentBlock() && BB != ASIUser->getParent())
    return nullptr;

  if (auto *IE = dyn_cast<InitExistentialAddrInst>(SingleWrite))
    return IE;

  if (auto *UCCA = dyn_cast<UnconditionalCheckedCastAddrInst>(SingleWrite)) {
    assert(isCopied && "isCopied not set for a unconditional_checked_cast_addr");
    return UCCA;
  }
  auto *CAI = cast<CopyAddrInst>(SingleWrite);
  assert(isCopied && "isCopied not set for a copy_addr");
  // Attempt to recurse to find a concrete type.
  if (auto *ASI = dyn_cast<AllocStackInst>(CAI->getSrc()))
    return getStackInitInst(ASI, CAI, isCopied);

  // Peek through a stack location holding an Enum.
  //   %stack_adr = alloc_stack
  //   %data_adr  = init_enum_data_addr %stk_adr
  //   %enum_adr  = inject_enum_addr %stack_adr
  //   %copy_src  = unchecked_take_enum_data_addr %enum_adr
  // Replace %copy_src with %data_adr and recurse.
  //
  // TODO: a general Optional elimination sil-combine could
  // supersede this check.
  if (auto *UTEDAI = dyn_cast<UncheckedTakeEnumDataAddrInst>(CAI->getSrc())) {
    if (InitEnumDataAddrInst *IEDAI = findInitAddressForTrivialEnum(UTEDAI))
      return getStackInitInst(IEDAI, CAI, isCopied);
  }

  // Check if the CAISrc is a global_addr.
  if (auto *GAI = dyn_cast<GlobalAddrInst>(CAI->getSrc()))
    return findInitExistentialFromGlobalAddr(GAI, CAI);

  // If the source of the copy cannot be determined, return the copy itself
  // because the caller may have special handling for the source address.
  return CAI;
}

/// Return the address of the value used to initialize the given stack location.
/// If the value originates from init_existential_addr, then it will be a
/// different type than \p allocStackAddr.
static SILValue getAddressOfStackInit(SILValue allocStackAddr,
                                      SILInstruction *ASIUser, bool &isCopied) {
  SILInstruction *initI = getStackInitInst(allocStackAddr, ASIUser, isCopied);
  if (!initI)
    return SILValue();

  if (auto *IEA = dyn_cast<InitExistentialAddrInst>(initI))
    return IEA;

  if (auto *CAI = dyn_cast<CopyAddrInst>(initI))
    return CAI->getSrc();

  return SILValue();
}

/// Check if the given operand originates from a recognized OpenArchetype
/// instruction. If so, return the Opened, otherwise return nullptr.
OpenedArchetypeInfo::OpenedArchetypeInfo(Operand &use) {
  SILValue openedVal = use.get();
  SILInstruction *user = use.getUser();
  if (auto *instance = dyn_cast<AllocStackInst>(openedVal)) {
    // Handle:
    //   %opened = open_existential_addr
    //   %instance = alloc $opened
    //   copy_addr %opened to %stack
    //   <opened_use> %instance
    if (auto stackInitVal =
            getAddressOfStackInit(instance, user, isOpenedValueCopied)) {
      openedVal = stackInitVal;
    }
  }
  if (auto *Open = dyn_cast<OpenExistentialAddrInst>(openedVal)) {
    OpenedArchetype = Open->getType().castTo<ArchetypeType>();
    OpenedArchetypeValue = Open;
    ExistentialValue = Open->getOperand();
    return;
  }
  if (auto *Open = dyn_cast<OpenExistentialRefInst>(openedVal)) {
    OpenedArchetype = Open->getType().castTo<ArchetypeType>();
    OpenedArchetypeValue = Open;
    ExistentialValue = Open->getOperand();
    return;
  }
  if (auto *Open = dyn_cast<OpenExistentialMetatypeInst>(openedVal)) {
    auto Ty = Open->getType().getASTType();
    while (auto Metatype = dyn_cast<MetatypeType>(Ty))
      Ty = Metatype.getInstanceType();
    OpenedArchetype = cast<ArchetypeType>(Ty);
    OpenedArchetypeValue = Open;
    ExistentialValue = Open->getOperand();
  }
}

/// Initialize ExistentialSubs from the given conformance list, using the
/// already initialized ExistentialType and ConcreteType.
void ConcreteExistentialInfo::initializeSubstitutionMap(
    ArrayRef<ProtocolConformanceRef> ExistentialConformances, SILModule *M) {

  // Construct a single-generic-parameter substitution map directly to the
  // ConcreteType with this existential's full list of conformances.
  CanGenericSignature ExistentialSig =
      M->getASTContext().getExistentialSignature(ExistentialType,
                                                 M->getSwiftModule());
  ExistentialSubs = SubstitutionMap::get(ExistentialSig, {ConcreteType},
                                         ExistentialConformances);
  assert(isValid());
}

/// If the ConcreteType is an opened existential, also initialize
/// ConcreteTypeDef to the definition of that type.
void ConcreteExistentialInfo::initializeConcreteTypeDef(
    SILInstruction *typeConversionInst) {
  if (!ConcreteType->isOpenedExistential())
    return;

  assert(isValid());

  // If the concrete type is another existential, we're "forwarding" an
  // opened existential type, so we must keep track of the original
  // defining instruction.
  if (!typeConversionInst->getTypeDependentOperands().empty()) {
    ConcreteTypeDef = cast<SingleValueInstruction>(
        typeConversionInst->getTypeDependentOperands()[0].get());
    return;
  }

  auto typeOperand =
      cast<InitExistentialMetatypeInst>(typeConversionInst)->getOperand();
  assert(typeOperand->getType().hasOpenedExistential()
         && "init_existential is supposed to have a typedef operand");
  ConcreteTypeDef = cast<SingleValueInstruction>(typeOperand);
}

/// Construct this ConcreteExistentialInfo based on the given existential use.
///
/// Finds the init_existential, or an address with concrete type used to
/// initialize the given \p openedUse. If the value is copied
/// from another stack location, \p isCopied is set to true.
///
/// If successful, ConcreteExistentialInfo will be valid upon return, with the
/// following fields assigned:
/// - ExistentialType
/// - isCopied
/// - ConcreteType
/// - ConcreteValue
/// - ConcreteTypeDef
/// - ExistentialSubs
ConcreteExistentialInfo::ConcreteExistentialInfo(SILValue existential,
                                                 SILInstruction *user) {
  if (existential->getType().isAddress()) {
    auto *ASI = dyn_cast<AllocStackInst>(existential);
    if (!ASI)
      return;

    SILInstruction *stackInit =
        getStackInitInst(ASI, user, isConcreteValueCopied);
    if (!stackInit)
      return;

    if (auto *IE = dyn_cast<InitExistentialAddrInst>(stackInit)) {
      ExistentialType = IE->getOperand()->getType().getASTType();
      ConcreteType = IE->getFormalConcreteType();
      ConcreteValue = IE;
      initializeSubstitutionMap(IE->getConformances(), &IE->getModule());
      initializeConcreteTypeDef(IE);
      return;
    }
    // TODO: Once we have a way to introduce more constrained archetypes, handle
    // any unconditional_checked_cast that wasn't already statically eliminated.
    //
    // Unexpected stack write.
    return;
  }

  if (auto *IER = dyn_cast<InitExistentialRefInst>(existential)) {
    ExistentialType = IER->getType().getASTType();
    ConcreteType = IER->getFormalConcreteType();
    ConcreteValue = IER->getOperand();
    initializeSubstitutionMap(IER->getConformances(), &IER->getModule());
    initializeConcreteTypeDef(IER);
    return;
  }

  if (auto *IEM = dyn_cast<InitExistentialMetatypeInst>(existential)) {
    ExistentialType = IEM->getType().getASTType();
    ConcreteValue = IEM->getOperand();
    ConcreteType = ConcreteValue->getType().getASTType();
    while (auto InstanceType =
               dyn_cast<ExistentialMetatypeType>(ExistentialType)) {
      ExistentialType = InstanceType.getInstanceType();
      ConcreteType = cast<MetatypeType>(ConcreteType).getInstanceType();
    }
    initializeSubstitutionMap(IEM->getConformances(), &IEM->getModule());
    initializeConcreteTypeDef(IEM);
    return;
  }
  // Unrecognized opened existential producer.
  return;
}

/// Initialize a ConcreteExistentialInfo based on a concrete type and protocol
/// declaration that has already been computed via whole module type
/// inference. A cast instruction will be introduced to produce the concrete
/// value from the opened value.
///
/// The simpler constructor taking only the existential value is preferred
/// because it generates simpler SIL and does not require an extra
/// cast. However, if that constructor fails to produce a valid
/// ConcreteExistentialInfo, this constructor may succeed because it doesn't
/// needs to rediscover the whole-module inferred ConcreteTypeCandidate.
ConcreteExistentialInfo::ConcreteExistentialInfo(SILValue existential,
                                                 SILInstruction *user,
                                                 CanType ConcreteTypeCandidate,
                                                 ProtocolDecl *Protocol) {
  SILModule *M = existential->getModule();

  // We have the open_existential; we still need the conformance.
  auto ConformanceRef =
      M->getSwiftModule()->conformsToProtocol(ConcreteTypeCandidate, Protocol);
  if (!ConformanceRef)
    return;

  // Assert that the conformance is complete.
  auto *ConcreteConformance = ConformanceRef.getValue().getConcrete();
  assert(ConcreteConformance->isComplete());

  ConcreteType = ConcreteTypeCandidate;
  // There is no ConcreteValue in this case.

  /// Determine the ExistentialConformances and SubstitutionMap.
  ExistentialType = Protocol->getDeclaredType()->getCanonicalType();
  initializeSubstitutionMap(ProtocolConformanceRef(ConcreteConformance), M);

  assert(isValid());
}

ConcreteOpenedExistentialInfo::ConcreteOpenedExistentialInfo(Operand &use)
    : OAI(use) {
  if (!OAI.isValid())
    return;

  CEI.emplace(OAI.ExistentialValue, OAI.OpenedArchetypeValue);
  if (!CEI->isValid()) {
    CEI.reset();
    return;
  }
  CEI->isConcreteValueCopied |= OAI.isOpenedValueCopied;
}

ConcreteOpenedExistentialInfo::ConcreteOpenedExistentialInfo(
    Operand &use, CanType concreteType, ProtocolDecl *protocol)
    : OAI(use) {
  if (!OAI.isValid())
    return;

  CEI.emplace(OAI.ExistentialValue, OAI.OpenedArchetypeValue, concreteType,
              protocol);
  if (!CEI->isValid()) {
    CEI.reset();
    return;
  }
  CEI->isConcreteValueCopied |= OAI.isOpenedValueCopied;
}