AST: Refactor HasImplementationOnlyImportsRequest to handle querying for imports with any ImportFlag.

Author: tshortli

include/swift/AST/Import.h

@@ -94,6 +94,8 @@ enum class ImportFlags {
 /// \see ImportFlags
 using ImportOptions = OptionSet<ImportFlags>;
 
+void simple_display(llvm::raw_ostream &out, ImportOptions options);
+
 // MARK: - Import Paths
 
 namespace detail {

include/swift/AST/SourceFile.h

@@ -174,6 +174,16 @@ class SourceFile final : public FileUnit {
   ParserStatePtr DelayedParserState =
       ParserStatePtr(/*ptr*/ nullptr, /*deleter*/ nullptr);
 
+  friend class HasImportsMatchingFlagRequest;
+
+  /// Indicates which import options have a valid caches. Storage for
+  /// \c HasImportsMatchingFlagRequest.
+  ImportOptions validCachedImportOptions;
+
+  /// The cached computation of which import flags are present in the file.
+  /// Storage for \c HasImportsMatchingFlagRequest.
+  ImportOptions cachedImportOptions;
+
   friend ASTContext;
 
 public:

include/swift/AST/TypeCheckRequests.h

@@ -3202,21 +3202,24 @@ class ModuleImplicitImportsRequest
   bool isCached() const { return true; }
 };
 
-/// Checks whether a file performs an implementation-only import.
-class HasImplementationOnlyImportsRequest
-    : public SimpleRequest<HasImplementationOnlyImportsRequest,
-                           bool(SourceFile *), RequestFlags::Cached> {
+/// Checks whether a file contains any import declarations with the given flag.
+class HasImportsMatchingFlagRequest
+    : public SimpleRequest<HasImportsMatchingFlagRequest,
+                           bool(SourceFile *, ImportFlags),
+                           RequestFlags::SeparatelyCached> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
-  bool evaluate(Evaluator &evaluator, SourceFile *SF) const;
+  bool evaluate(Evaluator &evaluator, SourceFile *SF, ImportFlags flag) const;
 
 public:
   // Cached.
   bool isCached() const { return true; }
+  Optional<bool> getCachedResult() const;
+  void cacheResult(bool value) const;
 };
 
 /// Get the library level of a module.

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -169,8 +169,8 @@ SWIFT_REQUEST(TypeChecker, HasCircularInheritedProtocolsRequest,
               bool(ProtocolDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, HasCircularRawValueRequest,
               bool(EnumDecl *), Cached, NoLocationInfo)
-SWIFT_REQUEST(TypeChecker, HasImplementationOnlyImportsRequest,
-              bool(SourceFile *), Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, HasImportsMatchingFlagRequest,
+              bool(SourceFile *, ImportOptions), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ModuleLibraryLevelRequest,
               LibraryLevel(ModuleDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, InferredGenericSignatureRequest,

lib/AST/Module.cpp

@@ -2389,19 +2389,66 @@ void SourceFile::setImportUsedPreconcurrency(
   PreconcurrencyImportsUsed.insert(import);
 }
 
-bool HasImplementationOnlyImportsRequest::evaluate(Evaluator &evaluator,
-                                                   SourceFile *SF) const {
-  return llvm::any_of(SF->getImports(),
-                      [](AttributedImport<ImportedModule> desc) {
-    return desc.options.contains(ImportFlags::ImplementationOnly);
-  });
+bool HasImportsMatchingFlagRequest::evaluate(Evaluator &evaluator,
+                                             SourceFile *SF,
+                                             ImportFlags flag) const {
+  for (auto desc : SF->getImports()) {
+    if (desc.options.contains(flag))
+      return true;
+  }
+  return false;
+}
+
+Optional<bool> HasImportsMatchingFlagRequest::getCachedResult() const {
+  SourceFile *sourceFile = std::get<0>(getStorage());
+  ImportFlags flag = std::get<1>(getStorage());
+  if (sourceFile->validCachedImportOptions.contains(flag))
+    return sourceFile->cachedImportOptions.contains(flag);
+
+  return None;
+}
+
+void HasImportsMatchingFlagRequest::cacheResult(bool value) const {
+  SourceFile *sourceFile = std::get<0>(getStorage());
+  ImportFlags flag = std::get<1>(getStorage());
+
+  sourceFile->validCachedImportOptions |= flag;
+  if (value)
+    sourceFile->cachedImportOptions |= flag;
+}
+
+void swift::simple_display(llvm::raw_ostream &out, ImportOptions options) {
+  using Flag = std::pair<ImportFlags, StringRef>;
+  Flag possibleFlags[] = {
+#define FLAG(Name) {ImportFlags::Name, #Name},
+    FLAG(Exported)
+    FLAG(Testable)
+    FLAG(PrivateImport)
+    FLAG(ImplementationOnly)
+    FLAG(SPIAccessControl)
+    FLAG(Preconcurrency)
+    FLAG(WeakLinked)
+    FLAG(Reserved)
+#undef FLAG
+  };
+
+  auto flagsToPrint = llvm::make_filter_range(
+      possibleFlags, [&](Flag flag) { return options & flag.first; });
+
+  out << "{ ";
+  interleave(
+      flagsToPrint, [&](Flag flag) { out << flag.second; },
+      [&] { out << ", "; });
+  out << " }";
 }
 
 bool SourceFile::hasImplementationOnlyImports() const {
   auto &ctx = getASTContext();
   auto *mutableThis = const_cast<SourceFile *>(this);
-  return evaluateOrDefault(
-      ctx.evaluator, HasImplementationOnlyImportsRequest{mutableThis}, false);
+  return evaluateOrDefault(ctx.evaluator,
+                           HasImportsMatchingFlagRequest{
+                               mutableThis, ImportFlags::ImplementationOnly},
+                           false);
 }
 
 bool SourceFile::hasTestableOrPrivateImport(

lib/AST/TypeCheckRequests.cpp

@@ -1474,6 +1474,12 @@ void swift::simple_display(llvm::raw_ostream &out,
     out << ")";
   }
 
+  if (import.options.contains(ImportFlags::Preconcurrency))
+    out << " preconcurrency";
+
+  if (import.options.contains(ImportFlags::WeakLinked))
+    out << " weak-linked";
+
   out << " ]";
 }