Merge pull request #14900 from jrose-apple/mux

Wire up FileSpecificDiagnosticConsumer for serialized diagnostics

lib/FrontendTool/FrontendTool.cpp

@@ -1335,7 +1335,7 @@ static bool performCompileStepsPostSILGen(
 
   // Just because we had an AST error it doesn't mean we can't performLLVM.
   bool HadError = Instance.getASTContext().hadError();
-  
+
   // If the AST Context has no errors but no IRModule is available,
   // parallelIRGen happened correctly, since parallel IRGen produces multiple
   // modules.
@@ -1509,6 +1509,50 @@ computeStatsReporter(const CompilerInvocation &Invocation, CompilerInstance *Ins
       ProfileEvents, ProfileEntities);
 }
 
+/// Creates a diagnostic consumer that handles serializing diagnostics, based on
+/// the supplementary output paths specified by \p inputsAndOutputs.
+///
+/// The returned consumer will handle producing multiple serialized diagnostics
+/// files if necessary, by using sub-consumers for each file and dispatching to
+/// the right one.
+///
+/// If no serialized diagnostics are being produced, returns null.
+static std::unique_ptr<DiagnosticConsumer>
+createSerializedDiagnosticConsumerIfNeeded(
+    const FrontendInputsAndOutputs &inputsAndOutputs) {
+
+  // The "4" here is somewhat arbitrary. In practice we're going to have one
+  // sub-consumer for each .dia file we're trying to output, which (again in
+  // practice) is going to be 1 in WMO mode and equal to the number of primary
+  // inputs in batch mode. That in turn is going to be "the number of files we
+  // need to recompile in this build, divided by the number of jobs". So a
+  // value of "4" here means that there would be no heap allocation on a clean
+  // build of a module with up to 32 files on an 8-core machine, if the user
+  // doesn't customize anything.
+  SmallVector<FileSpecificDiagnosticConsumer::ConsumerPair, 4>
+      serializedConsumers;
+
+  inputsAndOutputs.forEachInputProducingSupplementaryOutput(
+      [&](const InputFile &Input) -> bool {
+    std::string serializedDiagnosticsPath = Input.serializedDiagnosticsPath();
+    if (serializedDiagnosticsPath.empty())
+      return false;
+
+    serializedConsumers.emplace_back(
+        Input.file(),
+        serialized_diagnostics::createConsumer(serializedDiagnosticsPath));
+
+    // Continue for other inputs.
+    return false;
+  });
+
+  if (serializedConsumers.empty())
+    return nullptr;
+  if (serializedConsumers.size() == 1)
+    return std::move(serializedConsumers.front()).second;
+  return llvm::make_unique<FileSpecificDiagnosticConsumer>(serializedConsumers);
+}
+
 int swift::performFrontend(ArrayRef<const char *> Args,
                            const char *Argv0, void *MainAddr,
                            FrontendObserver *observer) {
@@ -1621,20 +1665,13 @@ int swift::performFrontend(ArrayRef<const char *> Args,
   // arguments are generated by the driver, not directly by a user. The driver
   // is responsible for emitting diagnostics for its own errors. See SR-2683
   // for details.
-  //
-  // FIXME: Do we need one SerializedConsumer per primary? Owned by the
-  // SourceFile?
-  std::unique_ptr<DiagnosticConsumer> SerializedConsumer;
-  {
-    const std::string &SerializedDiagnosticsPath =
-        Invocation.getSerializedDiagnosticsPathForAtMostOnePrimary();
-    if (!SerializedDiagnosticsPath.empty()) {
-      SerializedConsumer.reset(
-          serialized_diagnostics::createConsumer(SerializedDiagnosticsPath));
-      Instance->addDiagnosticConsumer(SerializedConsumer.get());
-    }
-  }
+  std::unique_ptr<DiagnosticConsumer> SerializedConsumerDispatcher =
+      createSerializedDiagnosticConsumerIfNeeded(
+        Invocation.getFrontendOptions().InputsAndOutputs);
+  if (SerializedConsumerDispatcher)
+    Instance->addDiagnosticConsumer(SerializedConsumerDispatcher.get());
 
+  // FIXME: The fix-its need to be multiplexed like the serialized diagnostics.
   std::unique_ptr<DiagnosticConsumer> FixitsConsumer;
   {
     const std::string &FixitsOutputPath =