[embedded] Initial support for generic classes in embedded Swift

- VTableSpecializer, a new pass that synthesizes a new vtable per each observed concrete type used
- Don't use full type metadata refs in embedded Swift
- Lazily emit specialized class metadata (LazySpecializedClassMetadata) in IRGen
- Don't emit regular class metadata for a class decl if it's generic (only emit the specialized metadata)

Author: kubamracek

include/swift/IRGen/Linking.h

@@ -90,7 +90,7 @@ inline bool isEmbedded(CanType t) {
 }
 
 inline bool isMetadataAllowedInEmbedded(CanType t) {
-  return isa<ClassType>(t);
+  return isa<ClassType>(t) || isa<BoundGenericClassType>(t);
 }
 
 inline bool isEmbedded(Decl *d) {

include/swift/SIL/SILModule.h

@@ -243,6 +243,9 @@ class SILModule {
   /// Lookup table for SIL vtables from class decls.
   llvm::DenseMap<const ClassDecl *, SILVTable *> VTableMap;
 
+  /// Lookup table for specialized SIL vtables from types.
+  llvm::DenseMap<SILType, SILVTable *> SpecializedVTableMap;
+
   /// The list of SILVTables in the module.
   std::vector<SILVTable*> vtables;
 
@@ -827,6 +830,9 @@ class SILModule {
   /// Look up the VTable mapped to the given ClassDecl. Returns null on failure.
   SILVTable *lookUpVTable(const ClassDecl *C, bool deserializeLazily = true);
 
+  /// Look up a specialized VTable
+  SILVTable *lookUpSpecializedVTable(SILType classTy);
+
   /// Attempt to lookup the function corresponding to \p Member in the class
   /// hierarchy of \p Class.
   SILFunction *lookUpFunctionInVTable(ClassDecl *Class, SILDeclRef Member);

include/swift/SIL/SILVTable.h

@@ -81,6 +81,7 @@ class SILVTableEntry {
   void setNonOverridden(bool value) { IsNonOverridden = value; }
 
   SILFunction *getImplementation() const { return ImplAndKind.getPointer(); }
+  void setImplementation(SILFunction *f);
 
   void print(llvm::raw_ostream &os) const;
 
@@ -114,6 +115,8 @@ class SILVTable final : public SILAllocated<SILVTable>,
   /// The ClassDecl mapped to this VTable.
   ClassDecl *Class;
 
+  SILType classType;
+
   /// Whether or not this vtable is serialized, which allows
   /// devirtualization from another module.
   bool Serialized : 1;
@@ -122,11 +125,19 @@ class SILVTable final : public SILAllocated<SILVTable>,
   unsigned NumEntries : 31;
 
   /// Private constructor. Create SILVTables by calling SILVTable::create.
-  SILVTable(ClassDecl *c, IsSerialized_t serialized, ArrayRef<Entry> entries);
+  SILVTable(ClassDecl *c, SILType classType, IsSerialized_t serialized,
+            ArrayRef<Entry> entries);
 
-public:
+ public:
   ~SILVTable();
 
+  /// Create a new SILVTable with the given method-to-implementation mapping.
+  /// The SILDeclRef keys should reference the most-overridden members available
+  /// through the class.
+  static SILVTable *create(SILModule &M, ClassDecl *Class, SILType classType,
+                           IsSerialized_t Serialized,
+                           ArrayRef<Entry> Entries);
+
   /// Create a new SILVTable with the given method-to-implementation mapping.
   /// The SILDeclRef keys should reference the most-overridden members available
   /// through the class.
@@ -137,6 +148,11 @@ class SILVTable final : public SILAllocated<SILVTable>,
   /// Return the class that the vtable represents.
   ClassDecl *getClass() const { return Class; }
 
+  bool isSpecialized() const {
+    return !classType.isNull();
+  }
+  SILType getClassType() const { return classType; }
+
   /// Returns true if this vtable is going to be (or was) serialized.
   IsSerialized_t isSerialized() const {
     return Serialized ? IsSerialized : IsNotSerialized;

include/swift/SILOptimizer/PassManager/Passes.def

@@ -247,6 +247,8 @@ SWIFT_FUNCTION_PASS(DeadStoreElimination, "dead-store-elimination",
      "Dead Store Elimination")
 PASS(GenericSpecializer, "generic-specializer",
      "Generic Function Specialization on Static Types")
+PASS(VTableSpecializer, "vtable-specializer",
+     "Generic Specialization on VTables")
 PASS(ExistentialSpecializer, "existential-specializer",
      "Existential Specializer")
 PASS(SILSkippingChecker, "check-sil-skipping",

include/swift/SILOptimizer/Utils/Generics.h

@@ -82,15 +82,15 @@ class ReabstractionInfo {
 
   /// If set, indirect to direct conversions should be performed by the generic
   /// specializer.
-  bool ConvertIndirectToDirect;
+  bool ConvertIndirectToDirect = true;
 
   /// If true, drop metatype arguments.
   /// See `droppedMetatypeArgs`.
   bool dropMetatypeArgs = false;
 
   /// The first NumResults bits in Conversions refer to formal indirect
   /// out-parameters.
-  unsigned NumFormalIndirectResults;
+  unsigned NumFormalIndirectResults = 0;
 
   /// The function type after applying the substitutions used to call the
   /// specialized function.
@@ -101,7 +101,7 @@ class ReabstractionInfo {
   CanSILFunctionType SpecializedType;
 
   /// The generic environment to be used by the specialization.
-  GenericEnvironment *SpecializedGenericEnv;
+  GenericEnvironment *SpecializedGenericEnv = nullptr;
 
   /// The generic signature of the specialization.
   /// It is nullptr if the specialization is not polymorphic.
@@ -125,7 +125,7 @@ class ReabstractionInfo {
   SubstitutionMap ClonerParamSubMap;
 
   // Reference to the original generic non-specialized callee function.
-  SILFunction *Callee;
+  SILFunction *Callee = nullptr;
 
   // The module the specialization is created in.
   ModuleDecl *TargetModule = nullptr;
@@ -136,7 +136,7 @@ class ReabstractionInfo {
   ApplySite Apply;
 
   // Set if a specialized function has unbound generic parameters.
-  bool HasUnboundGenericParams;
+  bool HasUnboundGenericParams = false;
 
   // Substitutions to be used for creating a new function type
   // for the specialized function.
@@ -149,7 +149,7 @@ class ReabstractionInfo {
   bool isPrespecialization = false;
 
   // Is the generated specialization going to be serialized?
-  IsSerialized_t Serialized;
+  IsSerialized_t Serialized = IsNotSerialized;
 
   enum TypeCategory {
     NotLoadable,
@@ -162,7 +162,9 @@ class ReabstractionInfo {
                                            SubstitutionMap SubstMap,
                                            bool HasUnboundGenericParams);
 
+public:
   void createSubstitutedAndSpecializedTypes();
+private:
 
   TypeCategory getReturnTypeCategory(const SILResultInfo &RI,
                                      const SILFunctionConventions &substConv,
@@ -205,6 +207,12 @@ class ReabstractionInfo {
                     SILFunction *Callee, GenericSignature SpecializedSig,
                     bool isPrespecialization = false);
 
+  ReabstractionInfo(CanSILFunctionType substitutedType,
+                    SILModule &M) :
+    SubstitutedType(substitutedType),
+    isWholeModule(M.isWholeModule()) {}
+
+
   bool isPrespecialized() const { return isPrespecialization; }
 
   IsSerialized_t isSerialized() const {
@@ -400,6 +408,64 @@ class GenericFuncSpecializer {
 /// prespecialization for -Onone support.
 bool isKnownPrespecialization(StringRef SpecName);
 
+class TypeReplacements {
+private:
+  llvm::Optional<SILType> resultType;
+  llvm::MapVector<unsigned, CanType> indirectResultTypes;
+  llvm::MapVector<unsigned, CanType> paramTypeReplacements;
+  llvm::MapVector<unsigned, CanType> yieldTypeReplacements;
+
+public:
+  llvm::Optional<SILType> getResultType() const { return resultType; }
+
+  void setResultType(SILType type) { resultType = type; }
+
+  bool hasResultType() const { return resultType.has_value(); }
+
+  const llvm::MapVector<unsigned, CanType> &getIndirectResultTypes() const {
+    return indirectResultTypes;
+  }
+
+  void addIndirectResultType(unsigned index, CanType type) {
+    indirectResultTypes.insert(std::make_pair(index, type));
+  }
+
+  bool hasIndirectResultTypes() const { return !indirectResultTypes.empty(); }
+
+  const llvm::MapVector<unsigned, CanType> &getParamTypeReplacements() const {
+    return paramTypeReplacements;
+  }
+
+  void addParameterTypeReplacement(unsigned index, CanType type) {
+    paramTypeReplacements.insert(std::make_pair(index, type));
+  }
+
+  bool hasParamTypeReplacements() const {
+    return !paramTypeReplacements.empty();
+  }
+
+  const llvm::MapVector<unsigned, CanType> &getYieldTypeReplacements() const {
+    return yieldTypeReplacements;
+  }
+
+  void addYieldTypeReplacement(unsigned index, CanType type) {
+    yieldTypeReplacements.insert(std::make_pair(index, type));
+  }
+
+  bool hasYieldTypeReplacements() const {
+    return !yieldTypeReplacements.empty();
+  }
+
+  bool hasTypeReplacements() const {
+    return hasResultType() || hasParamTypeReplacements() ||
+           hasIndirectResultTypes() || hasYieldTypeReplacements();
+  }
+};
+
+ApplySite replaceWithSpecializedCallee(
+    ApplySite applySite, SILValue callee, const ReabstractionInfo &reInfo,
+    const TypeReplacements &typeReplacements = {});
+
 /// Checks if all OnoneSupport pre-specializations are included in the module
 /// as public functions.
 ///

lib/IRGen/ClassMetadataVisitor.h

@@ -63,6 +63,9 @@ template <class Impl> class ClassMetadataVisitor
     : super(IGM), Target(target),
       VTable(IGM.getSILModule().lookUpVTable(target, /*deserialize*/ false)) {}
 
+  ClassMetadataVisitor(IRGenModule &IGM, ClassDecl *target, SILVTable *vtable)
+    : super(IGM), Target(target), VTable(vtable) {}
+
 public:
   void layout() {
     static_assert(MetadataAdjustmentIndex::Class == 3,

lib/IRGen/GenClass.cpp

@@ -1079,7 +1079,9 @@ void IRGenModule::emitClassDecl(ClassDecl *D) {
     emitClassMetadata(*this, D, fragileLayout, resilientLayout);
     emitFieldDescriptor(D);
   } else {
-    emitEmbeddedClassMetadata(*this, D, fragileLayout);
+    if (!D->isGenericContext()) {
+      emitEmbeddedClassMetadata(*this, D, fragileLayout);
+    }
   }
 
   IRGen.addClassForEagerInitialization(D);

lib/IRGen/GenDecl.cpp

@@ -1427,7 +1427,8 @@ void IRGenerator::emitLazyDefinitions() {
          !LazyOpaqueTypeDescriptors.empty() || !LazyFieldDescriptors.empty() ||
          !LazyFunctionDefinitions.empty() || !LazyWitnessTables.empty() ||
          !LazyCanonicalSpecializedMetadataAccessors.empty() ||
-         !LazyMetadataAccessors.empty()) {
+         !LazyMetadataAccessors.empty() ||
+         !LazySpecializedClassMetadata.empty()) {
     // Emit any lazy type metadata we require.
     while (!LazyTypeMetadata.empty()) {
       NominalTypeDecl *type = LazyTypeMetadata.pop_back_val();
@@ -1492,8 +1493,9 @@ void IRGenerator::emitLazyDefinitions() {
     while (!LazyFunctionDefinitions.empty()) {
       SILFunction *f = LazyFunctionDefinitions.pop_back_val();
       CurrentIGMPtr IGM = getGenModule(f);
-      assert(!f->isPossiblyUsedExternally()
-             && "function with externally-visible linkage emitted lazily?");
+      // XXX TODO
+      //assert(!f->isPossiblyUsedExternally()
+      //       && "function with externally-visible linkage emitted lazily?");
       IGM->emitSILFunction(f);
     }
 
@@ -1519,6 +1521,12 @@ void IRGenerator::emitLazyDefinitions() {
       CurrentIGMPtr IGM = getGenModule(nominal->getDeclContext());
       emitLazyMetadataAccessor(*IGM.get(), nominal);
     }
+
+    while (!LazySpecializedClassMetadata.empty()) {
+      CanType classType = LazySpecializedClassMetadata.pop_back_val();
+      CurrentIGMPtr IGM = getGenModule(classType->getClassOrBoundGenericClass());
+      emitLazySpecializedClassMetadata(*IGM.get(), classType);
+    }
   }
 
   FinishedEmittingLazyDefinitions = true;
@@ -1601,6 +1609,12 @@ bool IRGenerator::hasLazyMetadata(TypeDecl *type) {
   return isLazy;
 }
 
+void IRGenerator::noteUseOfSpecializedClassMetadata(CanType classType) {
+  if (LazilyEmittedSpecializedClassMetadata.insert(classType.getPointer()).second) {
+    LazySpecializedClassMetadata.push_back(classType);
+  }
+}
+
 void IRGenerator::noteUseOfTypeGlobals(NominalTypeDecl *type,
                                        bool isUseOfMetadata,
                                        RequireMetadata_t requireMetadata) {
@@ -5061,8 +5075,9 @@ IRGenModule::getAddrOfTypeMetadata(CanType concreteType,
   // Foreign classes and prespecialized generic types do not use an alias into
   // the full metadata and therefore require a GEP.
   bool fullMetadata =
-      foreign || (concreteType->getAnyGeneric() &&
-                  concreteType->getAnyGeneric()->isGenericContext());
+      !Context.LangOpts.hasFeature(Feature::Embedded) &&
+      (foreign || (concreteType->getAnyGeneric() &&
+                  concreteType->getAnyGeneric()->isGenericContext()));
 
   llvm::Type *defaultVarTy;
   unsigned adjustmentIndex;
@@ -5099,6 +5114,11 @@ IRGenModule::getAddrOfTypeMetadata(CanType concreteType,
   // trigger lazy emission of the metadata.
   if (NominalTypeDecl *nominal = concreteType->getAnyNominal()) {
     IRGen.noteUseOfTypeMetadata(nominal);
+    if (auto *classDecl = dyn_cast<ClassDecl>(nominal)) {
+      if (classDecl->isGenericContext()) {
+        IRGen.noteUseOfSpecializedClassMetadata(concreteType);
+      }
+    }
   }
 
   if (shouldPrespecializeGenericMetadata()) {

lib/IRGen/GenMeta.cpp

@@ -3772,6 +3772,14 @@ namespace {
         FieldLayout(fieldLayout),
         MetadataLayout(IGM.getClassMetadataLayout(theClass)) {}
 
+    ClassMetadataBuilderBase(IRGenModule &IGM, ClassDecl *theClass,
+                             ConstantStructBuilder &builder,
+                             const ClassLayout &fieldLayout,
+                             SILVTable *vtable)
+      : super(IGM, theClass, vtable), B(builder),
+        FieldLayout(fieldLayout),
+        MetadataLayout(IGM.getClassMetadataLayout(theClass)) {}
+
   public:
     const ClassLayout &getFieldLayout() const { return FieldLayout; }
 
@@ -4209,6 +4217,12 @@ namespace {
                               const ClassLayout &fieldLayout)
       : super(IGM, theClass, builder, fieldLayout) {}
 
+    FixedClassMetadataBuilder(IRGenModule &IGM, ClassDecl *theClass,
+                              ConstantStructBuilder &builder,
+                              const ClassLayout &fieldLayout,
+                              SILVTable *vtable)
+      : super(IGM, theClass, builder, fieldLayout, vtable) {}
+
     void addFieldOffset(VarDecl *var) {
       if (asImpl().getFieldLayout().hasObjCImplementation())
         return;
@@ -5003,6 +5017,37 @@ void irgen::emitEmbeddedClassMetadata(IRGenModule &IGM, ClassDecl *classDecl,
   (void)var;
 }
 
+void irgen::emitLazySpecializedClassMetadata(IRGenModule &IGM,
+                                             CanType classTy) {
+  auto &context = classTy->getNominalOrBoundGenericNominal()->getASTContext();
+  PrettyStackTraceType stackTraceRAII(
+    context, "emitting lazy specialized class metadata for", classTy);
+
+  SILType classType = SILType::getPrimitiveObjectType(classTy);
+  auto &classTI = IGM.getTypeInfo(classType).as<ClassTypeInfo>();
+  
+  auto &fragileLayout =
+    classTI.getClassLayout(IGM, classType, /*forBackwardDeployment=*/true);
+
+  ClassDecl *classDecl = classType.getClassOrBoundGenericClass();
+
+  ConstantInitBuilder initBuilder(IGM);
+  auto init = initBuilder.beginStruct();
+  init.setPacked(true);
+
+  SILVTable *vtable = IGM.getSILModule().lookUpSpecializedVTable(classType);
+  assert(vtable);
+
+  FixedClassMetadataBuilder builder(IGM, classDecl, init, fragileLayout, vtable);
+  builder.layout();
+  bool canBeConstant = builder.canBeConstant();
+
+  StringRef section{};
+  auto var = IGM.defineTypeMetadata(classTy, false, canBeConstant,
+                                    init.finishAndCreateFuture(), section);
+  (void)var;
+}
+
 void irgen::emitSpecializedGenericClassMetadata(IRGenModule &IGM, CanType type,
                                                 ClassDecl &decl) {
   assert(decl.isGenericContext());

lib/IRGen/GenMeta.h

@@ -83,6 +83,8 @@ namespace irgen {
   /// Emit the type metadata accessor for a type for which it might be used.
   void emitLazyMetadataAccessor(IRGenModule &IGM, NominalTypeDecl *type);
 
+  void emitLazySpecializedClassMetadata(IRGenModule &IGM, CanType classType);
+
   void emitLazyCanonicalSpecializedMetadataAccessor(IRGenModule &IGM,
                                                     CanType theType);