Make sure the code completion takes both private imports and testable

imports into account separately.

Author: aschwaighofer

include/swift/AST/Module.h

@@ -990,7 +990,18 @@ class SourceFile final : public FileUnit {
 
   void addImports(ArrayRef<ImportedModuleDesc> IM);
 
-  bool hasTestableOrPrivateImport(AccessLevel accessLevel, const ValueDecl *ofDecl) const;
+  enum ImportQueryKind {
+    /// Return the resuls for testable or private imports.
+    TestableAndPrivate,
+    /// Return the results only for testable imports.
+    TestableOnly,
+    /// Return the results only for private imports.
+    PrivateOnly
+  };
+
+  bool
+  hasTestableOrPrivateImport(AccessLevel accessLevel, const ValueDecl *ofDecl,
+                             ImportQueryKind kind = TestableAndPrivate) const;
 
   void clearLookupCache();
 

include/swift/IDE/CodeCompletionCache.h

@@ -41,6 +41,7 @@ class CodeCompletionCache {
     std::vector<std::string> AccessPath;
     bool ResultsHaveLeadingDot;
     bool ForTestableLookup;
+    bool ForPrivateImportLookup;
     bool CodeCompleteInitsInPostfixExpr;
 
     friend bool operator==(const Key &LHS, const Key &RHS) {
@@ -49,6 +50,7 @@ class CodeCompletionCache {
         LHS.AccessPath == RHS.AccessPath &&
         LHS.ResultsHaveLeadingDot == RHS.ResultsHaveLeadingDot &&
         LHS.ForTestableLookup == RHS.ForTestableLookup &&
+        LHS.ForPrivateImportLookup == RHS.ForPrivateImportLookup &&
         LHS.CodeCompleteInitsInPostfixExpr == RHS.CodeCompleteInitsInPostfixExpr;
     }
   };
@@ -106,10 +108,10 @@ template<>
 struct DenseMapInfo<swift::ide::CodeCompletionCache::Key> {
   using KeyTy = swift::ide::CodeCompletionCache::Key;
   static inline KeyTy getEmptyKey() {
-    return KeyTy{"", "", {}, false, false, false};
+    return KeyTy{"", "", {}, false, false, false, false};
   }
   static inline KeyTy getTombstoneKey() {
-    return KeyTy{"", "", {}, true, false, false};
+    return KeyTy{"", "", {}, true, false, false, false};
   }
   static unsigned getHashValue(const KeyTy &Val) {
     size_t H = 0;
@@ -119,6 +121,7 @@ struct DenseMapInfo<swift::ide::CodeCompletionCache::Key> {
       H ^= std::hash<std::string>()(Piece);
     H ^= std::hash<bool>()(Val.ResultsHaveLeadingDot);
     H ^= std::hash<bool>()(Val.ForTestableLookup);
+    H ^= std::hash<bool>()(Val.ForPrivateImportLookup);
     return static_cast<unsigned>(H);
   }
   static bool isEqual(const KeyTy &LHS, const KeyTy &RHS) {

lib/AST/Module.cpp

@@ -1389,7 +1389,8 @@ void SourceFile::addImports(ArrayRef<ImportedModuleDesc> IM) {
 }
 
 bool SourceFile::hasTestableOrPrivateImport(
-    AccessLevel accessLevel, const swift::ValueDecl *ofDecl) const {
+    AccessLevel accessLevel, const swift::ValueDecl *ofDecl,
+    SourceFile::ImportQueryKind queryKind) const {
   auto *module = ofDecl->getModuleContext();
   switch (accessLevel) {
   case AccessLevel::Internal:
@@ -1399,10 +1400,19 @@ bool SourceFile::hasTestableOrPrivateImport(
     // decls).
     return std::any_of(
         Imports.begin(), Imports.end(),
-        [module](ImportedModuleDesc desc) -> bool {
-          return desc.module.second == module &&
-                 (desc.importOptions.contains(ImportFlags::PrivateImport) ||
-                  desc.importOptions.contains(ImportFlags::Testable));
+        [module, queryKind](ImportedModuleDesc desc) -> bool {
+          if (queryKind == ImportQueryKind::TestableAndPrivate)
+            return desc.module.second == module &&
+                   (desc.importOptions.contains(ImportFlags::PrivateImport) ||
+                    desc.importOptions.contains(ImportFlags::Testable));
+          else if (queryKind == ImportQueryKind::TestableOnly)
+            return desc.module.second == module &&
+                   desc.importOptions.contains(ImportFlags::Testable);
+          else {
+            assert(queryKind == ImportQueryKind::PrivateOnly);
+            return desc.module.second == module &&
+                   desc.importOptions.contains(ImportFlags::PrivateImport);
+          }
         });
   case AccessLevel::Open:
     return true;
@@ -1412,6 +1422,9 @@ bool SourceFile::hasTestableOrPrivateImport(
     break;
   }
 
+  if (queryKind == ImportQueryKind::TestableOnly)
+    return false;
+
   auto *DC = ofDecl->getDeclContext();
   if (!DC)
     return false;

lib/IDE/CodeCompletion.cpp

@@ -5946,12 +5946,15 @@ void CodeCompletionCallbacksImpl::doneParsing() {
       if (!ModuleFilename.empty()) {
         auto &Ctx = TheModule->getASTContext();
         CodeCompletionCache::Key K{
-            ModuleFilename,
-            TheModule->getName().str(),
-            AccessPath,
-            Request.NeedLeadingDot,
-            SF.hasTestableOrPrivateImport(AccessLevel::Internal, TheModule),
-            Ctx.LangOpts.CodeCompleteInitsInPostfixExpr};
+          ModuleFilename, TheModule->getName().str(), AccessPath,
+              Request.NeedLeadingDot,
+              SF.hasTestableOrPrivateImport(
+                  AccessLevel::Internal, TheModule,
+                  SourceFile::ImportQueryKind::TestableOnly),
+              SF.hasTestableOrPrivateImport(
+                  AccessLevel::Internal, TheModule,
+                  SourceFile::ImportQueryKind::PrivateOnly),
+          Ctx.LangOpts.CodeCompleteInitsInPostfixExpr};
 
         using PairType = llvm::DenseSet<swift::ide::CodeCompletionCache::Key,
             llvm::DenseMapInfo<CodeCompletionCache::Key>>::iterator;

lib/IDE/CodeCompletionCache.cpp

@@ -488,7 +488,8 @@ OnDiskCodeCompletionCache::getFromFile(StringRef filename) {
     return None;
 
   // Make up a key for readCachedModule.
-  CodeCompletionCache::Key K{filename, "<module-name>", {}, false, false, false};
+  CodeCompletionCache::Key K{filename, "<module-name>", {},   false,
+                             false,    false,           false};
 
   // Read the cached results.
   auto V = CodeCompletionCache::createValue();