[Serialization] Filter Decl to deserialize by their attributes

Add an alternative to getTopLevelDecls and getDeclChecked to limit which
decls are deserialized by first looking at their attributes. If the
attributes are accepted by a function passed as argument the decl is
fully deserialized, otherwise it is ignored.

The filter is included in the signature of existing functions in the
Serilalization services, but I’ve added new methods for it in FileUnit
and its subclasses to leave existing implementations untouched.

Author: xymus

include/swift/AST/FileUnit.h

@@ -152,6 +152,20 @@ 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
+  /// filter which decls to fully deserialize. Only decls with accepted
+  /// attributes are deserialized and added to Results.
+  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/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.
@@ -4181,14 +4191,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);