Generalize and cleanup code for existential specialization.

Generalizes the ConcreteExistentialInfo abstraction so it can be used
both by the ExistentialSpecializer and SILCombine, allowing redundant
code in ExistentialSpecializer.cpp to be deleted.

Splits OpenedArchetypeInfo from ConcreteExistentialInfo. Adds a
ConcreteOpenedArchetypeInfo convenience wrapper around them both, for
use wherever we were originally using ConcreteExistentialInfo.

Splits getAddressOfStackInit into getStackInitInst, This is cleaner and
allows both the ExistentialSpecializer and SILCombine to handle more
interesting cases in the future, like unconditional_checked_cast.

Creates utilities, initializeSubstitutionMap, and
initializeConcreteTypeDef to simplify an generalize
ConcreteExistentialInfo.

While rewriting ExistentialSpecializer to use the new
abstraction, I fixed a latent bug in which is was using a SIL
argument index as a function type parameter index (this would
have broken up if/when we decide to enable calls with indirect
results).

Author: atrick

include/swift/SILOptimizer/Utils/Existential.h

@@ -21,33 +21,35 @@
 
 namespace swift {
 
-/// Find InitExistential from global_addr and copy_addr.
-SILValue findInitExistentialFromGlobalAddrAndCopyAddr(GlobalAddrInst *GAI,
-                                                      CopyAddrInst *CAI);
-
-/// Find InitExistential from global_addr and an apply argument.
-SILValue findInitExistentialFromGlobalAddrAndApply(GlobalAddrInst *GAI,
-                                                   ApplySite Apply, int ArgIdx);
-
-/// Returns the address of an object with which the stack location \p ASI is
-/// initialized. This is either a init_existential_addr or the destination of a
-/// copy_addr. Returns a null value if the address does not dominate the
-/// alloc_stack user \p ASIUser.
-/// If the value is copied from another stack location, \p isCopied is set to
-/// true.
-SILValue getAddressOfStackInit(SILValue allocStackAddr, SILInstruction *ASIUser,
-                               bool &isCopied);
-
-/// Find the init_existential, which could be used to determine a concrete
-/// type of the value used by \p openedUse.
-/// If the value is copied from another stack location, \p isCopied is set to
-/// true.
+/// Record information about an opened archetype.
 ///
-/// FIXME: replace all uses of this with ConcreteExistentialInfo.
-SILInstruction *findInitExistential(Operand &openedUse,
-                                    ArchetypeType *&OpenedArchetype,
-                                    SILValue &OpenedArchetypeDef,
-                                    bool &isCopied);
+/// This is used to determine whether a generic call argument originates from
+/// an opened existential. For example:
+/// %o = open_existential_ref %e : $P & Q to $@opened("PQ") P & Q
+/// %r = apply %f<@opened("PQ") P & Q>(%o)
+///   : $@convention(method) <τ_0_0 where τ_0_0 : P, τ_0_0 : Q>
+///     (@guaranteed τ_0_0) -> @owned τ_0_0
+///
+/// When successfull, ConcreteExistentialInfo can be used to determine the
+/// concrete type of the opened existential.
+struct OpenedArchetypeInfo {
+  ArchetypeType *OpenedArchetype = nullptr;
+  // The opened value.
+  SingleValueInstruction *OpenedArchetypeValue;
+  // The existential value.
+  SILValue ExistentialValue;
+  // True if the openedValue is copied from another stack location
+  bool isOpenedValueCopied = false;
+
+  // Construct a valid instance if the given use originates from a recognizable
+  // OpenedArchetype instruction.
+  OpenedArchetypeInfo(Operand &use);
+
+  bool isValid() const {
+    assert(!OpenedArchetype || (OpenedArchetypeValue && ExistentialValue));
+    return OpenedArchetype;
+  }
+};
 
 /// Record conformance and concrete type info derived from an init_existential
 /// value that is reopened before it's use. This is useful for finding the
@@ -60,20 +62,16 @@ SILInstruction *findInitExistential(Operand &openedUse,
 ///   : $@convention(method) <τ_0_0 where τ_0_0 : P, τ_0_0 : Q>
 ///     (@guaranteed τ_0_0) -> @owned τ_0_0
 struct ConcreteExistentialInfo {
-  // The opened type passed as self. `$@opened("PQ")` above.
-  // This is also the replacement type of the method's Self type.
-  ArchetypeType *OpenedArchetype = nullptr;
-  // The definition of the OpenedArchetype.
-  SILValue OpenedArchetypeDef;
-  // True if the openedValue is copied from another stack location
-  bool isCopied;
-  // The init_existential instruction that produces the opened existential.
-  SILInstruction *InitExistential = nullptr;
   // The existential type of the self argument before it is opened,
   // produced by an init_existential.
   CanType ExistentialType;
   // The concrete type of self from the init_existential. `$C` above.
   CanType ConcreteType;
+  // The concrete value used to initialize the opened existential.
+  // `%c` in the above comment.
+  SILValue ConcreteValue;
+  // True if the ConcreteValue is copied from another stack location
+  bool isConcreteValueCopied = false;
   // When ConcreteType is itself an opened existential, record the type
   // definition. May be nullptr for a valid AppliedConcreteType.
   SingleValueInstruction *ConcreteTypeDef = nullptr;
@@ -82,31 +80,20 @@ struct ConcreteExistentialInfo {
   // and includes the full list of existential conformances.
   // signature: <P & Q>, replacement: $C : conformances: [$P, $Q]
   SubstitutionMap ExistentialSubs;
-  // The value of concrete type used to initialize the existential. `%c` above.
-  SILValue ConcreteValue;
 
-  // Search for a recognized pattern in which the given value is an opened
-  // existential that was previously initialized to a concrete type.
-  // Constructs a valid ConcreteExistentialInfo object if successfull.
-  ConcreteExistentialInfo(Operand &openedUse);
+  // Search for a recognized pattern in which the given existential value is
+  // initialized to a concrete type. Constructs a valid ConcreteExistentialInfo
+  // object if successfull.
+  ConcreteExistentialInfo(SILValue existential, SILInstruction *user);
 
   // This constructor initializes a ConcreteExistentialInfo based on already
-  // known ConcreteType and ProtocolDecl pair. It determines the
-  // OpenedArchetypeDef for the ArgOperand that will be used by unchecked_cast
-  // instructions to cast OpenedArchetypeDef to ConcreteType.
-  ConcreteExistentialInfo(Operand &ArgOperand, CanType ConcreteType,
-                          ProtocolDecl *Protocol);
+  // known ConcreteType and ProtocolDecl pair.
+  ConcreteExistentialInfo(SILValue existential, SILInstruction *user,
+                          CanType ConcreteType, ProtocolDecl *Protocol);
 
   /// For scenerios where ConcreteExistentialInfo is created using a known
-  /// ConcreteType and ProtocolDecl, both of InitExistential
-  /// and ConcreteValue can be null. So there is no need for explicit check for
-  /// not null for them instead we assert on (!InitExistential ||
-  /// ConcreteValue). 
-  bool isValid() const {
-    assert(!InitExistential || ConcreteValue);
-    return OpenedArchetype && OpenedArchetypeDef && ConcreteType &&
-           !ExistentialSubs.empty();
-  }
+  /// ConcreteType and ProtocolDecl, the ConcreteValue can be null.
+  bool isValid() const { return ConcreteType && !ExistentialSubs.empty(); }
 
   // Do a conformance lookup on ConcreteType with the given requirement, P. If P
   // is satisfiable based on the existential's conformance, return the new
@@ -116,6 +103,35 @@ struct ConcreteExistentialInfo {
     CanType selfTy = P->getSelfInterfaceType()->getCanonicalType();
     return ExistentialSubs.lookupConformance(selfTy, P);
   }
+
+private:
+  void initializeSubstitutionMap(
+      ArrayRef<ProtocolConformanceRef> ExistentialConformances, SILModule *M);
+
+  void initializeConcreteTypeDef(SILInstruction *typeConversionInst);
+};
+
+// Convenience for tracking both the OpenedArchetypeInfo and
+// ConcreteExistentialInfo from the same SILValue.
+struct ConcreteOpenedExistentialInfo {
+  OpenedArchetypeInfo OAI;
+  // If CEI has a value, it must be valid.
+  Optional<ConcreteExistentialInfo> CEI;
+
+  ConcreteOpenedExistentialInfo(Operand &use);
+
+  // Provide a whole module type-inferred ConcreteType to fall back on if the
+  // concrete type cannot be determined from data flow.
+  ConcreteOpenedExistentialInfo(Operand &use, CanType concreteType,
+                                ProtocolDecl *protocol);
+
+  bool isValid() const {
+    if (!CEI)
+      return false;
+
+    assert(CEI->isValid());
+    return true;
+  }
 };
 
 } // end namespace swift

lib/SILOptimizer/FunctionSignatureTransforms/ExistentialSpecializer.cpp

@@ -38,12 +38,12 @@ class ExistentialSpecializer : public SILFunctionTransform {
   /// Determine if the current function is a target for existential
   /// specialization of args.
   bool canSpecializeExistentialArgsInFunction(
-      ApplySite &Apply,
+      FullApplySite &Apply,
       llvm::SmallDenseMap<int, ExistentialTransformArgumentDescriptor>
           &ExistentialArgDescriptor);
 
   /// Can Callee be specialized?
-  bool canSpecializeCalleeFunction(ApplySite &Apply);
+  bool canSpecializeCalleeFunction(FullApplySite &Apply);
 
   /// Specialize existential args in function F.
   void specializeExistentialArgsInAppliesWithinFunction(SILFunction &F);
@@ -70,82 +70,6 @@ class ExistentialSpecializer : public SILFunctionTransform {
 };
 } // namespace
 
-/// Find concrete type from init_existential_refs/addrs.
-static bool findConcreteTypeFromInitExistential(SILValue Arg,
-                                                CanType &ConcreteType) {
-  if (auto *IER = dyn_cast<InitExistentialRefInst>(Arg)) {
-    ConcreteType = IER->getFormalConcreteType();
-    return true;
-  } else if (auto *IE = dyn_cast<InitExistentialAddrInst>(Arg)) {
-    ConcreteType = IE->getFormalConcreteType();
-    return true;
-  }
-  return false;
-}
-
-/// Find the concrete type of the existential argument. Wrapper
-/// for findInitExistential in Local.h/cpp. In future, this code
-/// can move to Local.h/cpp.
-static bool findConcreteType(ApplySite AI, int ArgIdx, CanType &ConcreteType) {
-  bool isCopied = false;
-  auto Arg = AI.getArgument(ArgIdx);
-
-  /// Ignore any unconditional cast instructions. Is it Safe? Do we need to
-  /// also add UnconditionalCheckedCastAddrInst? TODO.
-  if (auto *Instance = dyn_cast<UnconditionalCheckedCastInst>(Arg)) {
-    Arg = Instance->getOperand();
-  }
-
-  /// Return init_existential if the Arg is global_addr.
-  if (auto *GAI = dyn_cast<GlobalAddrInst>(Arg)) {
-    SILValue InitExistential =
-        findInitExistentialFromGlobalAddrAndApply(GAI, AI, ArgIdx);
-    /// If the Arg is already init_existential, return the concrete type.
-    if (InitExistential &&
-        findConcreteTypeFromInitExistential(InitExistential, ConcreteType)) {
-      return true;
-    }
-  }
-
-  /// Handle AllocStack instruction separately.
-  if (auto *Instance = dyn_cast<AllocStackInst>(Arg)) {
-    if (SILValue Src =
-            getAddressOfStackInit(Instance, AI.getInstruction(), isCopied)) {
-      Arg = Src;
-    }
-  }
-
-  /// If the Arg is already init_existential after getAddressofStackInit
-  /// call, return the concrete type.
-  if (findConcreteTypeFromInitExistential(Arg, ConcreteType)) {
-    return true;
-  }
-
-  /// Call findInitExistential and determine the init_existential.
-  ArchetypeType *OpenedArchetype = nullptr;
-  SILValue OpenedArchetypeDef;
-  auto FAS = FullApplySite::isa(AI.getInstruction());
-  if (!FAS)
-    return false;
-  SILInstruction *InitExistential =
-      findInitExistential(FAS.getArgumentOperands()[ArgIdx], OpenedArchetype,
-                          OpenedArchetypeDef, isCopied);
-  if (!InitExistential) {
-    LLVM_DEBUG(llvm::dbgs() << "ExistentialSpecializer Pass: Bail! Due to "
-                               "findInitExistential\n";);
-    return false;
-  }
-
-  /// Return the concrete type from init_existential returned from
-  /// findInitExistential.
-  if (auto *SingleVal = InitExistential->castToSingleValueInstruction()) {
-    if (findConcreteTypeFromInitExistential(SingleVal, ConcreteType)) {
-      return true;
-    }
-  }
-  return false;
-}
-
 /// Check if the argument Arg is used in a destroy_use instruction.
 static void
 findIfCalleeUsesArgInDestroyUse(SILValue Arg,
@@ -162,17 +86,21 @@ findIfCalleeUsesArgInDestroyUse(SILValue Arg,
 /// Check if any apply argument meets the criteria for existential
 /// specialization.
 bool ExistentialSpecializer::canSpecializeExistentialArgsInFunction(
-    ApplySite &Apply,
+    FullApplySite &Apply,
     llvm::SmallDenseMap<int, ExistentialTransformArgumentDescriptor>
         &ExistentialArgDescriptor) {
   auto *F = Apply.getReferencedFunction();
-  auto Args = F->begin()->getFunctionArguments();
+  auto CalleeArgs = F->begin()->getFunctionArguments();
   bool returnFlag = false;
 
-  /// Analyze the argument for protocol conformance.
-  for (unsigned Idx = 0, Num = Args.size(); Idx < Num; ++Idx) {
-    auto Arg = Args[Idx];
-    auto ArgType = Arg->getType();
+  /// Analyze the argument for protocol conformance.  Iterator over the callee's
+  /// function arguments.  The same SIL argument index is used for both caller
+  /// and callee side arguments.
+  auto origCalleeConv = Apply.getOrigCalleeConv();
+  assert(Apply.getCalleeArgIndexOfFirstAppliedArg() == 0);
+  for (unsigned Idx = 0, Num = CalleeArgs.size(); Idx < Num; ++Idx) {
+    auto CalleeArg = CalleeArgs[Idx];
+    auto ArgType = CalleeArg->getType();
     auto SwiftArgType = ArgType.getASTType();
 
     /// Checking for AnyObject and Any is added to ensure that we do not blow up
@@ -184,18 +112,21 @@ bool ExistentialSpecializer::canSpecializeExistentialArgsInFunction(
       continue;
 
     auto ExistentialRepr =
-        Arg->getType().getPreferredExistentialRepresentation(F->getModule());
+        CalleeArg->getType().getPreferredExistentialRepresentation(
+            F->getModule());
     if (ExistentialRepr != ExistentialRepresentation::Opaque &&
           ExistentialRepr != ExistentialRepresentation::Class)
       continue;
 
     /// Find the concrete type.
-    CanType ConcreteType;
-    if (!findConcreteType(Apply, Idx, ConcreteType)) {
+    Operand &ArgOper = Apply.getArgumentRef(Idx);
+    CanType ConcreteType =
+        ConcreteExistentialInfo(ArgOper.get(), ArgOper.getUser()).ConcreteType;
+    if (!ConcreteType) {
       LLVM_DEBUG(
           llvm::dbgs()
-              << "ExistentialSpecializer Pass: Bail! Due to findConcreteType "
-                 "for callee:"
+              << "ExistentialSpecializer Pass: Bail! cannot find ConcreteType "
+                 "for call argument to:"
               << F->getName() << " in caller:"
               << Apply.getInstruction()->getParent()->getParent()->getName()
               << "\n";);
@@ -205,15 +136,15 @@ bool ExistentialSpecializer::canSpecializeExistentialArgsInFunction(
     /// Determine attributes of the existential addr arguments such as
     /// destroy_use, immutable_access. 
     ExistentialTransformArgumentDescriptor ETAD;
-    ETAD.AccessType =
-        Apply.getOrigCalleeType()->getParameters()[Idx].isIndirectMutating() ||
-                Apply.getOrigCalleeType()->getParameters()[Idx].isConsumed()
-            ? OpenedExistentialAccess::Mutable
-            : OpenedExistentialAccess::Immutable;
+    auto paramInfo = origCalleeConv.getParamInfoForSILArg(Idx);
+    ETAD.AccessType = (paramInfo.isIndirectMutating() || paramInfo.isConsumed())
+                          ? OpenedExistentialAccess::Mutable
+                          : OpenedExistentialAccess::Immutable;
     ETAD.DestroyAddrUse = false;
-    if ((Args[Idx]->getType().getPreferredExistentialRepresentation(
-            F->getModule())) != ExistentialRepresentation::Class)
-      findIfCalleeUsesArgInDestroyUse(Arg, ETAD);
+    if ((CalleeArgs[Idx]->getType().getPreferredExistentialRepresentation(
+            F->getModule()))
+        != ExistentialRepresentation::Class)
+      findIfCalleeUsesArgInDestroyUse(CalleeArg, ETAD);
 
     /// Save the attributes
     ExistentialArgDescriptor[Idx] = ETAD;
@@ -226,12 +157,7 @@ bool ExistentialSpecializer::canSpecializeExistentialArgsInFunction(
 }
 
 /// Determine if this callee function can be specialized or not.
-bool ExistentialSpecializer::canSpecializeCalleeFunction(ApplySite &Apply) {
-
-  /// Do not handle partial applies.
-  if (isa<PartialApplyInst>(Apply.getInstruction())) {
-    return false;
-  }
+bool ExistentialSpecializer::canSpecializeCalleeFunction(FullApplySite &Apply) {
 
   /// Determine the caller of the apply.
   auto *Callee = Apply.getReferencedFunction();
@@ -285,7 +211,7 @@ void ExistentialSpecializer::specializeExistentialArgsInAppliesWithinFunction(
       auto *I = &*It;
 
       /// Is it an apply site?
-      ApplySite Apply = ApplySite::isa(I);
+      FullApplySite Apply = FullApplySite::isa(I);
       if (!Apply)
         continue;