[BatchMode] Fix interaction of bridging PCH and file specific diag consumer.

lib/AST/DiagnosticConsumer.cpp

@@ -19,6 +19,7 @@
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/SourceManager.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Debug.h"
@@ -116,6 +117,20 @@ FileSpecificDiagnosticConsumer::consumerForLocation(SourceManager &SM,
   // FileSpecificDiagnosticConsumer to be set up before the source files are
   // actually loaded.
   if (ConsumersOrderedByRange.empty()) {
+
+    // It's possible to get here while a bridging header PCH is being
+    // attached-to, if there's some sort of AST-reader warning or error, which
+    // happens before CompilerInstance::setUpInputs(), at which point _no_
+    // source buffers are loaded in yet. In that case we return nullptr, rather
+    // than trying to build a nonsensical map (and actually crashing since we
+    // can't find buffers for the inputs).
+    assert(!SubConsumers.empty());
+    if (!SM.getIDForBufferIdentifier(SubConsumers.begin()->first).hasValue()) {
+      assert(llvm::none_of(SubConsumers, [&](const ConsumerPair &pair) {
+            return SM.getIDForBufferIdentifier(pair.first).hasValue();
+          }));
+      return nullptr;
+    }
     auto *mutableThis = const_cast<FileSpecificDiagnosticConsumer*>(this);
     mutableThis->computeConsumersOrderedByRange(SM);
   }

test/Driver/batch_mode_bridging_pch.swift

@@ -5,6 +5,18 @@
 //
 // RUN: %swiftc_driver -enable-bridging-pch -v -import-objc-header %t/foo-bridging-header.h -enable-batch-mode -c -emit-module -module-name main -j 2 %t/file-01.swift %t/file-02.swift %t/file-03.swift %t/main.swift %s 2>&1 | %FileCheck %s
 //
+// Next we make a module map with an unknown attribute, which will cause an
+// AST-reader warning while (re)parsing the module map, while attaching a PCH.
+// We turn on serialized diagnostics in the frontends, and check that that
+// warning, issued before the batch-mode multi-file diagnostic multiplexor has
+// its file mappings established, does not crash the multiplexor.
+//
+// RUN: %empty-directory(%t/MyModule)
+// RUN: echo 'module MyModule [DefinitelyNotAnAttribute] { header "header.h" export * }' >%t/MyModule/module.modulemap
+// RUN: touch %t/MyModule/header.h
+// RUN: echo '#include "MyModule/header.h"' >>%t/foo-bridging-header.h
+// RUN: %swiftc_driver -enable-bridging-pch -v -I %t -import-objc-header %t/foo-bridging-header.h -enable-batch-mode -c -emit-module -module-name main -j 2 %t/file-01.swift %t/file-02.swift %t/file-03.swift %t/main.swift -serialize-diagnostics %s 2>&1 | %FileCheck %s
+//
 // CHECK: -emit-pch
 // CHECK: -primary-file {{.*}}/file-01.swift -primary-file {{.*}}/file-02.swift