Merge pull request #28415 from xymus/filter-by-attrs

[Serialization & ClangImporter] Filter decl to deserialize by their `@objc` attributes

Author: xymus

include/swift/AST/FileUnit.h

@@ -152,6 +152,21 @@ class FileUnit : public DeclContext {
   /// The order of the results is not guaranteed to be meaningful.
   virtual void getTopLevelDecls(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.
+  /// The order of the results is not guaranteed to be meaningful.
+  ///
+  /// \param Results Vector collecting the decls.
+  ///
+  /// \param matchAttributes Check on the attributes of a decl to keep only
+  /// decls with matching attributes. The subclass SerializedASTFile checks the
+  /// attributes first to only deserialize decls with accepted attributes,
+  /// limiting deserialization work.
+  virtual void
+  getTopLevelDeclsWhereAttributesMatch(
+              SmallVectorImpl<Decl*> &Results,
+              llvm::function_ref<bool(DeclAttributes)> matchAttributes) const;
 
   /// Finds all precedence group decls in this file.
   ///

include/swift/AST/Module.h

@@ -476,6 +476,20 @@ class ModuleDecl : public DeclContext, public TypeDecl {
   /// The order of the results is not guaranteed to be meaningful.
   void getTopLevelDecls(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.
+  /// The order of the results is not guaranteed to be meaningful.
+  ///
+  /// \param Results Vector collecting the decls.
+  ///
+  /// \param matchAttributes Check on the attributes of a decl to
+  /// filter which decls to fully deserialize. Only decls with accepted
+  /// attributes are deserialized and added to Results.
+  void getTopLevelDeclsWhereAttributesMatch(
+               SmallVectorImpl<Decl*> &Results,
+               llvm::function_ref<bool(DeclAttributes)> matchAttributes) const;
+
   /// Finds all local type decls of this module.
   ///
   /// This does a simple local lookup, not recursively looking through imports.

include/swift/Serialization/SerializedModuleLoader.h

@@ -351,6 +351,11 @@ class SerializedASTFile final : public LoadedFile {
 
   virtual void getTopLevelDecls(SmallVectorImpl<Decl*> &results) const override;
 
+  virtual void
+  getTopLevelDeclsWhereAttributesMatch(
+      SmallVectorImpl<Decl*> &Results,
+      llvm::function_ref<bool(DeclAttributes)> matchAttributes) const override;
+
   virtual void
   getPrecedenceGroups(SmallVectorImpl<PrecedenceGroupDecl*> &Results) const override;
 

lib/AST/Module.cpp

@@ -644,6 +644,12 @@ void ModuleDecl::getTopLevelDecls(SmallVectorImpl<Decl*> &Results) const {
   FORWARD(getTopLevelDecls, (Results));
 }
 
+void ModuleDecl::getTopLevelDeclsWhereAttributesMatch(
+              SmallVectorImpl<Decl*> &Results,
+              llvm::function_ref<bool(DeclAttributes)> matchAttributes) const {
+  FORWARD(getTopLevelDeclsWhereAttributesMatch, (Results, matchAttributes));
+}
+
 void SourceFile::getTopLevelDecls(SmallVectorImpl<Decl*> &Results) const {
   Results.append(Decls.begin(), Decls.end());
 }
@@ -1873,6 +1879,22 @@ void *FileUnit::operator new(size_t Bytes, ASTContext &C, unsigned Alignment) {
   return C.Allocate(Bytes, Alignment);
 }
 
+void FileUnit::getTopLevelDeclsWhereAttributesMatch(
+            SmallVectorImpl<Decl*> &Results,
+            llvm::function_ref<bool(DeclAttributes)> matchAttributes) const {
+  auto prevSize = Results.size();
+  getTopLevelDecls(Results);
+
+  // Filter out unwanted decls that were just added to Results.
+  // Note: We could apply this check in all implementations of
+  // getTopLevelDecls instead or in everything that creates a Decl.
+  auto newEnd = std::remove_if(Results.begin() + prevSize, Results.end(),
+                               [&matchAttributes](const Decl *D) -> bool {
+      return !matchAttributes(D->getAttrs());
+    });
+  Results.erase(newEnd, Results.end());
+}
+
 StringRef LoadedFile::getFilename() const {
   return "";
 }

lib/ClangImporter/ImportDecl.cpp

@@ -4563,17 +4563,25 @@ namespace {
 
       if (!found) {
         // Try harder to find a match looking at just custom Objective-C names.
-        SmallVector<Decl *, 64> allTopLevelDecls;
-        overlay->getTopLevelDecls(allTopLevelDecls);
-        for (auto result : allTopLevelDecls) {
+        // Limit what we deserialize to decls with an @objc attribute.
+        SmallVector<Decl *, 4> matchingTopLevelDecls;
+
+        // Get decls with a matching @objc attribute
+        overlay->getTopLevelDeclsWhereAttributesMatch(
+          matchingTopLevelDecls,
+          [&name](const DeclAttributes attrs) -> bool {
+            if (auto objcAttr = attrs.getAttribute<ObjCAttr>())
+              if (auto objcName = objcAttr->getName())
+                return objcName->getSimpleName() == name;
+            return false;
+          });
+
+        // Filter by decl kind
+        for (auto result : matchingTopLevelDecls) {
           if (auto singleResult = dyn_cast<T>(result)) {
-            // The base name _could_ match but it's irrelevant here.
-            if (isMatch(singleResult, /*baseNameMatches=*/false,
-                        /*allowObjCMismatch=*/false)) {
-              if (found)
-                return nullptr;
-              found = singleResult;
-            }
+            if (found)
+              return nullptr;
+            found = singleResult;
           }
         }
       }

lib/Serialization/Deserialization.cpp

@@ -136,6 +136,8 @@ const char TypeError::ID = '\0';
 void TypeError::anchor() {}
 const char ExtensionError::ID = '\0';
 void ExtensionError::anchor() {}
+const char DeclAttributesDidNotMatch::ID = '\0';
+void DeclAttributesDidNotMatch::anchor() {}
 
 /// Skips a single record in the bitstream.
 ///
@@ -2280,7 +2282,8 @@ class swift::DeclDeserializer {
   /// passing each one to AddAttribute.
   llvm::Error deserializeDeclAttributes();
 
-  Expected<Decl *> getDeclCheckedImpl();
+  Expected<Decl *> getDeclCheckedImpl(
+    llvm::function_ref<bool(DeclAttributes)> matchAttributes = nullptr);
 
   Expected<Decl *> deserializeTypeAlias(ArrayRef<uint64_t> scratch,
                                         StringRef blobData) {
@@ -3855,7 +3858,9 @@ class swift::DeclDeserializer {
 };
 
 Expected<Decl *>
-ModuleFile::getDeclChecked(DeclID DID) {
+ModuleFile::getDeclChecked(
+    DeclID DID,
+    llvm::function_ref<bool(DeclAttributes)> matchAttributes) {
   if (DID == 0)
     return nullptr;
 
@@ -3868,9 +3873,14 @@ ModuleFile::getDeclChecked(DeclID DID) {
     fatalIfNotSuccess(DeclTypeCursor.JumpToBit(declOrOffset));
 
     Expected<Decl *> deserialized =
-      DeclDeserializer(*this, declOrOffset).getDeclCheckedImpl();
+      DeclDeserializer(*this, declOrOffset).getDeclCheckedImpl(
+        matchAttributes);
     if (!deserialized)
       return deserialized;
+  } else if (matchAttributes) {
+    // Decl was cached but we may need to filter it
+    if (!matchAttributes(declOrOffset.get()->getAttrs()))
+      return llvm::make_error<DeclAttributesDidNotMatch>();
   }
 
   // Tag every deserialized ValueDecl coming out of getDeclChecked with its ID.
@@ -4203,14 +4213,24 @@ llvm::Error DeclDeserializer::deserializeDeclAttributes() {
 }
 
 Expected<Decl *>
-DeclDeserializer::getDeclCheckedImpl() {
-  if (auto s = ctx.Stats)
-    s->getFrontendCounters().NumDeclsDeserialized++;
+DeclDeserializer::getDeclCheckedImpl(
+  llvm::function_ref<bool(DeclAttributes)> matchAttributes) {
 
   auto attrError = deserializeDeclAttributes();
   if (attrError)
     return std::move(attrError);
 
+  if (matchAttributes) {
+    // Deserialize the full decl only if matchAttributes finds a match.
+    DeclAttributes attrs = DeclAttributes();
+    attrs.setRawAttributeChain(DAttrs);
+    if (!matchAttributes(attrs))
+      return llvm::make_error<DeclAttributesDidNotMatch>();
+  }
+
+  if (auto s = ctx.Stats)
+    s->getFrontendCounters().NumDeclsDeserialized++;
+
   // FIXME: @_dynamicReplacement(for:) includes a reference to another decl,
   // usually in the same type, and that can result in this decl being
   // re-entrantly deserialized. If that happens, don't fail here.

lib/Serialization/DeserializationErrors.h

@@ -411,6 +411,26 @@ class SILEntityError : public llvm::ErrorInfo<SILEntityError> {
   }
 };
 
+// Decl was not deserialized because its attributes did not match the filter.
+//
+// \sa getDeclChecked
+class DeclAttributesDidNotMatch : public llvm::ErrorInfo<DeclAttributesDidNotMatch> {
+  friend ErrorInfo;
+  static const char ID;
+  void anchor() override;
+
+public:
+  DeclAttributesDidNotMatch() {}
+
+  void log(raw_ostream &OS) const override {
+    OS << "Decl attributes did not match filter";
+  }
+
+  std::error_code convertToErrorCode() const override {
+    return llvm::inconvertibleErrorCode();
+  }
+};
+
 LLVM_NODISCARD
 static inline std::unique_ptr<llvm::ErrorInfoBase>
 takeErrorInfo(llvm::Error error) {

lib/Serialization/ModuleFile.cpp

@@ -2539,11 +2539,20 @@ ModuleFile::collectLinkLibraries(ModuleDecl::LinkLibraryCallback callback) const
     callback(LinkLibrary(Name, LibraryKind::Framework));
 }
 
-void ModuleFile::getTopLevelDecls(SmallVectorImpl<Decl *> &results) {
+void ModuleFile::getTopLevelDecls(
+       SmallVectorImpl<Decl *> &results,
+       llvm::function_ref<bool(DeclAttributes)> matchAttributes) {
   PrettyStackTraceModuleFile stackEntry(*this);
   for (DeclID entry : OrderedTopLevelDecls) {
-    Expected<Decl *> declOrError = getDeclChecked(entry);
+    Expected<Decl *> declOrError = getDeclChecked(entry, matchAttributes);
     if (!declOrError) {
+      if (declOrError.errorIsA<DeclAttributesDidNotMatch>()) {
+        // Decl rejected by matchAttributes, ignore it.
+        assert(matchAttributes);
+        consumeError(declOrError.takeError());
+        continue;
+      }
+
       if (!getContext().LangOpts.EnableDeserializationRecovery)
         fatal(declOrError.takeError());
       consumeError(declOrError.takeError());

lib/Serialization/ModuleFile.h

@@ -786,7 +786,15 @@ class ModuleFile
   void collectLinkLibraries(ModuleDecl::LinkLibraryCallback callback) const;
 
   /// Adds all top-level decls to the given vector.
-  void getTopLevelDecls(SmallVectorImpl<Decl*> &Results);
+  ///
+  /// \param Results Vector collecting the decls.
+  ///
+  /// \param matchAttributes Optional check on the attributes of a decl to
+  /// filter which decls to fully deserialize. Only decls with accepted
+  /// attributes are deserialized and added to Results.
+  void getTopLevelDecls(
+         SmallVectorImpl<Decl*> &Results,
+         llvm::function_ref<bool(DeclAttributes)> matchAttributes = nullptr);
 
   /// Adds all precedence groups to the given vector.
   void getPrecedenceGroups(SmallVectorImpl<PrecedenceGroupDecl*> &Results);
@@ -888,8 +896,15 @@ class ModuleFile
   /// Returns the decl with the given ID, deserializing it if needed.
   ///
   /// \param DID The ID for the decl within this module.
+  ///
+  /// \param matchAttributes Optional check on the attributes of the decl to
+  /// determine if it should be fully deserialized and returned. If the
+  /// attributes fail the check, the decl is not deserialized and
+  /// \c DeclAttributesDidNotMatch is returned.
   llvm::Expected<Decl *>
-  getDeclChecked(serialization::DeclID DID);
+  getDeclChecked(
+    serialization::DeclID DID,
+    llvm::function_ref<bool(DeclAttributes)> matchAttributes = nullptr);
 
   /// Returns the decl context with the given ID, deserializing it if needed.
   DeclContext *getDeclContext(serialization::DeclContextID DID);

lib/Serialization/SerializedModuleLoader.cpp

@@ -1120,6 +1120,12 @@ SerializedASTFile::getTopLevelDecls(SmallVectorImpl<Decl*> &results) const {
   File.getTopLevelDecls(results);
 }
 
+void SerializedASTFile::getTopLevelDeclsWhereAttributesMatch(
+              SmallVectorImpl<Decl*> &results,
+              llvm::function_ref<bool(DeclAttributes)> matchAttributes) const {
+  File.getTopLevelDecls(results, matchAttributes);
+}
+
 void SerializedASTFile::getPrecedenceGroups(
        SmallVectorImpl<PrecedenceGroupDecl*> &results) const {
   File.getPrecedenceGroups(results);