Merge pull request #69674 from eeckstein/fix-generic-embedded-classes

Fix specialization of generic classes in embedded swift

Author: kubamracek

include/swift/SILOptimizer/Utils/Generics.h

@@ -124,9 +124,14 @@ class ReabstractionInfo {
   // callee.
   SubstitutionMap ClonerParamSubMap;
 
-  // Reference to the original generic non-specialized callee function.
+  // Reference to the original generic non-specialized callee function, if available
   SILFunction *Callee = nullptr;
 
+  // The method to specialize. This must be not null if Callee is null.
+  SILDeclRef methodDecl;
+
+  SILModule *M = nullptr;
+
   // The module the specialization is created in.
   ModuleDecl *TargetModule = nullptr;
 
@@ -186,7 +191,7 @@ class ReabstractionInfo {
       FunctionSignaturePartialSpecializer &FSPS);
 
 public:
-  ReabstractionInfo() {}
+  ReabstractionInfo(SILModule &M) : M(&M) {}
 
   /// Constructs the ReabstractionInfo for generic function \p Callee with
   /// substitutions \p ParamSubs.
@@ -208,9 +213,11 @@ class ReabstractionInfo {
                     bool isPrespecialization = false);
 
   ReabstractionInfo(CanSILFunctionType substitutedType,
+                    SILDeclRef methodDecl,
                     SILModule &M) :
     SubstitutedType(substitutedType),
-    isWholeModule(M.isWholeModule()) {}
+    methodDecl(methodDecl),
+    M(&M), isWholeModule(M.isWholeModule()) {}
 
 
   bool isPrespecialized() const { return isPrespecialization; }
@@ -326,13 +333,10 @@ class ReabstractionInfo {
 
   SILFunction *getNonSpecializedFunction() const { return Callee; }
 
-  /// Map type into a context of the specialized function.
-  Type mapTypeIntoContext(Type type) const;
-
   /// Map SIL type into a context of the specialized function.
   SILType mapTypeIntoContext(SILType type) const;
 
-  SILModule &getModule() const { return Callee->getModule(); }
+  SILModule &getModule() const { return *M; }
 
   /// Returns true if generic specialization is possible.
   bool canBeSpecialized() const;

lib/IRGen/ClassMetadataVisitor.h

@@ -67,6 +67,18 @@ template <class Impl> class ClassMetadataVisitor
     : super(IGM), Target(target), VTable(vtable) {}
 
 public:
+
+  // Layout in embedded mode while considering the class type.
+  // This is important for adding the right superclass pointer.
+  // The regular `layout` method can be used for layout tasks for which the
+  // actual superclass pointer is not relevant.
+  void layoutEmbedded(CanType classTy) {
+    asImpl().noteAddressPoint();
+    asImpl().addEmbeddedSuperclass(classTy);
+    asImpl().addDestructorFunction();
+    addEmbeddedClassMembers(Target);
+  }
+
   void layout() {
     static_assert(MetadataAdjustmentIndex::Class == 3,
                   "Adjustment index must be synchronized with this layout");

lib/IRGen/GenMeta.cpp

@@ -3903,6 +3903,16 @@ namespace {
       llvm_unreachable("covered switch");
     }
 
+    void addEmbeddedSuperclass(CanType classTy) {
+      CanType superclass = asImpl().getSuperclassTypeForMetadata();
+      if (!superclass) {
+        B.addNullPointer(IGM.TypeMetadataPtrTy);
+        return;
+      }
+      CanType superTy = classTy->getSuperclass()->getCanonicalType();
+      B.add(IGM.getAddrOfTypeMetadata(superTy));
+    }
+
     void addSuperclass() {
       if (asImpl().shouldAddNullSuperclass()) {
         B.addNullPointer(IGM.TypeMetadataPtrTy);
@@ -5025,7 +5035,7 @@ static void emitEmbeddedVTable(IRGenModule &IGM, CanType classTy,
 
   FixedClassMetadataBuilder builder(IGM, classDecl, init, fragileLayout,
                                     vtable);
-  builder.layout();
+  builder.layoutEmbedded(classTy);
   bool canBeConstant = builder.canBeConstant();
 
   StringRef section{};

lib/SIL/Verifier/SILVerifier.cpp

@@ -3857,7 +3857,11 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     auto member = CMI->getMember();
     auto overrideTy =
         TC.getConstantOverrideType(F.getTypeExpansionContext(), member);
-    if (CMI->getModule().getStage() != SILStage::Lowered) {
+
+    SILModule &mod = CMI->getModule();
+    bool embedded = mod.getASTContext().LangOpts.hasFeature(Feature::Embedded);
+
+    if (mod.getStage() != SILStage::Lowered && !embedded) {
       requireSameType(
           CMI->getType(), SILType::getPrimitiveObjectType(overrideTy),
           "result type of class_method must match abstracted type of method");

lib/SILOptimizer/Transforms/EagerSpecializer.cpp

@@ -857,7 +857,7 @@ void EagerSpecializerTransform::run() {
       // That means we are loading it from another module. In this case, we
       // don't want to create a pre-specialization.
       SpecializedFuncs.push_back(nullptr);
-      ReInfoVec.emplace_back(ReabstractionInfo());
+      ReInfoVec.emplace_back(ReabstractionInfo(F.getModule()));
       continue;
     }
     ReInfoVec.emplace_back(FuncBuilder.getModule().getSwiftModule(),

lib/SILOptimizer/Transforms/VTableSpecializer.cpp

@@ -204,7 +204,7 @@ bool swift::specializeClassMethodInst(ClassMethodInst *cm) {
   SILType substitutedType =
       funcTy.substGenericArgs(m, subs, TypeExpansionContext::minimal());
 
-  ReabstractionInfo reInfo(substitutedType.getAs<SILFunctionType>(), m);
+  ReabstractionInfo reInfo(substitutedType.getAs<SILFunctionType>(), cm->getMember(), m);
   reInfo.createSubstitutedAndSpecializedTypes();
   CanSILFunctionType finalFuncTy = reInfo.getSpecializedType();
   SILType finalSILTy = SILType::getPrimitiveObjectType(finalFuncTy);

lib/SILOptimizer/Utils/Generics.cpp

@@ -649,7 +649,7 @@ bool ReabstractionInfo::prepareAndCheck(ApplySite Apply, SILFunction *Callee,
 
 bool ReabstractionInfo::canBeSpecialized(ApplySite Apply, SILFunction *Callee,
                                          SubstitutionMap ParamSubs) {
-  ReabstractionInfo ReInfo;
+  ReabstractionInfo ReInfo(Callee->getModule());
   return ReInfo.prepareAndCheck(Apply, Callee, ParamSubs);
 }
 
@@ -659,6 +659,7 @@ ReabstractionInfo::ReabstractionInfo(
     bool ConvertIndirectToDirect, bool dropMetatypeArgs, OptRemark::Emitter *ORE)
     : ConvertIndirectToDirect(ConvertIndirectToDirect),
       dropMetatypeArgs(dropMetatypeArgs),
+      M(&Callee->getModule()),
       TargetModule(targetModule), isWholeModule(isWholeModule),
       Serialized(Serialized) {
   if (!prepareAndCheck(Apply, Callee, ParamSubs, ORE))
@@ -772,8 +773,6 @@ bool ReabstractionInfo::isPartialSpecialization() const {
 }
 
 void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
-  auto &M = Callee->getModule();
-
   // Find out how the function type looks like after applying the provided
   // substitutions.
   if (!SubstitutedType) {
@@ -792,7 +791,7 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
   TrivialArgs.resize(NumArgs);
   droppedMetatypeArgs.resize(NumArgs);
 
-  SILFunctionConventions substConv(SubstitutedType, M);
+  SILFunctionConventions substConv(SubstitutedType, getModule());
   TypeExpansionContext resilienceExp = getResilienceExpansion();
   TypeExpansionContext minimalExp(ResilienceExpansion::Minimal,
                                   TargetModule, isWholeModule);
@@ -863,26 +862,25 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
   // Produce a specialized type, which is the substituted type with
   // the parameters/results passing conventions adjusted according
   // to the conversions selected above.
-  SpecializedType = createSpecializedType(SubstitutedType, M);
+  SpecializedType = createSpecializedType(SubstitutedType, getModule());
 }
 
 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);
+  auto &TL = getModule().Types.getTypeLowering(ResultTy, typeExpansion);
 
   if (!TL.isLoadable())
     return NotLoadable;
 
-  if (RI.getReturnValueType(M, SubstitutedType, typeExpansion)
+  if (RI.getReturnValueType(getModule(), SubstitutedType, typeExpansion)
         ->isVoid())
     return NotLoadable;
 
-  if (!shouldExpand(M, ResultTy))
+  if (!shouldExpand(getModule(), ResultTy))
     return NotLoadable;
 
   return TL.isTrivial() ? LoadableAndTrivial : Loadable;
@@ -892,10 +890,9 @@ 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);
+  ParamTy = mapTypeIntoContext(ParamTy);
+  auto &TL = getModule().Types.getTypeLowering(ParamTy, typeExpansion);
 
   if (!TL.isLoadable())
     return NotLoadable;
@@ -908,7 +905,6 @@ CanSILFunctionType
 ReabstractionInfo::createSubstitutedType(SILFunction *OrigF,
                                          SubstitutionMap SubstMap,
                                          bool HasUnboundGenericParams) {
-  auto &M = OrigF->getModule();
   if ((SpecializedGenericSig &&
        SpecializedGenericSig->areAllParamsConcrete()) ||
       !HasUnboundGenericParams) {
@@ -920,8 +916,8 @@ ReabstractionInfo::createSubstitutedType(SILFunction *OrigF,
 
   auto lowered = OrigF->getLoweredFunctionType();
   auto genSub =
-      lowered->substGenericArgs(M, SubstMap, getResilienceExpansion());
-  auto unsub = genSub->getUnsubstitutedType(M);
+      lowered->substGenericArgs(getModule(), SubstMap, getResilienceExpansion());
+  auto unsub = genSub->getUnsubstitutedType(getModule());
   auto specialized = CanSpecializedGenericSig.getReducedType(unsub);
 
   // First substitute concrete types into the existing function type.
@@ -935,7 +931,7 @@ ReabstractionInfo::createSubstitutedType(SILFunction *OrigF,
       CanSpecializedGenericSig, FnTy->getExtInfo(), FnTy->getCoroutineKind(),
       FnTy->getCalleeConvention(), FnTy->getParameters(), FnTy->getYields(),
       FnTy->getResults(), FnTy->getOptionalErrorResult(),
-      FnTy->getPatternSubstitutions(), SubstitutionMap(), M.getASTContext(),
+      FnTy->getPatternSubstitutions(), SubstitutionMap(), getModule().getASTContext(),
       FnTy->getWitnessMethodConformanceOrInvalid());
 
   // This is an interface type. It should not have any archetypes.
@@ -972,6 +968,16 @@ CanSILFunctionType ReabstractionInfo::createThunkType(PartialApplyInst *forPAI)
   return newFnTy;
 }
 
+SILType ReabstractionInfo::mapTypeIntoContext(SILType type) const {
+  if (Callee) {
+    return Callee->mapTypeIntoContext(type);
+  }
+  assert(!methodDecl.isNull());
+  if (auto *genericEnv = M->Types.getConstantGenericEnvironment(methodDecl))
+    return genericEnv->mapTypeIntoContext(getModule(), type);
+  return type;
+}
+
 /// Convert the substituted function type into a specialized function type based
 /// on the ReabstractionInfo.
 CanSILFunctionType ReabstractionInfo::
@@ -1081,15 +1087,13 @@ void ReabstractionInfo::performFullSpecializationPreparation(
   assert((!EnablePartialSpecialization || !HasUnboundGenericParams) &&
          "Only full specializations are handled here");
 
-  SILModule &M = Callee->getModule();
-
   this->Callee = Callee;
 
   // Get the original substitution map.
   ClonerParamSubMap = ParamSubs;
 
   SubstitutedType = Callee->getLoweredFunctionType()->substGenericArgs(
-      M, ClonerParamSubMap, getResilienceExpansion());
+      getModule(), ClonerParamSubMap, getResilienceExpansion());
   CallerParamSubMap = {};
   createSubstitutedAndSpecializedTypes();
 }
@@ -1889,8 +1893,6 @@ void FunctionSignaturePartialSpecializer::createSpecializedGenericSignature(
 void ReabstractionInfo::performPartialSpecializationPreparation(
     SILFunction *Caller, SILFunction *Callee,
     SubstitutionMap ParamSubs) {
-  SILModule &M = Callee->getModule();
-
   // Caller is the SILFunction containing the apply instruction.
   CanGenericSignature CallerGenericSig;
   GenericEnvironment *CallerGenericEnv = nullptr;
@@ -1910,7 +1912,7 @@ void ReabstractionInfo::performPartialSpecializationPreparation(
              llvm::dbgs() << "Callee generic signature is:\n";
              CalleeGenericSig->print(llvm::dbgs()));
 
-  FunctionSignaturePartialSpecializer FSPS(M,
+  FunctionSignaturePartialSpecializer FSPS(getModule(),
                                            CallerGenericSig, CallerGenericEnv,
                                            CalleeGenericSig, CalleeGenericEnv,
                                            ParamSubs);
@@ -1968,7 +1970,7 @@ ReabstractionInfo::ReabstractionInfo(ModuleDecl *targetModule,
                                      bool isWholeModule, SILFunction *Callee,
                                      GenericSignature SpecializedSig,
                                      bool isPrespecialization)
-    : TargetModule(targetModule), isWholeModule(isWholeModule),
+    : M(&Callee->getModule()), TargetModule(targetModule), isWholeModule(isWholeModule),
       isPrespecialization(isPrespecialization) {
   Serialized =
       this->isPrespecialization ? IsNotSerialized : Callee->isSerialized();
@@ -1979,13 +1981,11 @@ ReabstractionInfo::ReabstractionInfo(ModuleDecl *targetModule,
   this->Callee = Callee;
   ConvertIndirectToDirect = true;
 
-  SILModule &M = Callee->getModule();
-
   auto CalleeGenericSig =
       Callee->getLoweredFunctionType()->getInvocationGenericSignature();
   auto *CalleeGenericEnv = Callee->getGenericEnvironment();
 
-  FunctionSignaturePartialSpecializer FSPS(M,
+  FunctionSignaturePartialSpecializer FSPS(getModule(),
                                            CalleeGenericSig, CalleeGenericEnv,
                                            SpecializedSig);
 
@@ -3091,7 +3091,7 @@ void swift::trySpecializeApplyOfGeneric(
 
   // Check if there is a pre-specialization available in a library.
   SpecializedFunction prespecializedF{};
-  ReabstractionInfo prespecializedReInfo;
+  ReabstractionInfo prespecializedReInfo(FuncBuilder.getModule());
   bool replacePartialApplyWithoutReabstraction = false;
 
   if (usePrespecialized(FuncBuilder, Apply, RefF, ReInfo, prespecializedReInfo,

test/embedded/generic-classes.swift

@@ -0,0 +1,36 @@
+// RUN: %target-run-simple-swift(%S/Inputs/print.swift -enable-experimental-feature Embedded -parse-as-library -runtime-compatibility-version none -wmo -Xfrontend -disable-objc-interop) | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: optimized_stdlib
+// REQUIRES: VENDOR=apple
+// REQUIRES: OS=macosx
+
+protocol Fooable {
+  func foo()
+}
+
+class GenericFooableClass<T>: Fooable {
+  func foo() { print("GenericFooableClass<T>.foo") }
+}
+
+class GenericFooableSubClass<T>: GenericFooableClass<T> {
+  override func foo() { print("GenericFooableSubClass<T>.foo") }
+}
+
+func makeItFoo<F: Fooable>(f: F) {
+  f.foo()
+}
+
+@main
+struct Main {
+  static func main() {
+    let f = GenericFooableClass<Int>()
+    makeItFoo(f: f)
+    let g: GenericFooableClass = GenericFooableSubClass<Int>()
+    makeItFoo(f: g)
+  }
+}
+
+// CHECK: GenericFooableClass<T>.foo
+// CHECK: GenericFooableSubClass<T>.foo
+