Serialize decls with _specialize(exported:true, target: somefunc(_:), ...) in a list for fast lookup

aschwaighofer · aschwaighofer · commit 2b6ba2830ce4 · 2020-12-11T09:58:32.000-08:00
rdar://71430641
diff --git a/include/swift/AST/FileUnit.h b/include/swift/AST/FileUnit.h
@@ -183,6 +183,9 @@ class FileUnit : public DeclContext {
   /// The order of the results is not guaranteed to be meaningful.
   virtual void getTopLevelDecls(SmallVectorImpl<Decl*> &results) const {}
 
+  virtual void
+  getExportedPrespecializations(SmallVectorImpl<Decl *> &results) const {}
+
   /// Finds top-level decls in this file filtered by their attributes.
   ///
   /// This does a simple local lookup, not recursively looking through imports.
diff --git a/include/swift/AST/Module.h b/include/swift/AST/Module.h
@@ -632,6 +632,8 @@ class ModuleDecl : public DeclContext, public TypeDecl {
   /// The order of the results is not guaranteed to be meaningful.
   void getTopLevelDecls(SmallVectorImpl<Decl*> &Results) const;
 
+  void getExportedPrespecializations(SmallVectorImpl<Decl *> &results) const;
+
   /// Finds top-level decls of this module filtered by their attributes.
   ///
   /// This does a simple local lookup, not recursively looking through imports.
diff --git a/include/swift/Serialization/SerializedModuleLoader.h b/include/swift/Serialization/SerializedModuleLoader.h
@@ -398,6 +398,9 @@ class SerializedASTFile final : public LoadedFile {
 
   virtual void getTopLevelDecls(SmallVectorImpl<Decl*> &results) const override;
 
+  virtual void getExportedPrespecializations(
+      SmallVectorImpl<Decl *> &results) const override;
+
   virtual void
   getTopLevelDeclsWhereAttributesMatch(
       SmallVectorImpl<Decl*> &Results,
diff --git a/lib/AST/Module.cpp b/lib/AST/Module.cpp
@@ -792,6 +792,11 @@ void ModuleDecl::getTopLevelDecls(SmallVectorImpl<Decl*> &Results) const {
   FORWARD(getTopLevelDecls, (Results));
 }
 
+void ModuleDecl::getExportedPrespecializations(
+    SmallVectorImpl<Decl *> &Results) const {
+  FORWARD(getExportedPrespecializations, (Results));
+}
+
 void ModuleDecl::getTopLevelDeclsWhereAttributesMatch(
               SmallVectorImpl<Decl*> &Results,
               llvm::function_ref<bool(DeclAttributes)> matchAttributes) const {
diff --git a/lib/SILGen/SILGen.cpp b/lib/SILGen/SILGen.cpp
@@ -2089,37 +2089,32 @@ swift::performASTLowering(FileUnit &sf, Lowering::TypeConverter &tc,
 
 static void transferSpecializeAttributeTargets(SILGenModule &SGM, SILModule &M,
                                                Decl *d) {
-  if (auto *asd = dyn_cast<AbstractStorageDecl>(d)) {
-    for (auto ad : asd->getAllAccessors()) {
-      transferSpecializeAttributeTargets(SGM, M, ad);
+  auto *vd = cast<AbstractFunctionDecl>(d);
+  for (auto *A : vd->getAttrs().getAttributes<SpecializeAttr>()) {
+    auto *SA = cast<SpecializeAttr>(A);
+    // Filter _spi.
+    auto spiGroups = SA->getSPIGroups();
+    auto hasSPIGroup = !spiGroups.empty();
+    if (hasSPIGroup) {
+      if (vd->getModuleContext() != M.getSwiftModule() &&
+          !M.getSwiftModule()->isImportedAsSPI(SA, vd)) {
+        continue;
+      }
     }
-  } else if (auto *vd = dyn_cast<AbstractFunctionDecl>(d)) {
-    for (auto *A : vd->getAttrs().getAttributes<SpecializeAttr>()) {
-      auto *SA = cast<SpecializeAttr>(A);
-      // Filter _spi.
-      auto spiGroups = SA->getSPIGroups();
-      auto hasSPIGroup = !spiGroups.empty();
+    if (auto *targetFunctionDecl = SA->getTargetFunctionDecl(vd)) {
+      auto target = SILDeclRef(targetFunctionDecl);
+      auto targetSILFunction = SGM.getFunction(target, NotForDefinition);
+      auto kind = SA->getSpecializationKind() ==
+                          SpecializeAttr::SpecializationKind::Full
+                      ? SILSpecializeAttr::SpecializationKind::Full
+                      : SILSpecializeAttr::SpecializationKind::Partial;
+      Identifier spiGroupIdent;
       if (hasSPIGroup) {
-        if (vd->getModuleContext() != M.getSwiftModule() &&
-            !M.getSwiftModule()->isImportedAsSPI(SA, vd)) {
-          continue;
-        }
-      }
-      if (auto *targetFunctionDecl = SA->getTargetFunctionDecl(vd)) {
-        auto target = SILDeclRef(targetFunctionDecl);
-        auto targetSILFunction = SGM.getFunction(target, NotForDefinition);
-        auto kind = SA->getSpecializationKind() ==
-                            SpecializeAttr::SpecializationKind::Full
-                        ? SILSpecializeAttr::SpecializationKind::Full
-                        : SILSpecializeAttr::SpecializationKind::Partial;
-        Identifier spiGroupIdent;
-        if (hasSPIGroup) {
-          spiGroupIdent = spiGroups[0];
-        }
-        targetSILFunction->addSpecializeAttr(SILSpecializeAttr::create(
-            M, SA->getSpecializedSignature(), SA->isExported(), kind, nullptr,
-            spiGroupIdent, vd->getModuleContext()));
+        spiGroupIdent = spiGroups[0];
       }
+      targetSILFunction->addSpecializeAttr(SILSpecializeAttr::create(
+          M, SA->getSpecializedSignature(), SA->isExported(), kind, nullptr,
+          spiGroupIdent, vd->getModuleContext()));
     }
   }
 }
@@ -2143,14 +2138,11 @@ void SILGenModule::visitImportDecl(ImportDecl *import) {
   if (module->isNonSwiftModule())
     return;
 
-  SmallVector<Decl*, 16> topLevelDecls;
-  module->getTopLevelDecls(topLevelDecls);
-  for (auto *t : topLevelDecls) {
-    if (auto *vd = dyn_cast<AbstractFunctionDecl>(t)) {
+  SmallVector<Decl*, 16> prespecializations;
+  module->getExportedPrespecializations(prespecializations);
+  for (auto *p : prespecializations) {
+    if (auto *vd = dyn_cast<AbstractFunctionDecl>(p)) {
       transferSpecializeAttributeTargets(*this, M, vd);
-    } else if (auto *extension = dyn_cast<ExtensionDecl>(t)) {
-      for (auto *d : extension->getMembers())
-        transferSpecializeAttributeTargets(*this, M, d);
     }
   }
 }
diff --git a/lib/Serialization/ModuleFile.cpp b/lib/Serialization/ModuleFile.cpp
@@ -866,6 +866,20 @@ void ModuleFile::getTopLevelDecls(
   }
 }
 
+void ModuleFile::getExportedPrespecializations(
+    SmallVectorImpl<Decl *> &results) {
+  for (DeclID entry : Core->ExportedPrespecializationDecls) {
+    Expected<Decl *> declOrError = getDeclChecked(entry);
+    if (!declOrError) {
+      if (!getContext().LangOpts.EnableDeserializationRecovery)
+        fatal(declOrError.takeError());
+      consumeError(declOrError.takeError());
+      continue;
+    }
+    results.push_back(declOrError.get());
+  }
+}
+
 void ModuleFile::getOperatorDecls(SmallVectorImpl<OperatorDecl *> &results) {
   PrettyStackTraceModuleFile stackEntry(*this);
   if (!Core->OperatorDecls)
diff --git a/lib/Serialization/ModuleFile.h b/lib/Serialization/ModuleFile.h
@@ -611,6 +611,8 @@ class ModuleFile
          SmallVectorImpl<Decl*> &Results,
          llvm::function_ref<bool(DeclAttributes)> matchAttributes = nullptr);
 
+  void getExportedPrespecializations(SmallVectorImpl<Decl *> &results);
+
   /// Adds all operators to the given vector.
   void getOperatorDecls(SmallVectorImpl<OperatorDecl *> &Results);
 
diff --git a/lib/Serialization/ModuleFileSharedCore.cpp b/lib/Serialization/ModuleFileSharedCore.cpp
@@ -685,6 +685,9 @@ bool ModuleFileSharedCore::readIndexBlock(llvm::BitstreamCursor &cursor) {
       case index_block::ORDERED_TOP_LEVEL_DECLS:
         allocateBuffer(OrderedTopLevelDecls, scratch);
         break;
+      case index_block::EXPORTED_PRESPECIALIZATION_DECLS:
+        allocateBuffer(ExportedPrespecializationDecls, scratch);
+        break;
       case index_block::LOCAL_TYPE_DECLS:
         LocalTypeDecls = readLocalDeclTable(scratch, blobData);
         break;
diff --git a/lib/Serialization/ModuleFileSharedCore.h b/lib/Serialization/ModuleFileSharedCore.h
@@ -263,6 +263,7 @@ class ModuleFileSharedCore {
   std::unique_ptr<SerializedObjCMethodTable> ObjCMethods;
 
   ArrayRef<serialization::DeclID> OrderedTopLevelDecls;
+  ArrayRef<serialization::DeclID> ExportedPrespecializationDecls;
 
   class DeclCommentTableInfo;
   using SerializedDeclCommentTable =
diff --git a/lib/Serialization/ModuleFormat.h b/lib/Serialization/ModuleFormat.h
@@ -56,7 +56,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 588; // change type of get_async_continuation[_addr]
+const uint16_t SWIFTMODULE_VERSION_MINOR = 589; // cache prespecialization decls.
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -2016,7 +2016,8 @@ namespace index_block {
 
     SUBSTITUTION_MAP_OFFSETS,
     CLANG_TYPE_OFFSETS,
-    LastRecordKind = CLANG_TYPE_OFFSETS,
+    EXPORTED_PRESPECIALIZATION_DECLS,
+    LastRecordKind = EXPORTED_PRESPECIALIZATION_DECLS,
   };
 
   constexpr const unsigned RecordIDFieldWidth = 5;
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
@@ -840,6 +840,7 @@ void Serializer::writeBlockInfoBlock() {
   BLOCK_RECORD(index_block, DECL_MEMBER_NAMES);
   BLOCK_RECORD(index_block, DECL_FINGERPRINTS);
   BLOCK_RECORD(index_block, ORDERED_TOP_LEVEL_DECLS);
+  BLOCK_RECORD(index_block, EXPORTED_PRESPECIALIZATION_DECLS);
 
   BLOCK(DECL_MEMBER_TABLES_BLOCK);
   BLOCK_RECORD(decl_member_tables_block, DECL_MEMBERS);
@@ -2223,8 +2224,8 @@ static SmallVector<Type, 4> collectDependenciesFromType(Type ty) {
 class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
   Serializer &S;
   DeclID id;
+  SmallVectorImpl<DeclID> &exportedPrespecializationDecls;
   bool didVerifyAttrs = false;
-
   template <typename DeclKind>
   void verifyAttrSerializable(const DeclKind *D) {
     ::verifyAttrSerializable(D);
@@ -2921,7 +2922,10 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
   }
 
 public:
-  DeclSerializer(Serializer &S, DeclID id) : S(S), id(id) {}
+  DeclSerializer(Serializer &S, DeclID id,
+                 SmallVectorImpl<DeclID> &exportedPrespecializationDecls)
+      : S(S), id(id),
+        exportedPrespecializationDecls(exportedPrespecializationDecls) {}
   ~DeclSerializer() {
     assert(didVerifyAttrs);
   }
@@ -2934,9 +2938,24 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
     if (auto *value = dyn_cast<ValueDecl>(D))
       writeDiscriminatorsIfNeeded(value);
 
+    if (auto *afd = dyn_cast<AbstractFunctionDecl>(D)) {
+      noteUseOfExportedPrespecialization(afd);
+    }
+
     DeclVisitor<DeclSerializer>::visit(const_cast<Decl *>(D));
   }
 
+  void noteUseOfExportedPrespecialization(const AbstractFunctionDecl *afd) {
+    for (auto *A : afd->getAttrs().getAttributes<SpecializeAttr>()) {
+      auto *SA = cast<SpecializeAttr>(A);
+      if (!SA->isExported())
+        continue;
+      if (auto *targetFunctionDecl = SA->getTargetFunctionDecl(afd)) {
+        exportedPrespecializationDecls.push_back(S.addDeclRef(afd));
+      }
+    }
+  }
+
   /// If this gets referenced, we forgot to handle a decl.
   void visitDecl(const Decl *) = delete;
 
@@ -3864,7 +3883,9 @@ void Serializer::writeASTBlockEntity(const Decl *D) {
 
   assert(!D->hasClangNode() && "imported decls should use cross-references");
 
-  DeclSerializer(*this, DeclsToSerialize.addRef(D)).visit(D);
+  DeclSerializer(*this, DeclsToSerialize.addRef(D),
+                 exportedPrespecializationDecls)
+      .visit(D);
 }
 
 #define SIMPLE_CASE(TYPENAME, VALUE) \
@@ -5300,6 +5321,11 @@ void Serializer::writeAST(ModuleOrSourceFile DC) {
     OrderedDecls.emit(ScratchRecord, index_block::ORDERED_TOP_LEVEL_DECLS,
                       orderedTopLevelDecls);
 
+    index_block::OrderedDeclsLayout ExportedPrespecializationDecls(Out);
+    ExportedPrespecializationDecls.emit(
+        ScratchRecord, index_block::EXPORTED_PRESPECIALIZATION_DECLS,
+        exportedPrespecializationDecls);
+
     index_block::ObjCMethodTableLayout ObjCMethodTable(Out);
     writeObjCMethodTable(ObjCMethodTable, objcMethods);
 
diff --git a/lib/Serialization/Serialization.h b/lib/Serialization/Serialization.h
@@ -105,6 +105,8 @@ class Serializer : public SerializerBase {
   uint32_t /*IdentifierID*/ LastUniquedStringID =
       serialization::NUM_SPECIAL_IDS - 1;
 
+  SmallVector<DeclID, 16> exportedPrespecializationDecls;
+
   /// Helper for serializing entities in the AST block object graph.
   ///
   /// Keeps track of assigning IDs to newly-seen entities, and collecting
diff --git a/lib/Serialization/SerializedModuleLoader.cpp b/lib/Serialization/SerializedModuleLoader.cpp
@@ -1267,6 +1267,10 @@ SerializedASTFile::getTopLevelDecls(SmallVectorImpl<Decl*> &results) const {
   File.getTopLevelDecls(results);
 }
 
+void SerializedASTFile::getExportedPrespecializations(
+    SmallVectorImpl<Decl *> &results) const {
+  File.getExportedPrespecializations(results);
+}
 void SerializedASTFile::getTopLevelDeclsWhereAttributesMatch(
               SmallVectorImpl<Decl*> &results,
               llvm::function_ref<bool(DeclAttributes)> matchAttributes) const {
