Merge pull request #17186 from jrose-apple/bridge-over-troubled-imports (#17300)

[Serialization] Always list the bridging header before any imports

rdar://problem/40471329
(cherry picked from commit e5cec5f)

Author: jrose-apple

include/swift/Serialization/ModuleFile.h

@@ -625,6 +625,13 @@ class ModuleFile
     return static_cast<Status>(Bits.Status);
   }
 
+  /// Transfers ownership of a buffer that might contain source code where
+  /// other parts of the compiler could have emitted diagnostics, to keep them
+  /// alive even if the ModuleFile is destroyed.
+  ///
+  /// Should only be called when getStatus() indicates a failure.
+  std::unique_ptr<llvm::MemoryBuffer> takeBufferForDiagnostics();
+
   /// Returns the list of modules this module depends on.
   ArrayRef<Dependency> getDependencies() const {
     return Dependencies;

include/swift/Serialization/SerializedModuleLoader.h

@@ -29,6 +29,8 @@ class SerializedModuleLoader : public ModuleLoader {
   using LoadedModulePair = std::pair<std::unique_ptr<ModuleFile>, unsigned>;
   std::vector<LoadedModulePair> LoadedModuleFiles;
 
+  SmallVector<std::unique_ptr<llvm::MemoryBuffer>, 2> OrphanedMemoryBuffers;
+
   explicit SerializedModuleLoader(ASTContext &ctx, DependencyTracker *tracker);
 
 public:

lib/Serialization/ModuleFile.cpp

@@ -1468,6 +1468,17 @@ Status ModuleFile::associateWithFileContext(FileUnit *file,
   return getStatus();
 }
 
+std::unique_ptr<llvm::MemoryBuffer> ModuleFile::takeBufferForDiagnostics() {
+  assert(getStatus() != Status::Valid);
+
+  // Today, the only buffer that might have diagnostics in them is the input
+  // buffer, and even then only if it has imported module contents.
+  if (!importedHeaderInfo.contents.empty())
+    return std::move(ModuleInputBuffer);
+
+  return nullptr;
+}
+
 ModuleFile::~ModuleFile() { }
 
 void ModuleFile::lookupValue(DeclName name,

lib/Serialization/Serialization.cpp

@@ -1070,28 +1070,38 @@ void Serializer::writeInputBlock(const SerializationOptions &options) {
   publicImportSet.insert(publicImports.begin(), publicImports.end());
 
   removeDuplicateImports(allImports);
+
   auto clangImporter =
     static_cast<ClangImporter *>(M->getASTContext().getClangModuleLoader());
-  ModuleDecl *importedHeaderModule = clangImporter->getImportedHeaderModule();
+  ModuleDecl *bridgingHeaderModule = clangImporter->getImportedHeaderModule();
+  ModuleDecl::ImportedModule bridgingHeaderImport{{}, bridgingHeaderModule};
+
+  // Make sure the bridging header module is always at the top of the import
+  // list, mimicking how it is processed before any module imports when
+  // compiling source files.
+  if (llvm::is_contained(allImports, bridgingHeaderImport)) {
+    off_t importedHeaderSize = 0;
+    time_t importedHeaderModTime = 0;
+    std::string contents;
+    if (!options.ImportedHeader.empty()) {
+      contents = clangImporter->getBridgingHeaderContents(
+          options.ImportedHeader, importedHeaderSize, importedHeaderModTime);
+    }
+    assert(publicImportSet.count(bridgingHeaderImport));
+    ImportedHeader.emit(ScratchRecord,
+                        publicImportSet.count(bridgingHeaderImport),
+                        importedHeaderSize, importedHeaderModTime,
+                        options.ImportedHeader);
+    if (!contents.empty()) {
+      contents.push_back('\0');
+      ImportedHeaderContents.emit(ScratchRecord, contents);
+    }
+  }
+
   ModuleDecl *theBuiltinModule = M->getASTContext().TheBuiltinModule;
   for (auto import : allImports) {
-    if (import.second == theBuiltinModule)
-      continue;
-
-    if (import.second == importedHeaderModule) {
-      off_t importedHeaderSize = 0;
-      time_t importedHeaderModTime = 0;
-      std::string contents;
-      if (!options.ImportedHeader.empty())
-        contents = clangImporter->getBridgingHeaderContents(
-            options.ImportedHeader, importedHeaderSize, importedHeaderModTime);
-      ImportedHeader.emit(ScratchRecord, publicImportSet.count(import),
-                          importedHeaderSize, importedHeaderModTime,
-                          options.ImportedHeader);
-      if (!contents.empty()) {
-        contents.push_back('\0');
-        ImportedHeaderContents.emit(ScratchRecord, contents);
-      }
+    if (import.second == theBuiltinModule ||
+        import.second == bridgingHeaderModule) {
       continue;
     }
 

lib/Serialization/SerializedModuleLoader.cpp

@@ -223,7 +223,14 @@ FileUnit *SerializedModuleLoader::loadAST(
     M.removeFile(*fileUnit);
   }
 
-  // This is the failure path. If we have a location, diagnose the issue.
+  // From here on is the failure path.
+
+  // Even though the module failed to load, it's possible its contents include
+  // a source buffer that need to survive because it's already been used for
+  // diagnostics.
+  if (auto orphanedBuffer = loadedModuleFile->takeBufferForDiagnostics())
+    OrphanedMemoryBuffers.push_back(std::move(orphanedBuffer));
+
   if (!diagLoc)
     return nullptr;
 

validation-test/Serialization/Inputs/bridging-header-first/AmbivalentProtocol.h

@@ -0,0 +1,2 @@
+@protocol AmbivalentProtocol
+@end

validation-test/Serialization/Inputs/bridging-header-first/Module/modular.h

@@ -0,0 +1 @@
+#import <AmbivalentProtocol.h>

validation-test/Serialization/Inputs/bridging-header-first/Module/module.modulemap

@@ -0,0 +1,4 @@
+module Module {
+    umbrella header "modular.h"
+    module * { export * }
+}

validation-test/Serialization/Inputs/bridging-header-first/bridging.h

@@ -0,0 +1 @@
+#import "AmbivalentProtocol.h"

validation-test/Serialization/Inputs/bridging-header-first/mangled.txt

@@ -0,0 +1 @@
+$S4main1CCACycfc

validation-test/Serialization/bridging-header-first.swift

@@ -0,0 +1,32 @@
+// RUN: %empty-directory(%t)
+
+// RUN: %target-build-swift -emit-module -emit-executable %s -g -I %S/Inputs/bridging-header-first/ -import-objc-header %S/Inputs/bridging-header-first/bridging.h -o %t/main
+// RUN: llvm-bcanalyzer -dump %t/main.swiftmodule | %FileCheck -check-prefix CHECK-DUMP %s
+// RUN: %lldb-moduleimport-test %t/main -type-from-mangled %S/Inputs/bridging-header-first/mangled.txt 2>&1 | %FileCheck -check-prefix CHECK-RESOLVED-TYPE %s
+
+// RUN: %target-build-swift -emit-module -emit-executable %s -g -I %S/Inputs/bridging-header-first/ -import-objc-header %S/Inputs/bridging-header-first/bridging.h -o %t/main -whole-module-optimization
+// RUN: llvm-bcanalyzer -dump %t/main.swiftmodule | %FileCheck -check-prefix CHECK-DUMP %s
+// RUN: %lldb-moduleimport-test %t/main -type-from-mangled %S/Inputs/bridging-header-first/mangled.txt 2>&1 | %FileCheck -check-prefix CHECK-RESOLVED-TYPE %s
+
+// REQUIRES: objc_interop
+
+// CHECK-DUMP-LABEL: CONTROL_BLOCK
+// CHECK-DUMP: MODULE_NAME
+// CHECK-DUMP-SAME: 'main'
+
+// CHECK-DUMP-LABEL: INPUT_BLOCK
+// CHECK-DUMP: IMPORTED_HEADER
+// CHECK-DUMP-SAME: '{{.+}}/bridging.h'
+// CHECK-DUMP: IMPORTED_MODULE
+// CHECK-DUMP-SAME: 'Module'
+// CHECK-DUMP: IMPORTED_MODULE
+// CHECK-DUMP-SAME: 'Swift'
+
+
+// CHECK-RESOLVED-TYPE: @convention(method) (C.Type) -> () -> C
+
+import Module
+class C {}
+extension C: AmbivalentProtocol {
+  func f() {}
+}