Merge pull request #9109 from slavapestov/substitution-map-devirt-fix

Small devirtualizer cleanup and fix

Author: slavapestov

include/swift/AST/GenericSignature.h

@@ -201,17 +201,6 @@ class alignas(1 << TypeAlignInBits) GenericSignature final
   Optional<ProtocolConformanceRef>
   lookupConformance(CanType depTy, ProtocolDecl *proto) const;
 
-  using GenericFunction = auto(CanType canType, Type conformingReplacementType,
-    ProtocolType *conformedProtocol)
-    ->Optional<ProtocolConformanceRef>;
-  using LookupConformanceFn = llvm::function_ref<GenericFunction>;
-
-  /// Build an array of substitutions from an interface type substitution map,
-  /// using the given function to look up conformances.
-  void getSubstitutions(TypeSubstitutionFn substitution,
-                        LookupConformanceFn lookupConformance,
-                        SmallVectorImpl<Substitution> &result) const;
-
   /// Build an array of substitutions from an interface type substitution map,
   /// using the given function to look up conformances.
   void getSubstitutions(const SubstitutionMap &subMap,

lib/AST/GenericEnvironment.cpp

@@ -342,17 +342,21 @@ Type GenericEnvironment::getSugaredType(Type type) const {
 
 SubstitutionList
 GenericEnvironment::getForwardingSubstitutions() const {
+  auto *genericSig = getGenericSignature();
+
+  SubstitutionMap subMap = genericSig->getSubstitutionMap(
+    QueryInterfaceTypeSubstitutions(this),
+    MakeAbstractConformanceForGenericType());
+
   SmallVector<Substitution, 4> result;
-  getGenericSignature()->getSubstitutions(QueryInterfaceTypeSubstitutions(this),
-                                          MakeAbstractConformanceForGenericType(),
-                                          result);
-  return getGenericSignature()->getASTContext().AllocateCopy(result);
+  genericSig->getSubstitutions(subMap, result);
+  return genericSig->getASTContext().AllocateCopy(result);
 }
 
 SubstitutionMap
 GenericEnvironment::
 getSubstitutionMap(TypeSubstitutionFn subs,
-                   GenericSignature::LookupConformanceFn lookupConformance) const {
+                   LookupConformanceFn lookupConformance) const {
   SubstitutionMap subMap(const_cast<GenericEnvironment *>(this));
 
   getGenericSignature()->enumeratePairedRequirements(

lib/AST/GenericSignature.cpp

@@ -405,7 +405,7 @@ GenericSignature::getSubstitutionMap(SubstitutionList subs) const {
 SubstitutionMap
 GenericSignature::
 getSubstitutionMap(TypeSubstitutionFn subs,
-                   GenericSignature::LookupConformanceFn lookupConformance) const {
+                   LookupConformanceFn lookupConformance) const {
   SubstitutionMap subMap(const_cast<GenericSignature *>(this));
 
   // Enumerate all of the requirements that require substitution.
@@ -437,33 +437,30 @@ getSubstitutionMap(TypeSubstitutionFn subs,
 }
 
 void GenericSignature::
-getSubstitutions(TypeSubstitutionFn subs,
-                 GenericSignature::LookupConformanceFn lookupConformance,
+getSubstitutions(const SubstitutionMap &subMap,
                  SmallVectorImpl<Substitution> &result) const {
 
   // Enumerate all of the requirements that require substitution.
   enumeratePairedRequirements([&](Type depTy, ArrayRef<Requirement> reqs) {
     auto &ctx = getASTContext();
 
     // Compute the replacement type.
-    Type currentReplacement = depTy.subst(subs, lookupConformance);
+    Type currentReplacement = depTy.subst(subMap);
     if (!currentReplacement)
       currentReplacement = ErrorType::get(depTy);
 
     // Collect the conformances.
     SmallVector<ProtocolConformanceRef, 4> currentConformances;
     for (auto req: reqs) {
       assert(req.getKind() == RequirementKind::Conformance);
-      auto protoType = req.getSecondType()->castTo<ProtocolType>();
-      if (auto conformance = lookupConformance(depTy->getCanonicalType(),
-                                               currentReplacement,
-                                               protoType)) {
+      auto protoDecl = req.getSecondType()->castTo<ProtocolType>()->getDecl();
+      if (auto conformance = subMap.lookupConformance(depTy->getCanonicalType(),
+                                                      protoDecl)) {
         currentConformances.push_back(*conformance);
       } else {
         if (!currentReplacement->hasError())
           currentReplacement = ErrorType::get(currentReplacement);
-        currentConformances.push_back(
-                                  ProtocolConformanceRef(protoType->getDecl()));
+        currentConformances.push_back(ProtocolConformanceRef(protoDecl));
       }
     }
 
@@ -477,14 +474,6 @@ getSubstitutions(TypeSubstitutionFn subs,
   });
 }
 
-void GenericSignature::
-getSubstitutions(const SubstitutionMap &subMap,
-                 SmallVectorImpl<Substitution> &result) const {
-  getSubstitutions(QuerySubstitutionMap{subMap},
-                   LookUpConformanceInSubstitutionMap(subMap),
-                   result);
-}
-
 bool GenericSignature::requiresClass(Type type, ModuleDecl &mod) {
   if (!type->isTypeParameter()) return false;
 

lib/SIL/TypeLowering.cpp

@@ -2518,17 +2518,17 @@ TypeConverter::getInterfaceBoxTypeForCapture(ValueDecl *captured,
                                SILField(loweredInterfaceType, isMutable));
 
   // Instantiate the layout with identity substitutions.
-  SmallVector<Substitution, 4> genericArgs;
-  signature->getSubstitutions(
-    [&](SubstitutableType* type) -> Type {
+  auto subMap = signature->getSubstitutionMap(
+    [&](SubstitutableType *type) -> Type {
       return signature->getCanonicalTypeInContext(type,
                                                   *M.getSwiftModule());
     },
     [](Type depTy, Type replacementTy, ProtocolType *conformedTy)
     -> ProtocolConformanceRef {
       return ProtocolConformanceRef(conformedTy->getDecl());
-    },
-    genericArgs);
+    });
+  SmallVector<Substitution, 4> genericArgs;
+  signature->getSubstitutions(subMap, genericArgs);
 
   auto boxTy = SILBoxType::get(C, layout, genericArgs);
 #ifndef NDEBUG

lib/SILGen/SILGenConstructor.cpp

@@ -926,8 +926,7 @@ void SILGenFunction::emitMemberInitializers(DeclContext *dc,
           // Generate a set of substitutions for the initialization function,
           // whose generic signature is that of the type context, and whose
           // replacement types are the archetypes of the initializer itself.
-          SmallVector<Substitution, 4> subsVec;
-          typeGenericSig->getSubstitutions(
+          auto subMap = typeGenericSig->getSubstitutionMap(
                        [&](SubstitutableType *type) {
                          if (auto gp = type->getAs<GenericTypeParamType>()) {
                            return genericEnv->mapTypeIntoContext(gp);
@@ -940,8 +939,9 @@ void SILGenFunction::emitMemberInitializers(DeclContext *dc,
                            ProtocolType *conformedProtocol) {
                          return ProtocolConformanceRef(
                                   conformedProtocol->getDecl());
-                       },
-                       subsVec);
+                       });
+          SmallVector<Substitution, 4> subsVec;
+          typeGenericSig->getSubstitutions(subMap, subsVec);
           subs = SGM.getASTContext().AllocateCopy(subsVec);
         }
 

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -778,23 +778,21 @@ getSubstitutionsForProtocolConformance(ProtocolConformanceRef CRef) {
 /// \param requirementSig The generic signature of the requirement
 /// \param witnessThunkSig The generic signature of the witness method
 /// \param origSubs The substitutions from the call instruction
-/// \param newSubs New substitutions are stored here
-static void getWitnessMethodSubstitutions(
-    SILModule &M,
+static SubstitutionMap
+getWitnessMethodSubstitutions(
     ProtocolConformanceRef conformanceRef,
     GenericSignature *requirementSig,
     GenericSignature *witnessThunkSig,
     SubstitutionList origSubs,
-    bool isDefaultWitness,
-    SmallVectorImpl<Substitution> &newSubs) {
+    bool isDefaultWitness) {
 
   if (witnessThunkSig == nullptr)
-    return;
+    return SubstitutionMap();
 
-  if (isDefaultWitness) {
-    newSubs.append(origSubs.begin(), origSubs.end());
-    return;
-  }
+  auto origSubMap = requirementSig->getSubstitutionMap(origSubs);
+
+  if (isDefaultWitness)
+    return origSubMap;
 
   assert(!conformanceRef.isAbstract());
   auto conformance = conformanceRef.getConcrete();
@@ -809,24 +807,19 @@ static void getWitnessMethodSubstitutions(
     baseDepth = witnessSig->getGenericParams().back()->getDepth() + 1;
 
   auto origDepth = 1;
-  auto origSubMap = requirementSig->getSubstitutionMap(origSubs);
 
-  auto subMap =
-    SubstitutionMap::combineSubstitutionMaps(baseSubMap,
-                                             origSubMap, 
-                                             CombineSubstitutionMaps::AtDepth,
-                                             baseDepth,
-                                             origDepth,
-                                             witnessThunkSig);
-
-  witnessThunkSig->getSubstitutions(subMap, newSubs);
+  return SubstitutionMap::combineSubstitutionMaps(
+      baseSubMap,
+      origSubMap,
+      CombineSubstitutionMaps::AtDepth,
+      baseDepth,
+      origDepth,
+      witnessThunkSig);
 }
 
-static void getWitnessMethodSubstitutions(ApplySite AI, SILFunction *F,
-                                          ProtocolConformanceRef CRef,
-                                          SmallVectorImpl<Substitution> &NewSubs) {
-  auto &Module = AI.getModule();
-
+static SubstitutionMap
+getWitnessMethodSubstitutions(SILModule &Module, ApplySite AI, SILFunction *F,
+                              ProtocolConformanceRef CRef) {
   auto requirementSig = AI.getOrigCalleeType()->getGenericSignature();
   auto witnessThunkSig = F->getLoweredFunctionType()->getGenericSignature();
 
@@ -839,8 +832,9 @@ static void getWitnessMethodSubstitutions(ApplySite AI, SILFunction *F,
                                                      *Module.getSwiftModule())
       == CRef.getRequirement();
 
-  getWitnessMethodSubstitutions(Module, CRef, requirementSig, witnessThunkSig,
-                                origSubs, isDefaultWitness, NewSubs);
+  return getWitnessMethodSubstitutions(
+      CRef, requirementSig, witnessThunkSig,
+      origSubs, isDefaultWitness);
 }
 
 /// Generate a new apply of a function_ref to replace an apply of a
@@ -858,14 +852,12 @@ devirtualizeWitnessMethod(ApplySite AI, SILFunction *F,
   // The complete set of substitutions may be different, e.g. because the found
   // witness thunk F may have been created by a specialization pass and have
   // additional generic parameters.
-  SmallVector<Substitution, 4> NewSubs;
-
-  getWitnessMethodSubstitutions(AI, F, C, NewSubs);
+  auto SubMap = getWitnessMethodSubstitutions(Module, AI, F, C);
 
   // Figure out the exact bound type of the function to be called by
   // applying all substitutions.
   auto CalleeCanType = F->getLoweredFunctionType();
-  auto SubstCalleeCanType = CalleeCanType->substGenericArgs(Module, NewSubs);
+  auto SubstCalleeCanType = CalleeCanType->substGenericArgs(Module, SubMap);
 
   // Collect arguments from the apply instruction.
   auto Arguments = SmallVector<SILValue, 4>();
@@ -894,6 +886,10 @@ devirtualizeWitnessMethod(ApplySite AI, SILFunction *F,
   auto ResultSILType = substConv.getSILResultType();
   ApplySite SAI;
 
+  SmallVector<Substitution, 4> NewSubs;
+  if (auto GenericSig = CalleeCanType->getGenericSignature())
+    GenericSig->getSubstitutions(SubMap, NewSubs);
+
   SILValue ResultValue;
   if (auto *A = dyn_cast<ApplyInst>(AI)) {
     auto *NewAI =

lib/SILOptimizer/Utils/Generics.cpp

@@ -1299,9 +1299,7 @@ FunctionSignaturePartialSpecializer::
 
 void FunctionSignaturePartialSpecializer::computeClonerParamSubs(
     SubstitutionList &ClonerParamSubs) {
-  SmallVector<Substitution, 4> List;
-
-  CalleeGenericSig->getSubstitutions(
+  auto SubMap = CalleeGenericSig->getSubstitutionMap(
       [&](SubstitutableType *type) -> Type {
         DEBUG(llvm::dbgs() << "\ngetSubstitution for ClonerParamSubs:\n"
                            << Type(type) << "\n"
@@ -1312,8 +1310,10 @@ void FunctionSignaturePartialSpecializer::computeClonerParamSubs(
         return SpecializedGenericEnv->mapTypeIntoContext(
             SpecializedInterfaceTy);
       },
-      LookUpConformanceInSignature(*SpecializedGenericSig), List);
+      LookUpConformanceInSignature(*SpecializedGenericSig));
 
+  SmallVector<Substitution, 4> List;
+  CalleeGenericSig->getSubstitutions(SubMap, List);
   ClonerParamSubs = Ctx.AllocateCopy(List);
   verifySubstitutionList(ClonerParamSubs, "ClonerParamSubs");
 }

lib/Sema/CSApply.cpp

@@ -95,8 +95,11 @@ void Solution::computeSubstitutions(
                                   ConformanceCheckFlags::Used));
   };
 
-  sig->getSubstitutions(QueryTypeSubstitutionMap{subs},
-                        lookupConformanceFn, result);
+  auto subMap = sig->getSubstitutionMap(
+    QueryTypeSubstitutionMap{subs},
+    lookupConformanceFn);
+
+  sig->getSubstitutions(subMap, result);
 }
 
 void Solution::computeSubstitutions(

lib/Sema/TypeCheckType.cpp

@@ -2479,7 +2479,7 @@ Type TypeResolver::resolveSILBoxType(SILBoxTypeRepr *repr,
     }
 
     bool ok = true;
-    genericSig->getSubstitutions(
+    auto subMap = genericSig->getSubstitutionMap(
       QueryTypeSubstitutionMap{genericArgMap},
       [&](CanType depTy, Type replacement, ProtocolType *proto)
       -> ProtocolConformanceRef {
@@ -2492,17 +2492,11 @@ Type TypeResolver::resolveSILBoxType(SILBoxTypeRepr *repr,
         }
 
         return *result;
-      },
-      genericArgs);
+      });
+    genericSig->getSubstitutions(subMap, genericArgs);
 
     if (!ok)
       return ErrorType::get(Context);
-    
-    // Canonicalize the replacement types.
-    for (auto &arg : genericArgs) {
-      arg = Substitution(arg.getReplacement()->getCanonicalType(),
-                         arg.getConformances());
-    }
   }
 
   auto layout = SILLayout::get(Context, genericSig, fields);