GenericSpecializer: use an alternative mangling if the function has re-abstracted resilient type parameters.

If the specialized function has a re-abstracted (= converted from indirect to direct) resilient argument or return types, use an alternative mangling: "TB" instead of "Tg".
Resilient parameters/returns can be converted from indirect to direct if the specialization is created within the type's resilience domain, i.e. in its module (where the type is loadable).
In this case we need to generate a different mangled name for the specialized function to distinguish it from specializations in other modules, which cannot re-abstract this resilient type.

This fixes a miscompile resulting from ODR-linking specializations from different modules, which in fact have different function signatures.

https://bugs.swift.org/browse/SR-13900
rdar://71914016

Author: eeckstein

include/swift/AST/TypeExpansionContext.h

@@ -89,6 +89,10 @@ class TypeExpansionContext {
       other.expansion == this->expansion;
   }
 
+  bool operator!=(const TypeExpansionContext &other) const {
+    return !operator==(other);
+  }
+
   bool operator<(const TypeExpansionContext other) const {
     assert(other.inContext != this->inContext ||
            other.isContextWholeModule == this->isContextWholeModule);

include/swift/SIL/GenericSpecializationMangler.h

@@ -89,8 +89,8 @@ class GenericSpecializationMangler : public SpecializationMangler {
                               Serialized, F) {}
 
   std::string mangleNotReabstracted(SubstitutionMap subs);
-                                    
-  std::string mangleReabstracted(SubstitutionMap subs);
+
+  std::string mangleReabstracted(SubstitutionMap subs, bool alternativeMangling);
 
   std::string mangleForDebugInfo(GenericSignature sig, SubstitutionMap subs,
                                  bool forInlining);

include/swift/SILOptimizer/Utils/Generics.h

@@ -61,6 +61,19 @@ class ReabstractionInfo {
   /// argument has a trivial type.
   SmallBitVector TrivialArgs;
 
+  /// Set to true if the function has a re-abstracted (= converted from
+  /// indirect to direct) resilient argument or return type. This can happen if
+  /// the function is compiled within the type's resilience domain, i.e. in
+  /// its module (where the type is loadable).
+  /// In this case we need to generate a different mangled name for the
+  /// function to distinguish it from functions in other modules, which cannot
+  /// re-abstract this resilient type.
+  /// Fortunately, a flag is sufficient to describe this: either a function has
+  /// re-abstracted resilient types or not. It cannot happen that two
+  /// functions have two different subsets of re-abstracted resilient parameter
+  /// types.
+  bool hasConvertedResilientParams = false;
+
   /// If set, indirect to direct conversions should be performed by the generic
   /// specializer.
   bool ConvertIndirectToDirect;
@@ -128,6 +141,12 @@ class ReabstractionInfo {
   // Is the generated specialization going to be serialized?
   IsSerialized_t Serialized;
 
+  enum TypeCategory {
+    NotLoadable,
+    Loadable,
+    LoadableAndTrivial
+  };
+  
   unsigned param2ArgIndex(unsigned ParamIdx) const  {
     return ParamIdx + NumFormalIndirectResults;
   }
@@ -138,6 +157,15 @@ class ReabstractionInfo {
                                            bool HasUnboundGenericParams);
 
   void createSubstitutedAndSpecializedTypes();
+  
+  TypeCategory getReturnTypeCategory(const SILResultInfo &RI,
+                                     const SILFunctionConventions &substConv,
+                                     TypeExpansionContext typeExpansion);
+
+  TypeCategory getParamTypeCategory(const SILParameterInfo &PI,
+                                    const SILFunctionConventions &substConv,
+                                    TypeExpansionContext typeExpansion);
+
   bool prepareAndCheck(ApplySite Apply, SILFunction *Callee,
                        SubstitutionMap ParamSubs,
                        OptRemark::Emitter *ORE = nullptr);
@@ -175,6 +203,12 @@ class ReabstractionInfo {
   IsSerialized_t isSerialized() const {
     return Serialized;
   }
+  
+  /// Returns true if the specialized function needs an alternative mangling.
+  /// See hasConvertedResilientParams.
+  bool needAlternativeMangling() const {
+    return hasConvertedResilientParams;
+  }
 
   TypeExpansionContext getResilienceExpansion() const {
     auto resilience = (Serialized ? ResilienceExpansion::Minimal

lib/SIL/Utils/GenericSpecializationMangler.cpp

@@ -103,10 +103,13 @@ mangleNotReabstracted(SubstitutionMap subs) {
 }
 
 std::string GenericSpecializationMangler::
-mangleReabstracted(SubstitutionMap subs) {
+mangleReabstracted(SubstitutionMap subs, bool alternativeMangling) {
   beginMangling();
   appendSubstitutions(getGenericSignature(), subs);
-  appendSpecializationOperator("Tg");
+  
+  // See ReabstractionInfo::hasConvertedResilientParams for why and when to use
+  // the alternative mangling.
+  appendSpecializationOperator(alternativeMangling ? "TB" : "Tg");
   return finalize();
 }
 

lib/SILOptimizer/IPO/UsePrespecialized.cpp

@@ -106,7 +106,8 @@ bool UsePrespecialized::replaceByPrespecialized(SILFunction &F) {
     // are serialized without bodies. Thus use IsNotSerialized here.
     Mangle::GenericSpecializationMangler NewGenericMangler(ReferencedF,
                                                            IsNotSerialized);
-    std::string ClonedName = NewGenericMangler.mangleReabstracted(Subs);
+    std::string ClonedName = NewGenericMangler.mangleReabstracted(Subs,
+       ReInfo.needAlternativeMangling());
 
     SILFunction *NewF = nullptr;
     // If we already have this specialization, reuse it.

lib/SILOptimizer/Utils/Generics.cpp

@@ -683,6 +683,9 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
   TrivialArgs.resize(NumArgs);
 
   SILFunctionConventions substConv(SubstitutedType, M);
+  TypeExpansionContext resilienceExp = getResilienceExpansion();
+  TypeExpansionContext minimalExp(ResilienceExpansion::Minimal,
+                                  TargetModule, isWholeModule);
 
   if (SubstitutedType->getNumDirectFormalResults() == 0) {
     // The original function has no direct result yet. Try to convert the first
@@ -693,18 +696,15 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
     for (SILResultInfo RI : SubstitutedType->getIndirectFormalResults()) {
       assert(RI.isFormalIndirect());
 
-      auto ResultTy = substConv.getSILType(RI, getResilienceExpansion());
-      ResultTy = Callee->mapTypeIntoContext(ResultTy);
-      auto &TL = M.Types.getTypeLowering(ResultTy,
-                                         getResilienceExpansion());
-
-      if (TL.isLoadable() &&
-          !RI.getReturnValueType(M, SubstitutedType, getResilienceExpansion())
-               ->isVoid() &&
-          shouldExpand(M, ResultTy)) {
+      TypeCategory tc = getReturnTypeCategory(RI, substConv, resilienceExp);
+      if (tc != NotLoadable) {
         Conversions.set(IdxForResult);
-        if (TL.isTrivial())
+        if (tc == LoadableAndTrivial)
           TrivialArgs.set(IdxForResult);
+        if (resilienceExp != minimalExp &&
+            getReturnTypeCategory(RI, substConv, minimalExp) == NotLoadable) {
+          hasConvertedResilientParams = true;
+        }
         break;
       }
       ++IdxForResult;
@@ -717,21 +717,20 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
     auto IdxToInsert = IdxForParam;
     ++IdxForParam;
 
-    auto ParamTy = substConv.getSILType(PI, getResilienceExpansion());
-    ParamTy = Callee->mapTypeIntoContext(ParamTy);
-    auto &TL = M.Types.getTypeLowering(ParamTy,
-                                       getResilienceExpansion());
-
-    if (!TL.isLoadable()) {
+    TypeCategory tc = getParamTypeCategory(PI, substConv, resilienceExp);
+    if (tc == NotLoadable)
       continue;
-    }
 
     switch (PI.getConvention()) {
     case ParameterConvention::Indirect_In:
     case ParameterConvention::Indirect_In_Guaranteed:
       Conversions.set(IdxToInsert);
-      if (TL.isTrivial())
+      if (tc == LoadableAndTrivial)
         TrivialArgs.set(IdxToInsert);
+      if (resilienceExp != minimalExp &&
+          getParamTypeCategory(PI, substConv, minimalExp) == NotLoadable) {
+        hasConvertedResilientParams = true;
+      }
       break;
     case ParameterConvention::Indirect_In_Constant:
     case ParameterConvention::Indirect_Inout:
@@ -749,6 +748,43 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
   SpecializedType = createSpecializedType(SubstitutedType, M);
 }
 
+ReabstractionInfo::TypeCategory ReabstractionInfo::
+getReturnTypeCategory(const SILResultInfo &RI,
+                  const SILFunctionConventions &substConv,
+                  TypeExpansionContext typeExpansion) {
+  auto &M = Callee->getModule();
+  auto ResultTy = substConv.getSILType(RI, typeExpansion);
+  ResultTy = Callee->mapTypeIntoContext(ResultTy);
+  auto &TL = M.Types.getTypeLowering(ResultTy, typeExpansion);
+
+  if (!TL.isLoadable())
+    return NotLoadable;
+    
+  if (RI.getReturnValueType(M, SubstitutedType, typeExpansion)
+        ->isVoid())
+    return NotLoadable;
+
+  if (!shouldExpand(M, ResultTy))
+    return NotLoadable;
+  
+  return TL.isTrivial() ? LoadableAndTrivial : Loadable;
+}
+
+ReabstractionInfo::TypeCategory ReabstractionInfo::
+getParamTypeCategory(const SILParameterInfo &PI,
+                  const SILFunctionConventions &substConv,
+                  TypeExpansionContext typeExpansion) {
+  auto &M = Callee->getModule();
+  auto ParamTy = substConv.getSILType(PI, typeExpansion);
+  ParamTy = Callee->mapTypeIntoContext(ParamTy);
+  auto &TL = M.Types.getTypeLowering(ParamTy, typeExpansion);
+
+  if (!TL.isLoadable())
+    return NotLoadable;
+    
+  return TL.isTrivial() ? LoadableAndTrivial : Loadable;
+}
+
 /// Create a new substituted type with the updated signature.
 CanSILFunctionType
 ReabstractionInfo::createSubstitutedType(SILFunction *OrigF,
@@ -1822,7 +1858,8 @@ GenericFuncSpecializer::GenericFuncSpecializer(
     if (ReInfo.isPrespecialized()) {
       ClonedName = Mangler.manglePrespecialized(ParamSubs);
     } else {
-      ClonedName = Mangler.mangleReabstracted(ParamSubs);
+      ClonedName = Mangler.mangleReabstracted(ParamSubs,
+                                              ReInfo.needAlternativeMangling());
     }
   }
   LLVM_DEBUG(llvm::dbgs() << "    Specialized function " << ClonedName << '\n');
@@ -2467,7 +2504,7 @@ usePrespecialized(SILOptFunctionBuilder &funcBuilder, ApplySite apply,
     Mangle::GenericSpecializationMangler mangler(refF, reInfo.isSerialized());
     std::string name = reInfo.isPrespecialized() ?
         mangler.manglePrespecialized(subs) :
-        mangler.mangleReabstracted(subs);
+        mangler.mangleReabstracted(subs, reInfo.needAlternativeMangling());
 
     prespecializedReInfo = reInfo;
     return lookupOrCreatePrespecialization(funcBuilder, refF, name, reInfo);

test/SILOptimizer/Inputs/specialization_and_resilience_module.swift

@@ -0,0 +1,23 @@
+public struct Mystruct {
+  var x: Int
+
+  public init(_ x: Int) { self.x = x }
+}
+
+@inline(never)
+@inlinable
+public func testParam<T>(_ t: T) {
+  print(t)
+}
+
+@inline(never)
+@inlinable
+public func testReturn<T>(_ a: [T]) -> T {
+  return a[0]
+}
+
+public func otherFunc() {
+  testParam(Mystruct(27))
+  print(testReturn([Mystruct(28)]))
+}
+

test/SILOptimizer/specialization_and_resilience.swift

@@ -1,19 +1,19 @@
-// RUN: %target-swift-frontend -parse-as-library -O -module-name=test %s -enable-library-evolution -emit-sil | %FileCheck %s
+// RUN: %empty-directory(%t) 
+// RUN: %target-build-swift -wmo -O -enable-library-evolution %S/Inputs/specialization_and_resilience_module.swift -DMODULE -parse-as-library -emit-module -emit-module-path=%t/Test.swiftmodule -module-name=Test -c -o %t/module.o
+// RUN: %target-build-swift -wmo -O %s -I%t -module-name=Main -c -o %t/main.o
+// RUN: %target-build-swift %t/main.o %t/module.o -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s
 
-public enum En {
-  case A
-  case B
-}
+// REQUIRES: executable_test
 
-@inlinable
-@inline(never)
-func genfunc<T>(_ t: T) -> T {
-  return t
-}
+import Test
+
+// CHECK: Mystruct(x: 100)
+testParam(Mystruct(100))
+// CHECK: Mystruct(x: 101)
+print(testReturn([Mystruct(101)]))
+// CHECK: Mystruct(x: 27)
+// CHECK: Mystruct(x: 28)
+otherFunc()
 
-// CHECK-LABEL: sil @$s4test11callGenFuncyyF : $@convention(thin) () -> () {
-// CHECK:  = function_ref @$s4test7genfuncyxxlFAA2EnO_Tg5 : $@convention(thin) (En) -> @out En
-// CHECK: } // end sil function '$s4test11callGenFuncyyF'
-public func callGenFunc() {
-  _ = genfunc(En.A)
-}

test/SILOptimizer/specialize_default_witness_resilience.sil

@@ -15,11 +15,11 @@ public struct ConformingStruct : ResilientProtocol {
   public func defaultB()
 }
 
-// CHECK-LABEL: sil shared @$s8defaultA4main16ConformingStructV_Tg5
+// CHECK-LABEL: sil shared @$s8defaultA4main16ConformingStructV_TB5
 // CHECK:       bb0(%0 : $ConformingStruct):
 // CHECK-NEXT:    [[TMP:%.*]] = alloc_stack $ConformingStruct
 // CHECK-NEXT:    store %0 to [[TMP]] : $*ConformingStruct
-// CHECK:         [[FN:%.*]] = function_ref @$s8defaultB4main16ConformingStructV_Tg5
+// CHECK:         [[FN:%.*]] = function_ref @$s8defaultB4main16ConformingStructV_TB5
 // CHECK-NEXT:    [[LOAD:%.*]] = load [[TMP]] : $*ConformingStruct
 // CHECK-NEXT:    [[RESULT:%.*]] = apply [[FN]]([[LOAD]])
 // CHECK-NEXT:    dealloc_stack [[TMP]] : $*ConformingStruct
@@ -32,7 +32,7 @@ bb0(%0 : $*Self):
   return %result : $()
 }
 
-// CHECK-LABEL: sil shared @$s8defaultB4main16ConformingStructV_Tg5
+// CHECK-LABEL: sil shared @$s8defaultB4main16ConformingStructV_TB5
 // CHECK:       bb0(%0 : $ConformingStruct):
 // CHECK-NEXT:    [[TMP:%.*]] = alloc_stack $ConformingStruct
 // CHECK-NEXT:    store %0 to [[TMP]] : $*ConformingStruct
@@ -48,7 +48,7 @@ bb0(%0 : $*Self):
 
 // CHECK-LABEL: sil hidden @test_specialize_default_witness_method
 // CHECK:       bb0(%0 : $*ConformingStruct):
-// CHECK:         [[FN:%.*]] = function_ref @$s8defaultA4main16ConformingStructV_Tg5
+// CHECK:         [[FN:%.*]] = function_ref @$s8defaultA4main16ConformingStructV_TB5
 // CHECK-NEXT:    [[VALUE:%.*]] = load %0 : $*ConformingStruct
 // CHECK-NEXT:    [[RESULT:%.*]] = apply [[FN]]([[VALUE]])
 // CHECK-NEXT:    return [[RESULT]]