[CodeComplete] Show global completions from modules that are imported as @_spi

ahoppen · ahoppen · commit a296e61dec7e · 2022-08-26T19:31:55.000+02:00
If a module is imported as `@_spi`, we didn’t receive any global completions from it.

rdar://99027179
diff --git a/include/swift/IDE/CodeCompletionCache.h b/include/swift/IDE/CodeCompletionCache.h
@@ -45,6 +45,8 @@ class CodeCompletionCache {
     bool ResultsHaveLeadingDot;
     bool ForTestableLookup;
     bool ForPrivateImportLookup;
+    /// Must be sorted alphabetically for stable identity.
+    llvm::SmallVector<std::string, 2> SpiGroups;
     bool AddInitsInToplevel;
     bool AddCallWithNoDefaultArgs;
     bool Annotated;
@@ -56,6 +58,7 @@ class CodeCompletionCache {
              LHS.ResultsHaveLeadingDot == RHS.ResultsHaveLeadingDot &&
              LHS.ForTestableLookup == RHS.ForTestableLookup &&
              LHS.ForPrivateImportLookup == RHS.ForPrivateImportLookup &&
+             LHS.SpiGroups == RHS.SpiGroups &&
              LHS.AddInitsInToplevel == RHS.AddInitsInToplevel &&
              LHS.AddCallWithNoDefaultArgs == RHS.AddCallWithNoDefaultArgs &&
              LHS.Annotated == RHS.Annotated;
@@ -125,16 +128,35 @@ template<>
 struct DenseMapInfo<swift::ide::CodeCompletionCache::Key> {
   using KeyTy = swift::ide::CodeCompletionCache::Key;
   static inline KeyTy getEmptyKey() {
-    return KeyTy{"", "", {}, false, false, false, false, false, false};
+    return KeyTy{/*ModuleFilename=*/"",
+                 /*ModuleName=*/"",
+                 /*AccessPath=*/{},
+                 /*ResultsHaveLeadingDot=*/false,
+                 /*ForTestableLookup=*/false,
+                 /*ForPrivateImportLookup=*/false,
+                 /*SpiGroups=*/{},
+                 /*AddInitsInToplevel=*/false,
+                 /*AddCallWithNoDefaultArgs=*/false,
+                 /*Annotated=*/false};
   }
   static inline KeyTy getTombstoneKey() {
-    return KeyTy{"", "", {}, true, false, false, false, false, false};
+    return KeyTy{/*ModuleFilename=*/"",
+                 /*ModuleName=*/"",
+                 /*AccessPath=*/{},
+                 /*ResultsHaveLeadingDot=*/true,
+                 /*ForTestableLookup=*/false,
+                 /*ForPrivateImportLookup=*/false,
+                 /*SpiGroups=*/{},
+                 /*AddInitsInToplevel=*/false,
+                 /*AddCallWithNoDefaultArgs=*/false,
+                 /*Annotated=*/false};
   }
   static unsigned getHashValue(const KeyTy &Val) {
     return llvm::hash_combine(
         Val.ModuleFilename, Val.ModuleName,
         llvm::hash_combine_range(Val.AccessPath.begin(), Val.AccessPath.end()),
         Val.ResultsHaveLeadingDot, Val.ForTestableLookup,
+        llvm::hash_combine_range(Val.SpiGroups.begin(), Val.SpiGroups.end()),
         Val.ForPrivateImportLookup, Val.AddInitsInToplevel,
         Val.AddCallWithNoDefaultArgs, Val.Annotated);
   }
diff --git a/lib/IDE/CodeCompletion.cpp b/lib/IDE/CodeCompletion.cpp
@@ -1291,6 +1291,16 @@ void swift::ide::deliverCompletionResults(
       // ModuleFilename can be empty if something strange happened during
       // module loading, for example, the module file is corrupted.
       if (!ModuleFilename.empty()) {
+        llvm::SmallVector<std::string, 2> spiGroups;
+        for (auto Import : SF.getImports()) {
+          if (Import.module.importedModule == TheModule) {
+            for (auto SpiGroup : Import.spiGroups) {
+              spiGroups.push_back(SpiGroup.str().str());
+            }
+            break;
+          }
+        }
+        llvm::sort(spiGroups);
         CodeCompletionCache::Key K{
             ModuleFilename.str(),
             std::string(TheModule->getName()),
@@ -1302,6 +1312,7 @@ void swift::ide::deliverCompletionResults(
             SF.hasTestableOrPrivateImport(
                 AccessLevel::Internal, TheModule,
                 SourceFile::ImportQueryKind::PrivateOnly),
+            spiGroups,
             CompletionContext.getAddInitsToTopLevel(),
             CompletionContext.addCallWithNoDefaultArgs(),
             CompletionContext.getAnnotateResult()};
@@ -1341,9 +1352,12 @@ void swift::ide::deliverCompletionResults(
       // Add results for all imported modules.
       SmallVector<ImportedModule, 4> Imports;
       SF.getImportedModules(
-          Imports, {ModuleDecl::ImportFilterKind::Exported,
-                    ModuleDecl::ImportFilterKind::Default,
-                    ModuleDecl::ImportFilterKind::ImplementationOnly});
+          Imports, {
+                       ModuleDecl::ImportFilterKind::Exported,
+                       ModuleDecl::ImportFilterKind::Default,
+                       ModuleDecl::ImportFilterKind::ImplementationOnly,
+                       ModuleDecl::ImportFilterKind::SPIAccessControl,
+                   });
 
       for (auto Imported : Imports) {
         for (auto Import : namelookup::getAllImports(Imported.importedModule))
diff --git a/lib/IDE/CodeCompletionCache.cpp b/lib/IDE/CodeCompletionCache.cpp
@@ -475,6 +475,12 @@ static std::string getName(StringRef cacheDirectory,
       << (K.AddInitsInToplevel ? "-inits" : "")
       << (K.AddCallWithNoDefaultArgs ? "-nodefaults" : "")
       << (K.Annotated ? "-annotated" : "");
+  if (K.SpiGroups.size() > 0) {
+    OSS << "-spi";
+    for (auto SpiGroup : K.SpiGroups) {
+      OSS << "-" << SpiGroup;
+    }
+  }
 
   // name[-access-path-components]
   for (StringRef component : K.AccessPath)
@@ -539,9 +545,16 @@ OnDiskCodeCompletionCache::getFromFile(StringRef filename) {
     return None;
 
   // Make up a key for readCachedModule.
-  CodeCompletionCache::Key K{filename.str(), "<module-name>", {},
-                             false,          false,           false,
-                             false,          false,           false};
+  CodeCompletionCache::Key K{/*ModuleFilename=*/filename.str(),
+                             /*ModuleName=*/"<module-name>",
+                             /*AccessPath=*/{},
+                             /*ResultsHaveLeadingDot=*/false,
+                             /*ForTestableLookup=*/false,
+                             /*ForPrivateImportLookup=*/false,
+                             /*SpiGroups=*/{},
+                             /*AddInitsInToplevel=*/false,
+                             /*AddCallWithNoDefaultArgs=*/false,
+                             /*Annotated=*/false};
 
   // Read the cached results.
   auto V = CodeCompletionCache::createValue();
diff --git a/test/IDE/complete_from_swift_module_spi.swift b/test/IDE/complete_from_swift_module_spi.swift
@@ -0,0 +1,61 @@
+// RUN: %empty-directory(%t/split)
+// RUN: %empty-directory(%t/build)
+// RUN: %{python} %utils/split_file.py -o %t/split %s
+
+// RUN: %target-swift-frontend -emit-module -o %t/build %t/split/pck.swift
+
+// First SPI completion then completion in file without SPI import
+// RUN: %empty-directory(%t/cc-cache)
+// RUN: %target-swift-ide-test -code-completion -completion-cache-path %t/cc-cache -source-filename %t/split/with-spi-import.swift -I %t/build -code-completion-token=COMPLETE | %FileCheck %s --check-prefix=WITH_SPI
+// RUN: %target-swift-ide-test -code-completion -completion-cache-path %t/cc-cache -source-filename %t/split/with-different-spi-import.swift -I %t/build -code-completion-token=COMPLETE | %FileCheck %s --check-prefix=WITHOUT_SPI
+// RUN: %target-swift-ide-test -code-completion -completion-cache-path %t/cc-cache -source-filename %t/split/with-normal-import.swift -I %t/build -code-completion-token=COMPLETE | %FileCheck %s --check-prefix=WITHOUT_SPI
+
+// First completion in file without SPI import, then with SPI import
+// RUN: %empty-directory(%t/cc-cache)
+// RUN: %target-swift-ide-test -code-completion -completion-cache-path %t/cc-cache -source-filename %t/split/with-normal-import.swift -I %t/build -code-completion-token=COMPLETE | %FileCheck %s --check-prefix=WITHOUT_SPI
+// RUN: %target-swift-ide-test -code-completion -completion-cache-path %t/cc-cache -source-filename %t/split/with-different-spi-import.swift -I %t/build -code-completion-token=COMPLETE | %FileCheck %s --check-prefix=WITHOUT_SPI
+// RUN: %target-swift-ide-test -code-completion -completion-cache-path %t/cc-cache -source-filename %t/split/with-spi-import.swift -I %t/build -code-completion-token=COMPLETE | %FileCheck %s --check-prefix=WITH_SPI
+
+// WITH_SPI: Begin completions
+// WITH_SPI-DAG: Decl[FreeFunction]/OtherModule[pck]: apiFunc()[#Void#]; name=apiFunc()
+// WITH_SPI-DAG: Decl[FreeFunction]/OtherModule[pck]: spiFunc()[#Void#]; name=spiFunc()
+// WITH_SPI: End completions
+
+// WITHOUT_SPI: Begin completions
+// WITHOUT_SPI-NOT: spiFunc
+// WITHOUT_SPI-DAG: Decl[FreeFunction]/OtherModule[pck]: apiFunc()[#Void#]; name=apiFunc()
+// WITHOUT_SPI-NOT: spiFunc
+// WITHOUT_SPI: End completions
+
+
+// BEGIN pck.swift
+
+public func apiFunc() {}
+
+@_spi(MySPI)
+public func spiFunc() {}
+
+// BEGIN with-spi-import.swift
+
+@_spi(MySPI) import pck
+
+func test() {
+  #^COMPLETE^#
+}
+
+// BEGIN with-different-spi-import.swift
+
+@_spi(OtherSPI) import pck
+
+func test() {
+  #^COMPLETE^#
+}
+
+
+// BEGIN with-normal-import.swift
+
+import pck
+
+func test() {
+  #^COMPLETE^#
+}
