From 07b200130fd934a9fc9d7f01eb9e27d77a1c706c Mon Sep 17 00:00:00 2001
From: Jordan Rose <jordan_rose@apple.com>
Date: Wed, 13 Jun 2018 16:17:27 -0700
Subject: [PATCH] [Serialization] Keep buffers alive if they may have
 diagnostics

This can only happen in one case today: a module imports a bridging
header, but the header on disk has disappeared, and now we need to
fall back to the (often inadequate) version that's stored inside the
swiftmodule file. Even if the module fails to load, the bridging
header has already been imported, and so anything else that happens
might still emit diagnostics and need that text to be alive, which
means we need to keep the buffer alive too.
---
 include/swift/Serialization/ModuleFile.h             |  7 +++++++
 include/swift/Serialization/SerializedModuleLoader.h |  2 ++
 lib/Serialization/ModuleFile.cpp                     | 11 +++++++++++
 lib/Serialization/SerializedModuleLoader.cpp         |  9 ++++++++-
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/include/swift/Serialization/ModuleFile.h b/include/swift/Serialization/ModuleFile.h
index 256a01f6eac..5928c2a281c 100644
--- a/include/swift/Serialization/ModuleFile.h
+++ b/include/swift/Serialization/ModuleFile.h
@@ -620,6 +620,13 @@ public:
     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;
diff --git a/include/swift/Serialization/SerializedModuleLoader.h b/include/swift/Serialization/SerializedModuleLoader.h
index 13db71116b1..00fa1d52a31 100644
--- a/include/swift/Serialization/SerializedModuleLoader.h
+++ b/include/swift/Serialization/SerializedModuleLoader.h
@@ -29,6 +29,8 @@ private:
   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:
diff --git a/lib/Serialization/ModuleFile.cpp b/lib/Serialization/ModuleFile.cpp
index 68576db2017..dc237301300 100644
--- a/lib/Serialization/ModuleFile.cpp
+++ b/lib/Serialization/ModuleFile.cpp
@@ -1457,6 +1457,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,
diff --git a/lib/Serialization/SerializedModuleLoader.cpp b/lib/Serialization/SerializedModuleLoader.cpp
index d7f8729b454..02edb2e7b45 100644
--- a/lib/Serialization/SerializedModuleLoader.cpp
+++ b/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;