[5.1][CodeCompletion] Complete Swift only module name after 'import'

include/swift/AST/ASTContext.h

@@ -892,9 +892,9 @@ class ASTContext final {
     return getIdentifier(getSwiftName(kind));
   }
 
-  /// Collect visible clang modules from the ClangModuleLoader. These modules are
-  /// not necessarily loaded.
-  void getVisibleTopLevelClangModules(SmallVectorImpl<clang::Module*> &Modules) const;
+  /// Populate \p names with visible top level module names.
+  /// This guarantees that resulted \p names doesn't have duplicated names.
+  void getVisibleTopLevelModuleNames(SmallVectorImpl<Identifier> &names) const;
 
 private:
   /// Register the given generic signature builder to be used as the canonical

include/swift/AST/ModuleLoader.h

@@ -82,6 +82,13 @@ class ModuleLoader {
 public:
   virtual ~ModuleLoader() = default;
 
+  /// Collect visible module names.
+  ///
+  /// Append visible module names to \p names. Note that names are possibly
+  /// duplicated, and not guaranteed to be ordered in any way.
+  virtual void collectVisibleTopLevelModuleNames(
+      SmallVectorImpl<Identifier> &names) const = 0;
+
   /// Check whether the module with a given name can be imported without
   /// importing it.
   ///

include/swift/ClangImporter/ClangImporter.h

@@ -119,6 +119,11 @@ class ClangImporter final : public ClangModuleLoader {
   static std::shared_ptr<clang::DependencyCollector>
   createDependencyCollector(bool TrackSystemDeps);
 
+  /// Append visible module names to \p names. Note that names are possibly
+  /// duplicated, and not guaranteed to be ordered in any way.
+  void collectVisibleTopLevelModuleNames(
+      SmallVectorImpl<Identifier> &names) const override;
+
   /// Check whether the module with a given name can be imported without
   /// importing it.
   ///
@@ -336,7 +341,7 @@ class ClangImporter final : public ClangModuleLoader {
   /// Calling this function does not load the module.
   void collectSubModuleNames(
       ArrayRef<std::pair<Identifier, SourceLoc>> path,
-      std::vector<std::string> &names);
+      std::vector<std::string> &names) const;
 
   /// Given a Clang module, decide whether this module is imported already.
   static bool isModuleImported(const clang::Module *M);

include/swift/DWARFImporter/DWARFImporter.h

@@ -66,6 +66,9 @@ class DWARFImporter final : public ClangModuleLoader {
 
   ~DWARFImporter();
 
+  void collectVisibleTopLevelModuleNames(
+      SmallVectorImpl<Identifier> &names) const override;
+
   /// Check whether the module with a given name can be imported without
   /// importing it.
   ///

include/swift/Frontend/ParseableInterfaceModuleLoader.h

@@ -154,6 +154,11 @@ class ParseableInterfaceModuleLoader : public SerializedModuleLoaderBase {
                                          RemarkOnRebuildFromInterface));
   }
 
+  /// Append visible module names to \p names. Note that names are possibly
+  /// duplicated, and not guaranteed to be ordered in any way.
+  void collectVisibleTopLevelModuleNames(
+      SmallVectorImpl<Identifier> &names) const override;
+
   /// Unconditionally build \p InPath (a swiftinterface file) to \p OutPath (as
   /// a swiftmodule file).
   ///

include/swift/Sema/SourceLoader.h

@@ -48,6 +48,11 @@ class SourceLoader : public ModuleLoader {
   SourceLoader &operator=(const SourceLoader &) = delete;
   SourceLoader &operator=(SourceLoader &&) = delete;
 
+  /// Append visible module names to \p names. Note that names are possibly
+  /// duplicated, and not guaranteed to be ordered in any way.
+  void collectVisibleTopLevelModuleNames(
+      SmallVectorImpl<Identifier> &names) const override;
+
   /// Check whether the module with a given name can be imported without
   /// importing it.
   ///

include/swift/Serialization/SerializedModuleLoader.h

@@ -44,6 +44,9 @@ class SerializedModuleLoaderBase : public ModuleLoader {
   SerializedModuleLoaderBase(ASTContext &ctx, DependencyTracker *tracker,
                              ModuleLoadingMode LoadMode);
 
+  void collectVisibleTopLevelModuleNamesImpl(SmallVectorImpl<Identifier> &names,
+                                             StringRef extension) const;
+
   using AccessPathElem = std::pair<Identifier, SourceLoc>;
   bool findModule(AccessPathElem moduleID,
                   std::unique_ptr<llvm::MemoryBuffer> *moduleBuffer,
@@ -169,6 +172,11 @@ class SerializedModuleLoader : public SerializedModuleLoaderBase {
 public:
   virtual ~SerializedModuleLoader();
 
+  /// Append visible module names to \p names. Note that names are possibly
+  /// duplicated, and not guaranteed to be ordered in any way.
+  void collectVisibleTopLevelModuleNames(
+      SmallVectorImpl<Identifier> &names) const override;
+
   /// Create a new importer that can load serialized Swift modules
   /// into the given ASTContext.
   static std::unique_ptr<SerializedModuleLoader>
@@ -220,6 +228,9 @@ class MemoryBufferSerializedModuleLoader : public SerializedModuleLoaderBase {
     MemoryBuffers[AccessPath] = std::move(input);
   }
 
+  void collectVisibleTopLevelModuleNames(
+      SmallVectorImpl<Identifier> &names) const override {}
+
   /// Create a new importer that can load serialized Swift modules
   /// into the given ASTContext.
   static std::unique_ptr<MemoryBufferSerializedModuleLoader>

lib/AST/ASTContext.cpp

@@ -46,10 +46,6 @@
 #include "swift/Syntax/SyntaxArena.h"
 #include "swift/Strings.h"
 #include "swift/Subsystems.h"
-#include "clang/AST/Attr.h"
-#include "clang/AST/DeclObjC.h"
-#include "clang/Lex/HeaderSearch.h"
-#include "clang/Lex/Preprocessor.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringMap.h"
@@ -1502,12 +1498,6 @@ ModuleDecl *ASTContext::getLoadedModule(Identifier ModuleName) const {
   return LoadedModules.lookup(ModuleName);
 }
 
-void ASTContext::getVisibleTopLevelClangModules(
-    SmallVectorImpl<clang::Module*> &Modules) const {
-  getClangModuleLoader()->getClangPreprocessor().getHeaderSearchInfo().
-    collectAllModules(Modules);
-}
-
 static AllocationArena getArena(GenericSignature *genericSig) {
   if (!genericSig)
     return AllocationArena::Permanent;
@@ -1749,6 +1739,19 @@ void ProtocolDecl::setDefaultAssociatedConformanceWitness(
   (void)pair;
 }
 
+void ASTContext::getVisibleTopLevelModuleNames(
+    SmallVectorImpl<Identifier> &names) const {
+  names.clear();
+  for (auto &importer : getImpl().ModuleLoaders)
+    importer->collectVisibleTopLevelModuleNames(names);
+
+  // Sort and unique.
+  std::sort(names.begin(), names.end(), [](Identifier LHS, Identifier RHS) {
+    return LHS.str().compare_lower(RHS.str()) < 0;
+  });
+  names.erase(std::unique(names.begin(), names.end()), names.end());
+}
+
 bool ASTContext::canImportModule(std::pair<Identifier, SourceLoc> ModulePath) {
   // If this module has already been successfully imported, it is importable.
   if (getLoadedModule(ModulePath) != nullptr)

lib/ClangImporter/ClangImporter.cpp

@@ -1494,9 +1494,22 @@ ClangImporter::emitBridgingPCH(StringRef headerPath,
   return false;
 }
 
+void ClangImporter::collectVisibleTopLevelModuleNames(
+    SmallVectorImpl<Identifier> &names) const {
+  SmallVector<clang::Module *, 32> Modules;
+  Impl.getClangPreprocessor().getHeaderSearchInfo().collectAllModules(Modules);
+  for (auto &M : Modules) {
+    if (!M->isAvailable())
+      continue;
+
+    names.push_back(
+        Impl.SwiftContext.getIdentifier(M->getTopLevelModuleName()));
+  }
+}
+
 void ClangImporter::collectSubModuleNames(
     ArrayRef<std::pair<Identifier, SourceLoc>> path,
-    std::vector<std::string> &names) {
+    std::vector<std::string> &names) const {
   auto &clangHeaderSearch = Impl.getClangPreprocessor().getHeaderSearchInfo();
 
   // Look up the top-level module first.

lib/DWARFImporter/DWARFImporter.cpp

@@ -142,6 +142,9 @@ DWARFImporter::DWARFImporter(ASTContext &ctx,
 
 DWARFImporter::~DWARFImporter() { delete &Impl; }
 
+void DWARFImporter::collectVisibleTopLevelModuleNames(
+    SmallVectorImpl<Identifier> &names) const {}
+
 bool DWARFImporter::canImportModule(std::pair<Identifier, SourceLoc> named) {
   return false;
 }

lib/Frontend/ParseableInterfaceModuleLoader.cpp

@@ -1377,3 +1377,10 @@ bool ParseableInterfaceModuleLoader::buildSwiftModuleFromSwiftInterface(
   return builder.buildSwiftModule(OutPath, /*shouldSerializeDeps*/true,
                                   /*ModuleBuffer*/nullptr);
 }
+
+void ParseableInterfaceModuleLoader::collectVisibleTopLevelModuleNames(
+    SmallVectorImpl<Identifier> &names) const {
+  collectVisibleTopLevelModuleNamesImpl(
+      names,
+      file_types::getExtension(file_types::TY_SwiftParseableInterfaceFile));
+}

lib/IDE/CodeCompletion.cpp

@@ -33,6 +33,7 @@
 #include "swift/Parse/CodeCompletionCallbacks.h"
 #include "swift/Sema/IDETypeChecking.h"
 #include "swift/Syntax/SyntaxKind.h"
+#include "swift/Strings.h"
 #include "swift/Subsystems.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Attr.h"
@@ -1755,35 +1756,21 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   }
 
   void addImportModuleNames() {
-    // FIXME: Add user-defined swift modules
-    SmallVector<StringRef, 20> ModuleNames;
-
-    // Collect clang module names.
-    {
-      SmallVector<clang::Module*, 20> ClangModules;
-      Ctx.getVisibleTopLevelClangModules(ClangModules);
-      for (auto *M : ClangModules) {
-        if (!M->isAvailable())
-          continue;
-        if (M->getTopLevelModuleName().startswith("_"))
-          continue;
-        if (M->getTopLevelModuleName() == Ctx.SwiftShimsModuleName.str())
-          continue;
-
-        ModuleNames.push_back(M->getTopLevelModuleName());
-      }
-    }
-
-    std::sort(ModuleNames.begin(), ModuleNames.end(),
-              [](StringRef LHS, StringRef RHS) {
-                return LHS.compare_lower(RHS) < 0;
-              });
+    SmallVector<Identifier, 0> ModuleNames;
+    Ctx.getVisibleTopLevelModuleNames(ModuleNames);
 
     llvm::StringSet<> ImportedModules;
     collectImportedModules(ImportedModules);
 
+    auto mainModuleName = CurrDeclContext->getParentModule()->getName();
     for (auto ModuleName : ModuleNames) {
-      auto MD = ModuleDecl::create(Ctx.getIdentifier(ModuleName), Ctx);
+      if (ModuleName.str().startswith("_") ||
+          ModuleName == mainModuleName ||
+          ModuleName == Ctx.SwiftShimsModuleName ||
+          ModuleName.str() == SWIFT_ONONE_SUPPORT)
+        continue;
+
+      auto MD = ModuleDecl::create(ModuleName, Ctx);
       CodeCompletionResultBuilder Builder(
           Sink,
           CodeCompletionResult::ResultKind::Declaration,

lib/Sema/SourceLoader.cpp

@@ -70,6 +70,11 @@ class SkipNonTransparentFunctions : public DelayedParsingCallbacks {
 
 } // unnamed namespace
 
+void SourceLoader::collectVisibleTopLevelModuleNames(
+    SmallVectorImpl<Identifier> &names) const {
+  // TODO: Implement?
+}
+
 bool SourceLoader::canImportModule(std::pair<Identifier, SourceLoc> ID) {
   // Search the memory buffers to see if we can find this file on disk.
   FileOrError inputFileOrError = findModule(Ctx, ID.first.str(),