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

In batch mode, when using serialized diagnostics, we have a
FileSpecificDiagnosticConsumer splitting diags out to separate .dia files (one
per primary). This constructs a mapping from filenames to consumers when it's
first given a diag. So far so good.

It assumes, however, that there are source buffers to be found for each input
file it knows the name of. This is true -- and useful to assert -- once the
source buffers are set up. It's also a prerequisite for building the map it
uses to direct diags.

Unfortunately there's a small window between the consumer being built and the
source manager getting buffers, and in this window we attach to a bridging PCH.
If that attaching generates a diag of itself (say due to a bogus module map or
such) then the diag consumer is asked to build its map without source buffers.
The thing to do here is just to fall back to the "no mapping found" case; the
only tricky part is identifying when we're in that window of time. I've chosen
to use the approximating condition of "none of the file buffers exist at all".

Author: graydon

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