[Serialization] Preserve @_implementationOnly through module merging

When we build incrementally, we produce "partial swiftmodules" for
each input source file, then merge them together into the final
compiled module that, among other things, gets used for debugging.
Without this, we'd drop @_implementationOnly imports and any types
from the modules that were imported during the module-merging step
and then be unable to debug those types

Author: jrose-apple

include/swift/Serialization/ModuleFile.h

@@ -659,11 +659,26 @@ class ModuleFile
   // Out of line to avoid instantiation OnDiskChainedHashTable here.
   ~ModuleFile();
 
-  /// Associates this module file with an AST module.
-  ///
-  /// Returns any error that occurred during association, including validation
-  /// that the module file is compatible with the module it's being loaded as.
-  Status associateWithFileContext(FileUnit *file, SourceLoc diagLoc);
+  /// Associates this module file with the AST node representing it.
+  ///
+  /// Checks that the file is compatible with the AST module it's being loaded
+  /// into, loads any dependencies needed to understand the module, and updates
+  /// the ASTContext and ClangImporter with search paths and other information
+  /// from the module.
+  ///
+  /// \param file The FileUnit that represents this file's place in the AST.
+  /// \param diagLoc A location used for diagnostics that occur during loading.
+  /// This does not include diagnostics about \e this file failing to load,
+  /// but rather other things that might be imported as part of bringing the
+  /// file into the AST.
+  /// \param treatAsPartialModule If true, processes implementation-only
+  /// information instead of assuming the client won't need it and shouldn't
+  /// see it.
+  ///
+  /// \returns any error that occurred during association, such as being
+  /// compiled for a different OS.
+  Status associateWithFileContext(FileUnit *file, SourceLoc diagLoc,
+                                  bool treatAsPartialModule);
 
   /// Checks whether this module can be used.
   Status getStatus() const {

include/swift/Serialization/SerializedModuleLoader.h

@@ -94,7 +94,7 @@ class SerializedModuleLoaderBase : public ModuleLoader {
   FileUnit *loadAST(ModuleDecl &M, Optional<SourceLoc> diagLoc,
                     std::unique_ptr<llvm::MemoryBuffer> moduleInputBuffer,
                     std::unique_ptr<llvm::MemoryBuffer> moduleDocInputBuffer,
-                    bool isFramework = false);
+                    bool isFramework, bool treatAsPartialModule);
 
   /// Check whether the module with a given name can be imported without
   /// importing it.

lib/Frontend/Frontend.cpp

@@ -868,7 +868,8 @@ bool CompilerInstance::parsePartialModulesAndLibraryFiles(
   for (auto &PM : PartialModules) {
     assert(PM.ModuleBuffer);
     if (!SML->loadAST(*MainModule, SourceLoc(), std::move(PM.ModuleBuffer),
-                      std::move(PM.ModuleDocBuffer)))
+                      std::move(PM.ModuleDocBuffer), /*isFramework*/false,
+                      /*treatAsPartialModule*/true))
       hadLoadError = true;
   }
 

lib/Serialization/ModuleFile.cpp

@@ -1460,7 +1460,8 @@ ModuleFile::ModuleFile(
 }
 
 Status ModuleFile::associateWithFileContext(FileUnit *file,
-                                            SourceLoc diagLoc) {
+                                            SourceLoc diagLoc,
+                                            bool treatAsPartialModule) {
   PrettyStackTraceModuleFile stackEntry(*this);
 
   assert(getStatus() == Status::Valid && "invalid module file");
@@ -1513,7 +1514,7 @@ Status ModuleFile::associateWithFileContext(FileUnit *file,
       continue;
     }
 
-    if (dependency.isImplementationOnly())
+    if (dependency.isImplementationOnly() && !treatAsPartialModule)
       continue;
 
     StringRef modulePathStr = dependency.RawPath;

lib/Serialization/SerializedModuleLoader.cpp

@@ -366,7 +366,7 @@ FileUnit *SerializedModuleLoaderBase::loadAST(
     ModuleDecl &M, Optional<SourceLoc> diagLoc,
     std::unique_ptr<llvm::MemoryBuffer> moduleInputBuffer,
     std::unique_ptr<llvm::MemoryBuffer> moduleDocInputBuffer,
-    bool isFramework) {
+    bool isFramework, bool treatAsPartialModule) {
   assert(moduleInputBuffer);
 
   StringRef moduleBufferID = moduleInputBuffer->getBufferIdentifier();
@@ -402,7 +402,8 @@ FileUnit *SerializedModuleLoaderBase::loadAST(
 
     auto diagLocOrInvalid = diagLoc.getValueOr(SourceLoc());
     loadInfo.status =
-        loadedModuleFile->associateWithFileContext(fileUnit, diagLocOrInvalid);
+        loadedModuleFile->associateWithFileContext(fileUnit, diagLocOrInvalid,
+                                                   treatAsPartialModule);
     if (loadInfo.status == serialization::Status::Valid) {
       Ctx.bumpGeneration();
       LoadedModuleFiles.emplace_back(std::move(loadedModuleFile),
@@ -656,7 +657,8 @@ ModuleDecl *SerializedModuleLoaderBase::loadModule(SourceLoc importLoc,
   SWIFT_DEFER { M->setHasResolvedImports(); };
 
   if (!loadAST(*M, moduleID.second, std::move(moduleInputBuffer),
-               std::move(moduleDocInputBuffer), isFramework)) {
+               std::move(moduleDocInputBuffer), isFramework,
+               /*treatAsPartialModule*/false)) {
     M->setFailedToLoad();
   }
 

test/Serialization/module-merging-implementation-only.swift

@@ -0,0 +1,23 @@
+// This is a very simple test that module merging does not eliminate
+// @_implementationOnly imports or declarations referenced from those imports.
+// More thorough tests exist in LLDB, which can look into those imports when
+// debugging a client of the module with @_implementationOnly imports.
+
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -o %t %S/Inputs/def_struct.swift
+
+// RUN: %target-swift-frontend -emit-module -I %t -o %t/main~partial.swiftmodule -module-name main %s
+// RUN: llvm-bcanalyzer -dump %t/main~partial.swiftmodule | %FileCheck %s
+// RUN: grep -q TwoInts %t/main~partial.swiftmodule
+
+// RUN: %target-swift-frontend -merge-modules -emit-module -I %t -o %t/main.swiftmodule %t/main~partial.swiftmodule
+// RUN: llvm-bcanalyzer -dump %t/main.swiftmodule | %FileCheck %s
+// RUN: grep -q TwoInts %t/main.swiftmodule
+
+@_implementationOnly import def_struct
+
+struct Container {
+  var wrapped: TwoInts
+}
+
+// CHECK: <IMPORTED_MODULE abbrevid={{[0-9]+}} op0=2 op1=0{{.*}}/> blob data = 'def_struct'

tools/SourceKit/lib/SwiftLang/SwiftIndexing.cpp

@@ -190,7 +190,9 @@ static void indexModule(llvm::MemoryBuffer *Input,
     // documentation file.
     // FIXME: refactor the frontend to provide an easy way to figure out the
     // correct filename here.
-    auto FUnit = Loader->loadAST(*Mod, None, std::move(Buf), nullptr);
+    auto FUnit = Loader->loadAST(*Mod, None, std::move(Buf), nullptr,
+                                 /*isFramework*/false,
+                                 /*treatAsPartialModule*/false);
 
     // FIXME: Not knowing what went wrong is pretty bad. loadModule() should be
     // more modular, rather than emitting diagnostics itself.