Propagate concrete types for all arguments of an apply instruction

Author: Raj Barik

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -295,9 +295,13 @@ class SILCombiner :
 private:
   FullApplySite rewriteApplyCallee(FullApplySite apply, SILValue callee);
 
-  SILInstruction *createApplyWithConcreteType(FullApplySite Apply,
-                                              const ConcreteExistentialInfo &CEI,
-                                              SILBuilderContext &BuilderCtx);
+  bool canReplaceArg(FullApplySite Apply, const ConcreteExistentialInfo &CEI,
+                     unsigned ArgIdx);
+  SILInstruction *createApplyWithConcreteType(
+      FullApplySite Apply,
+      const llvm::SmallDenseMap<unsigned, const ConcreteExistentialInfo *>
+          &CEIs,
+      SILBuilderContext &BuilderCtx);
 
   // Common utility function to replace the WitnessMethodInst using a
   // BuilderCtx.

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -693,22 +693,26 @@ SILCombiner::propagateSoleConformingType(FullApplySite Apply,
   SILBuilderContext BuilderCtx(M, Builder.getTrackingList());
   replaceWitnessMethodInst(WMI, BuilderCtx, ConcreteType,
                            *(CEI.ExistentialSubs.getConformances().begin()));
+  // Construct the map for Self to be used for createApplyWithConcreteType.
+  llvm::SmallDenseMap<unsigned, const ConcreteExistentialInfo *> CEIs;
+  CEIs.insert(std::pair<unsigned, const ConcreteExistentialInfo *>(
+      Apply.getCalleeArgIndex(Apply.getSelfArgumentOperand()), &CEI));
   /// Create the new apply instruction using the concrete type.
-  auto *NewAI = createApplyWithConcreteType(Apply, CEI, BuilderCtx);
+  auto *NewAI = createApplyWithConcreteType(Apply, CEIs, BuilderCtx);
   return NewAI;
 }
 
 /// Given an Apply and an argument value produced by InitExistentialAddrInst,
 /// return true if the argument can be replaced by a copy of its value.
 ///
 /// FIXME: remove this helper when we can assume SIL opaque values.
-static bool canReplaceCopiedSelf(FullApplySite Apply,
+static bool canReplaceCopiedArg(FullApplySite Apply,
                                  SILInstruction *InitExistential,
-                                 DominanceAnalysis *DA) {
-  // If the witness method mutates self, we cannot replace self with
+                                 DominanceAnalysis *DA, unsigned ArgIdx) {
+  // If the witness method mutates Arg, we cannot replace Arg with
   // the source of a copy. Otherwise the call would modify another value than
-  // the original self.
-  if (Apply.getOrigCalleeType()->getSelfParameter().isIndirectMutating())
+  // the original argument.
+  if (Apply.getOrigCalleeType()->getParameters()[ArgIdx].isIndirectMutating())
     return false;
 
   auto *DT = DA->get(Apply.getFunction());
@@ -751,6 +755,54 @@ static bool canReplaceCopiedSelf(FullApplySite Apply,
   return true;
 }
 
+// Check the legal conditions under which a Arg parameter (specified as ArgIdx)
+// can be replaced with a concrete type. Concrete type info is passed as CEI
+// argument.
+bool SILCombiner::canReplaceArg(FullApplySite Apply,
+                                const ConcreteExistentialInfo &CEI,
+                                unsigned ArgIdx) {
+
+  // Don't specialize apply instructions that return the callee's Arg type,
+  // because this optimization does not know how to substitute types in the
+  // users of this apply. In the function type substitution below, all
+  // references to OpenedArchetype will be substituted. So walk to type to
+  // find all possible references, such as returning Optional<Arg>.
+  if (Apply.getType().getASTType().findIf(
+          [&CEI](Type t) -> bool { return t->isEqual(CEI.OpenedArchetype); })) {
+    return false;
+  }
+  // Bail out if any other arguments or indirect result that refer to the
+  // OpenedArchetype. The following optimization substitutes all occurrences
+  // of OpenedArchetype in the function signature, but will only rewrite the
+  // Arg operand.
+  //
+  // Note that the language does not allow Self to occur in contravariant
+  // position. However, SIL does allow this and it can happen as a result of
+  // upstream transformations. Since this is bail-out logic, it must handle
+  // all verifiable SIL.
+
+  // This bailout check is also needed for non-Self arguments [including Self].
+  unsigned NumApplyArgs = Apply.getNumArguments();
+  for (unsigned Idx = 0; Idx < NumApplyArgs; Idx++) {
+    if (Idx == ArgIdx)
+      continue;
+    if (Apply.getArgument(Idx)->getType().getASTType().findIf(
+            [&CEI](Type t) -> bool {
+              return t->isEqual(CEI.OpenedArchetype);
+            })) {
+      return false;
+    }
+  }
+  // The apply can only be rewritten in terms of the concrete value if it is
+  // legal to pass that value as the Arg argument.
+  if (CEI.isCopied && (!CEI.InitExistential ||
+                       !canReplaceCopiedArg(Apply, CEI.InitExistential, DA, ArgIdx))) {
+    return false;
+  }
+  // It is safe to replace Arg.
+  return true;
+}
+
 /// Rewrite the given method apply instruction in terms of the provided conrete
 /// type information.
 ///
@@ -774,74 +826,61 @@ static bool canReplaceCopiedSelf(FullApplySite Apply,
 /// FIXME: Protocol methods (witness or default) that return Self will be given
 /// a new return type. This implementation fails to update the type signature of
 /// SSA uses in those cases. Currently we bail out on methods that return Self.
-SILInstruction *
-SILCombiner::createApplyWithConcreteType(FullApplySite Apply,
-                                         const ConcreteExistentialInfo &CEI,
-                                         SILBuilderContext &BuilderCtx) {
-  assert(Apply.getOrigCalleeType()->isPolymorphic());
+SILInstruction *SILCombiner::createApplyWithConcreteType(
+    FullApplySite Apply,
+    const llvm::SmallDenseMap<unsigned, const ConcreteExistentialInfo *> &CEIs,
+    SILBuilderContext &BuilderCtx) {
 
-  // Don't specialize apply instructions that return the callee's Self type,
-  // because this optimization does not know how to substitute types in the
-  // users of this apply. In the function type substitution below, all
-  // references to OpenedArchetype will be substituted. So walk to type to find
-  // all possible references, such as returning Optional<Self>.
-  if (Apply.getType().getASTType().findIf(
-          [&CEI](Type t) -> bool { return t->isEqual(CEI.OpenedArchetype); })) {
-    return nullptr;
-  }
-  // Bail out if any non-self arguments or indirect result that refer to the
-  // OpenedArchetype. The following optimization substitutes all occurrences of
-  // OpenedArchetype in the function signature, but will only rewrite the self
-  // operand.
-  //
-  // Note that the language does not allow Self to occur in contravariant
-  // position. However, SIL does allow this and it can happen as a result of
-  // upstream transformations. Since this is bail-out logic, it must handle all
-  // verifiable SIL.
-  for (auto Arg : Apply.getArgumentsWithoutSelf()) {
-    if (Arg->getType().getASTType().findIf([&CEI](Type t) -> bool {
-          return t->isEqual(CEI.OpenedArchetype);
-        })) {
-      return nullptr;
-    }
-  }
-  // The apply can only be rewritten in terms of the concrete value if it is
-  // legal to pass that value as the self argument.
-  if (CEI.isCopied && (!CEI.InitExistential ||
-                       !canReplaceCopiedSelf(Apply, CEI.InitExistential, DA))) {
-    return nullptr;
-  }
+  // Ensure that the callee is polymorphic.
+  assert(Apply.getOrigCalleeType()->isPolymorphic());
 
-  // Create a set of arguments.
+  // Create the new set of arguments to apply including their substitutions.
+  SubstitutionMap NewCallSubs = Apply.getSubstitutionMap();
   SmallVector<SILValue, 8> NewArgs;
-  for (auto Arg : Apply.getArgumentsWithoutSelf()) {
-    NewArgs.push_back(Arg);
+  unsigned NumApplyArgs = Apply.getNumArguments();
+  bool UpdatedArgs = false;
+  for (unsigned ArgIdx = 0; ArgIdx < NumApplyArgs; ArgIdx++) {
+    auto ArgIt = CEIs.find(ArgIdx);
+    if (ArgIt == CEIs.end()) {
+      // Use the old argument if it does not have a valid concrete existential.
+      NewArgs.push_back(Apply.getArgument(ArgIdx));
+      continue;
+    }
+    auto *CEI = ArgIt->second;
+    // Check for Arg's concrete type propagation legality.
+    if (!canReplaceArg(Apply, *CEI, ArgIdx)) {
+      NewArgs.push_back(Apply.getArgument(ArgIdx));
+      continue;
+    }
+    UpdatedArgs = true;
+    // Ensure that we have a concrete value to propagate.
+    assert(CEI->ConcreteValue);
+    NewArgs.push_back(CEI->ConcreteValue);
+    // Form a new set of substitutions where the argument is
+    // replaced with a concrete type.
+    NewCallSubs = NewCallSubs.subst(
+        [&](SubstitutableType *type) -> Type {
+          if (type == CEI->OpenedArchetype)
+            return CEI->ConcreteType;
+          return type;
+        },
+        [&](CanType origTy, Type substTy,
+            ProtocolDecl *proto) -> Optional<ProtocolConformanceRef> {
+          if (origTy->isEqual(CEI->OpenedArchetype)) {
+            assert(substTy->isEqual(CEI->ConcreteType));
+            // Do a conformance lookup on this witness requirement using the
+            // existential's conformances. The witness requirement may be a
+            // base type of the existential's requirements.
+            return CEI->lookupExistentialConformance(proto); 
+          }
+          return ProtocolConformanceRef(proto);
+        });
   }
-  NewArgs.push_back(CEI.ConcreteValue);
-
-  assert(Apply.getOrigCalleeType()->isPolymorphic());
 
-  // Form a new set of substitutions where Self is
-  // replaced by a concrete type.
-  SubstitutionMap OrigCallSubs = Apply.getSubstitutionMap();
-  SubstitutionMap NewCallSubs = OrigCallSubs.subst(
-      [&](SubstitutableType *type) -> Type {
-        if (type == CEI.OpenedArchetype)
-          return CEI.ConcreteType;
-        return type;
-      },
-      [&](CanType origTy, Type substTy,
-          ProtocolDecl *proto) -> Optional<ProtocolConformanceRef> {
-        if (origTy->isEqual(CEI.OpenedArchetype)) {
-          assert(substTy->isEqual(CEI.ConcreteType));
-          // Do a conformance lookup on this witness requirement using the
-          // existential's conformances. The witness requirement may be a base
-          // type of the existential's requirements.
-          return CEI.lookupExistentialConformance(proto).getValue();
-        }
-        return ProtocolConformanceRef(proto);
-      });
+  if (!UpdatedArgs)
+    return nullptr;
 
+  // Now create the new apply instruction.
   SILBuilderWithScope ApplyBuilder(Apply.getInstruction(), BuilderCtx);
   FullApplySite NewApply;
   if (auto *TAI = dyn_cast<TryApplyInst>(Apply))
@@ -920,8 +959,12 @@ SILCombiner::propagateConcreteTypeOfInitExistential(FullApplySite Apply,
     replaceWitnessMethodInst(WMI, BuilderCtx, CEI.ConcreteType,
                              SelfConformance);
   }
+  // Construct the map for Self to be used for createApplyWithConcreteType.
+  llvm::SmallDenseMap<unsigned, const ConcreteExistentialInfo *> CEIs;
+  CEIs.insert(std::pair<unsigned, const ConcreteExistentialInfo *>(
+      Apply.getCalleeArgIndex(Apply.getSelfArgumentOperand()), &CEI));
   // Try to rewrite the apply.
-  return createApplyWithConcreteType(Apply, CEI, BuilderCtx);
+  return createApplyWithConcreteType(Apply, CEIs, BuilderCtx);
 }
 
 /// Rewrite a protocol extension lookup type from an archetype to a concrete
@@ -936,27 +979,36 @@ SILCombiner::propagateConcreteTypeOfInitExistential(FullApplySite Apply,
 /// ==> apply %f<C : P>(%ref)
 SILInstruction *
 SILCombiner::propagateConcreteTypeOfInitExistential(FullApplySite Apply) {
-  // This optimization requires a generic self argument.
-  if (!Apply.hasSelfArgument() || !Apply.hasSubstitutions())
-    return nullptr;
-
-  // Try to derive the concrete type of self and a related conformance from
-  // the found init_existential.
-  const ConcreteExistentialInfo CEI(Apply.getSelfArgumentOperand());
-  if (!CEI.isValid())
+  // This optimization requires a generic argument.
+  if (!Apply.hasSubstitutions())
     return nullptr;
 
   SILBuilderContext BuilderCtx(Builder.getModule(), Builder.getTrackingList());
   SILOpenedArchetypesTracker OpenedArchetypesTracker(&Builder.getFunction());
   BuilderCtx.setOpenedArchetypesTracker(&OpenedArchetypesTracker);
-  if (CEI.ConcreteType->isOpenedExistential()) {
-    // Temporarily record this opened existential def in this local
-    // BuilderContext before rewriting the witness method.
-    OpenedArchetypesTracker.addOpenedArchetypeDef(
-        cast<ArchetypeType>(CEI.ConcreteType), CEI.ConcreteTypeDef);
+  llvm::SmallDenseMap<unsigned, const ConcreteExistentialInfo *> CEIs;
+  for (unsigned ArgIdx = 0; ArgIdx < Apply.getNumArguments(); ArgIdx++) {
+    auto ArgASTType = Apply.getArgument(ArgIdx)->getType().getASTType();
+    if (!ArgASTType->hasArchetype())
+      continue;
+    const ConcreteExistentialInfo CEI(Apply.getArgumentOperands()[ArgIdx]);
+    if (!CEI.isValid())
+      continue;
+
+    CEIs.insert(
+        std::pair<unsigned, const ConcreteExistentialInfo *>(ArgIdx, &CEI));
+
+    if (CEI.ConcreteType->isOpenedExistential()) {
+      // Temporarily record this opened existential def in this local
+      // BuilderContext before rewriting the witness method.
+      OpenedArchetypesTracker.addOpenedArchetypeDef(
+          cast<ArchetypeType>(CEI.ConcreteType), CEI.ConcreteTypeDef);
+    }
   }
-  // Perform the transformation by rewriting the apply.
-  return createApplyWithConcreteType(Apply, CEI, BuilderCtx);
+  // Bail, if no argument has a concrete existential to propagate.
+  if (CEIs.empty())
+    return nullptr;
+  return createApplyWithConcreteType(Apply, CEIs, BuilderCtx);
 }
 
 /// \brief Check that all users of the apply are retain/release ignoring one