SIL: Fix crash in remapParentFunction() due to missing generic signature

We need to pass down the generic signature of the caller to correctly
mangle the substitution map, because the replacement types in this
substitution maps are interface types for the caller's generic signature.

Fixes rdar://problem/161968922.

Author: slavapestov

include/swift/SIL/GenericSpecializationMangler.h

@@ -83,7 +83,9 @@ class GenericSpecializationMangler : public SpecializationMangler {
     return Function->getLoweredFunctionType()->getInvocationGenericSignature();
   }
 
-  void appendSubstitutions(GenericSignature sig, SubstitutionMap subs);
+  void appendSubstitutions(GenericSignature calleeSig,
+                           SubstitutionMap calleeSubs,
+                           GenericSignature callerSig);
 
   std::string manglePrespecialized(GenericSignature sig,
                                       SubstitutionMap subs);
@@ -112,7 +114,9 @@ class GenericSpecializationMangler : public SpecializationMangler {
   std::string mangleReabstracted(SubstitutionMap subs, bool alternativeMangling,
                                  const SmallBitVector &paramsRemoved = SmallBitVector());
 
-  std::string mangleForDebugInfo(GenericSignature sig, SubstitutionMap subs,
+  std::string mangleForDebugInfo(GenericSignature calleeSig,
+                                 SubstitutionMap calleeSubs,
+                                 GenericSignature callerSubs,
                                  bool forInlining);
 
   std::string manglePrespecialized(SubstitutionMap subs) {

include/swift/SIL/TypeSubstCloner.h

@@ -384,23 +384,40 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
   enum { ForInlining = true };
   /// Helper function to clone the parent function of a SILDebugScope if
   /// necessary when inlining said function into a new generic context.
-  /// \param SubsMap - the substitutions of the inlining/specialization process.
-  /// \param RemappedSig - the generic signature.
+  /// \param SubsMap - the substitutions for the call to the callee.
+  /// \param CalleeSig - the callee generic signature.
   template<typename FunctionBuilderTy>
   static SILFunction *remapParentFunction(FunctionBuilderTy &FuncBuilder,
                                           SILModule &M,
                                           SILFunction *ParentFunction,
                                           SubstitutionMap SubsMap,
-                                          GenericSignature RemappedSig,
+                                          GenericSignature CalleeSig,
                                           bool ForInlining = false) {
     // If the original, non-inlined version of the function had no generic
     // environment, there is no need to remap it.
-    auto *OriginalEnvironment = ParentFunction->getGenericEnvironment();
-    if (!RemappedSig || !OriginalEnvironment)
+    auto *CalleeEnv = ParentFunction->getGenericEnvironment();
+    if (!CalleeSig || !CalleeEnv)
       return ParentFunction;
 
-    if (SubsMap.getRecursiveProperties().hasPrimaryArchetype())
+    // FIXME: Pass the CallerSig down directly instead of doing this.
+    GenericSignature CallerSig;
+    if (SubsMap.getRecursiveProperties().hasPrimaryArchetype()) {
+      for (auto ReplacementType : SubsMap.getReplacementTypes()) {
+        ReplacementType.visit([&](Type t) {
+          if (t->is<PrimaryArchetypeType>() || t->is<PackArchetypeType>()) {
+            auto sig =
+              t->castTo<ArchetypeType>()->getGenericEnvironment()
+                ->getGenericSignature();
+            if (!CallerSig)
+              CallerSig = sig;
+            else
+              ASSERT(CallerSig.getPointer() == sig.getPointer());
+          }
+        });
+      }
+
       SubsMap = SubsMap.mapReplacementTypesOutOfContext();
+    }
 
     // One abstract function in the debug info can only have one set of variables
     // and types. We check if the function is called with non-identity substitutions
@@ -419,7 +436,7 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
     Mangle::GenericSpecializationMangler Mangler(M.getASTContext(), ParentFunction,
                                                  IsNotSerialized);
     std::string MangledName =
-      Mangler.mangleForDebugInfo(RemappedSig, SubsMap, ForInlining);
+      Mangler.mangleForDebugInfo(CalleeSig, SubsMap, CallerSig, ForInlining);
 
     if (ParentFunction->getName() == MangledName)
       return ParentFunction;
@@ -443,7 +460,7 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
         // undead.
         if (ParentFunction->empty()) {
           FuncBuilder.eraseFunction(ParentFunction);
-          ParentFunction->setGenericEnvironment(OriginalEnvironment);
+          ParentFunction->setGenericEnvironment(CalleeEnv);
         }
       }
     }

lib/SIL/Utils/GenericSpecializationMangler.cpp

@@ -91,14 +91,16 @@ std::string SpecializationMangler::finalize() {
 //===----------------------------------------------------------------------===//
 
 void GenericSpecializationMangler::
-appendSubstitutions(GenericSignature sig, SubstitutionMap subs) {
+appendSubstitutions(GenericSignature calleeSig,
+                    SubstitutionMap calleeSubs,
+                    GenericSignature callerSig) {
   bool First = true;
-  sig->forEachParam([&](GenericTypeParamType *ParamType, bool Canonical) {
+  calleeSig->forEachParam([&](GenericTypeParamType *ParamType, bool Canonical) {
     if (Canonical) {
       auto ty = Type(ParamType);
-      auto substTy = ty.subst(subs);
+      auto substTy = ty.subst(calleeSubs);
       auto canTy = substTy->getCanonicalType();
-      appendType(canTy, nullptr);
+      appendType(canTy, callerSig);
       appendListSeparator(First);
     }
   });
@@ -108,7 +110,7 @@ appendSubstitutions(GenericSignature sig, SubstitutionMap subs) {
 std::string GenericSpecializationMangler::
 manglePrespecialized(GenericSignature sig, SubstitutionMap subs) {
   beginMangling();
-  appendSubstitutions(sig, subs);
+  appendSubstitutions(sig, subs, GenericSignature());
   appendSpecializationOperator("Ts");
   return finalize();
 }
@@ -117,7 +119,7 @@ std::string GenericSpecializationMangler::
 mangleNotReabstracted(SubstitutionMap subs,
                       const SmallBitVector &paramsRemoved) {
   beginMangling();
-  appendSubstitutions(getGenericSignature(), subs);
+  appendSubstitutions(getGenericSignature(), subs, GenericSignature());
   appendOperator("T");
   appendRemovedParams(paramsRemoved);
   appendSpecializationOperator("G");
@@ -128,7 +130,7 @@ std::string GenericSpecializationMangler::
 mangleReabstracted(SubstitutionMap subs, bool alternativeMangling,
                    const SmallBitVector &paramsRemoved) {
   beginMangling();
-  appendSubstitutions(getGenericSignature(), subs);
+  appendSubstitutions(getGenericSignature(), subs, GenericSignature());
   appendOperator("T");
   appendRemovedParams(paramsRemoved);
 
@@ -147,9 +149,10 @@ void GenericSpecializationMangler::appendRemovedParams(const SmallBitVector &par
 }
 
 std::string GenericSpecializationMangler::
-mangleForDebugInfo(GenericSignature sig, SubstitutionMap subs, bool forInlining) {
+mangleForDebugInfo(GenericSignature calleeSig, SubstitutionMap calleeSubs,
+                   GenericSignature callerSig, bool forInlining) {
   beginMangling();
-  appendSubstitutions(sig, subs);
+  appendSubstitutions(calleeSig, calleeSubs, callerSig);
   appendSpecializationOperator(forInlining ? "Ti" : "TG");
   return finalize();
 }

test/SILOptimizer/specialize_output_generic_signature.swift

@@ -0,0 +1,17 @@
+// RUN: %target-swift-frontend -emit-sil -O %s
+
+extension Array: Comparable where Element: Comparable {
+  public static func < (lhs: Array<Element>, rhs: Array<Element>) -> Bool {
+    // Contents do not matter
+    for (l, r) in zip(lhs, rhs) { if l != r { return l < r } }
+    return lhs.count < rhs.count
+  }
+}
+
+public struct Wrapper<Symbol: Hashable & Comparable>: Hashable {
+  public var partitions: PartitionSet<Array<Symbol>>
+}
+
+public struct PartitionSet<Symbol: Hashable & Comparable>: Equatable, Hashable {
+  public var partitions: Set<Set<Symbol>>
+}