Merge pull request #7938 from swiftix/wip-partial-specialization-v77

Support partial generic specialization of closures

Author: swiftix

lib/SILOptimizer/Utils/Generics.cpp

@@ -130,10 +130,10 @@ bool ReabstractionInfo::prepareAndCheck(ApplySite Apply, SILFunction *Callee,
   bool HasConcreteGenericParams = false;
   bool HasNonArchetypeGenericParams = false;
   HasUnboundGenericParams = false;
-  for (auto DT : CalleeGenericSig->getGenericParams()) {
+  for (auto GP : CalleeGenericSig->getSubstitutableParams()) {
     // Check only the substitutions for the generic parameters.
     // Ignore any dependent types, etc.
-    auto Replacement = Type(DT).subst(InterfaceSubs);
+    auto Replacement = Type(GP).subst(InterfaceSubs);
     if (!Replacement->is<ArchetypeType>())
       HasNonArchetypeGenericParams = true;
 
@@ -146,7 +146,7 @@ bool ReabstractionInfo::prepareAndCheck(ApplySite Apply, SILFunction *Callee,
       if (CalleeGenericEnv) {
         if (auto Archetype = Replacement->getAs<ArchetypeType>()) {
           auto OrigArchetype =
-              CalleeGenericEnv->mapTypeIntoContext(DT)->castTo<ArchetypeType>();
+              CalleeGenericEnv->mapTypeIntoContext(GP)->castTo<ArchetypeType>();
           if (Archetype->requiresClass() && !OrigArchetype->requiresClass())
             HasNonArchetypeGenericParams = true;
           if (Archetype->getLayoutConstraint() &&
@@ -261,6 +261,7 @@ ReabstractionInfo::ReabstractionInfo(ApplySite Apply, SILFunction *Callee,
     SpecializedType = CanSILFunctionType();
     SubstitutedType = CanSILFunctionType();
     SpecializedGenericSig = nullptr;
+    SpecializedGenericEnv = nullptr;
     return;
   }
 
@@ -382,6 +383,12 @@ ReabstractionInfo::createSubstitutedType(SILFunction *OrigF,
     Lowering::GenericContextScope GenericScope(M.Types,
                                                CanSpecializedGenericSig);
     FnTy = OrigF->getLoweredFunctionType()->substGenericArgs(M, SubstMap);
+    // FIXME: Some of the added new requirements may not have been taken into
+    // account by the substGenericArgs. So, canonicalize in the context of the
+    // specialized signature.
+    FnTy = cast<SILFunctionType>(
+        CanSpecializedGenericSig->getCanonicalTypeInContext(
+            FnTy, *M.getSwiftModule()));
   }
   assert(FnTy);
 
@@ -454,7 +461,8 @@ getGenericEnvironmentAndSignature(GenericSignatureBuilder &Builder,
   GenericSignatureBuilder TmpBuilder(
       M.getASTContext(), LookUpConformanceInModule(M.getSwiftModule()));
   TmpBuilder.addGenericSignature(GenericSig);
-  TmpBuilder.finalize(SourceLoc(), GenericSig->getGenericParams());
+  TmpBuilder.finalize(SourceLoc(), GenericSig->getGenericParams(),
+                      /*allowConcreteGenericParams=*/true);
   GenericSig =
       TmpBuilder.getGenericSignature()->getCanonicalSignature().getPointer();
   GenericEnv = GenericSig->createGenericEnvironment(*M.getSwiftModule());
@@ -482,7 +490,8 @@ getSignatureWithRequirements(GenericSignature *OrigGenSig,
     Builder.addRequirement(Req, Source);
   }
 
-  Builder.finalize(SourceLoc(), OrigGenSig->getGenericParams());
+  Builder.finalize(SourceLoc(), OrigGenSig->getGenericParams(),
+                   /*allowConcreteGenericParams=*/true);
   return getGenericEnvironmentAndSignature(Builder, M);
 }
 
@@ -861,17 +870,14 @@ static void prepareCallArguments(ApplySite AI, SILBuilder &Builder,
 
 /// Return a substituted callee function type.
 static CanSILFunctionType
-getCalleeSubstFunctionType(SILValue Callee, const ReabstractionInfo &ReInfo) {
+getCalleeSubstFunctionType(SILValue Callee, SubstitutionList Subs) {
   // Create a substituted callee type.
   auto CanFnTy =
       dyn_cast<SILFunctionType>(Callee->getType().getSwiftRValueType());
   auto CalleeSubstFnTy = CanFnTy;
 
-  if (ReInfo.getSpecializedType()->isPolymorphic() &&
-      !ReInfo.getCallerParamSubstitutions().empty()) {
-    CalleeSubstFnTy = CanFnTy->substGenericArgs(
-        ReInfo.getNonSpecializedFunction()->getModule(),
-        ReInfo.getCallerParamSubstitutions());
+  if (CanFnTy->isPolymorphic() && !Subs.empty()) {
+    CalleeSubstFnTy = CanFnTy->substGenericArgs(*Callee->getModule(), Subs);
     assert(!CalleeSubstFnTy->isPolymorphic() &&
            "Substituted callee type should not be polymorphic");
     assert(!CalleeSubstFnTy->hasTypeParameter() &&
@@ -899,7 +905,7 @@ static ApplySite replaceWithSpecializedCallee(ApplySite AI,
     Subs = ReInfo.getCallerParamSubstitutions();
   }
 
-  auto CalleeSubstFnTy = getCalleeSubstFunctionType(Callee, ReInfo);
+  auto CalleeSubstFnTy = getCalleeSubstFunctionType(Callee, Subs);
   auto CalleeSILSubstFnTy = SILType::getPrimitiveObjectType(CalleeSubstFnTy);
   SILFunctionConventions substConv(CalleeSubstFnTy, Builder.getModule());
 
@@ -988,19 +994,35 @@ class ReabstractionThunkGenerator {
       Fragile = IsFragile;
 
     {
-      Mangle::Mangler M;
-      GenericSpecializationMangler OldMangler(
-          M, OrigF, ReInfo.getOriginalParamSubstitutions(), Fragile,
-          GenericSpecializationMangler::NotReabstracted);
-      OldMangler.mangle();
-      std::string Old = M.finalize();
-
-      NewMangling::GenericSpecializationMangler NewMangler(
-        OrigF, ReInfo.getOriginalParamSubstitutions(), Fragile,
-          /*isReAbstracted*/ false);
-
-      std::string New = NewMangler.mangle();
-      ThunkName = NewMangling::selectMangling(Old, New);
+      if (!ReInfo.isPartialSpecialization()) {
+        Mangle::Mangler M;
+        GenericSpecializationMangler OldMangler(
+            M, OrigF, ReInfo.getOriginalParamSubstitutions(), Fragile,
+            GenericSpecializationMangler::NotReabstracted);
+        OldMangler.mangle();
+        std::string Old = M.finalize();
+
+        NewMangling::GenericSpecializationMangler NewMangler(
+            OrigF, ReInfo.getOriginalParamSubstitutions(), Fragile,
+            /*isReAbstracted*/ false);
+
+        std::string New = NewMangler.mangle();
+        ThunkName = NewMangling::selectMangling(Old, New);
+      } else {
+        Mangle::Mangler M;
+        PartialSpecializationMangler OldMangler(
+            M, OrigF, ReInfo.getSpecializedType(), Fragile,
+            PartialSpecializationMangler::NotReabstracted);
+        OldMangler.mangle();
+        std::string Old = M.finalize();
+
+        NewMangling::PartialSpecializationMangler NewMangler(
+            OrigF, ReInfo.getSpecializedType(), Fragile,
+            /*isReAbstracted*/ false);
+
+        std::string New = NewMangler.mangle();
+        ThunkName = NewMangling::selectMangling(Old, New);
+      }
     }
   }
 
@@ -1020,6 +1042,8 @@ SILFunction *ReabstractionThunkGenerator::createThunk() {
   if (!Thunk->empty())
     return Thunk;
 
+  Thunk->setGenericEnvironment(ReInfo.getSpecializedGenericEnvironment());
+
   SILBasicBlock *EntryBB = Thunk->createBasicBlock();
   SILBuilder Builder(EntryBB);
 
@@ -1068,16 +1092,19 @@ SILValue ReabstractionThunkGenerator::createReabstractionThunkApply(
     SILBuilder &Builder) {
   SILFunction *Thunk = &Builder.getFunction();
   auto *FRI = Builder.createFunctionRef(Loc, SpecializedFunc);
+  auto Subs = Thunk->getForwardingSubstitutions();
+  auto CalleeSubstFnTy = getCalleeSubstFunctionType(FRI, Subs);
+  auto CalleeSILSubstFnTy = SILType::getPrimitiveObjectType(CalleeSubstFnTy);
   auto specConv = SpecializedFunc->getConventions();
   if (!SpecializedFunc->getLoweredFunctionType()->hasErrorResult()) {
-    return Builder.createApply(Loc, FRI, SpecializedFunc->getLoweredType(),
-                               specConv.getSILResultType(), {}, Arguments,
+    return Builder.createApply(Loc, FRI, CalleeSILSubstFnTy,
+                               specConv.getSILResultType(), Subs, Arguments,
                                false);
   }
   // Create the logic for calling a throwing function.
   SILBasicBlock *NormalBB = Thunk->createBasicBlock();
   SILBasicBlock *ErrorBB = Thunk->createBasicBlock();
-  Builder.createTryApply(Loc, FRI, SpecializedFunc->getLoweredType(), {},
+  Builder.createTryApply(Loc, FRI, CalleeSILSubstFnTy, Subs,
                          Arguments, NormalBB, ErrorBB);
   auto *ErrorVal = ErrorBB->createPHIArgument(specConv.getSILErrorType(),
                                               ValueOwnershipKind::Owned);
@@ -1208,11 +1235,6 @@ void swift::trySpecializeApplyOfGeneric(
   bool replacePartialApplyWithoutReabstraction = false;
   auto *PAI = dyn_cast<PartialApplyInst>(Apply);
 
-  // TODO: Partial specializations of partial applies are
-  // not supported yet.
-  if (PAI && ReInfo.getSpecializedType()->isPolymorphic())
-    return;
-
   if (PAI && ReInfo.hasConversions()) {
     // If we have a partial_apply and we converted some results/parameters from
     // indirect to direct there are 3 cases:
@@ -1289,11 +1311,12 @@ void swift::trySpecializeApplyOfGeneric(
     for (auto &Op : PAI->getArgumentOperands()) {
       Arguments.push_back(Op.get());
     }
+    auto Subs = ReInfo.getCallerParamSubstitutions();
+    auto CalleeSubstFnTy = getCalleeSubstFunctionType(FRI, Subs);
+    auto CalleeSILSubstFnTy = SILType::getPrimitiveObjectType(CalleeSubstFnTy);
     auto *NewPAI = Builder.createPartialApply(PAI->getLoc(), FRI,
-                                      PAI->getSubstCalleeSILType(),
-                                      {},
-                                      Arguments,
-                                      PAI->getType());
+                                              CalleeSILSubstFnTy, Subs,
+                                              Arguments, PAI->getType());
     PAI->replaceAllUsesWith(NewPAI);
     DeadApplies.insert(PAI);
     return;