Flush pending diagnostics on crash. (#4337)

This risks diagnsotic formatting crashing, but I think we more
frequently see cases where it'd be interesting to know what diagnostics
were being delayed as part of the default sorting.

toolchain/diagnostics/diagnostic_consumer.h

@@ -45,6 +45,8 @@ class StreamDiagnosticConsumer : public DiagnosticConsumer {
   auto HandleDiagnostic(Diagnostic diagnostic) -> void override;
   auto Flush() -> void override { stream_->flush(); }
 
+  auto set_stream(llvm::raw_ostream* stream) -> void { stream_ = stream; }
+
  private:
   auto Print(const DiagnosticMessage& message, llvm::StringRef prefix) -> void;
 

toolchain/driver/BUILD

@@ -103,6 +103,7 @@ cc_library(
         "//common:ostream",
         "//common:version",
         "//common:vlog",
+        "//toolchain/base:pretty_stack_trace_function",
         "//toolchain/base:value_store",
         "//toolchain/check",
         "//toolchain/codegen",

toolchain/driver/compile_subcommand.cpp

@@ -6,6 +6,7 @@
 
 #include "common/vlog.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "toolchain/base/pretty_stack_trace_function.h"
 #include "toolchain/check/check.h"
 #include "toolchain/codegen/codegen.h"
 #include "toolchain/diagnostics/sorting_diagnostic_consumer.h"
@@ -485,6 +486,10 @@ class CompilationUnit {
     consumer_->Flush();
   }
 
+  // Flushes diagnostics, specifically as part of generating stack trace
+  // information.
+  auto FlushForStackTrace() -> void { consumer_->Flush(); }
+
   auto input_filename() -> llvm::StringRef { return input_filename_; }
   auto success() -> bool { return success_; }
   auto has_source() -> bool { return source_.has_value(); }
@@ -660,6 +665,25 @@ auto CompileSubcommand::Run(DriverEnv& driver_env) -> DriverResult {
     stream_consumer.Flush();
   });
 
+  PrettyStackTraceFunction flush_on_crash([&](llvm::raw_ostream& out) {
+    // When crashing, flush diagnostics. If sorting diagnostics, they can be
+    // redirected to the crash stream; if streaming, the original stream is
+    // flushed.
+    // TODO: Eventually we'll want to limit the count.
+    if (options_.stream_errors) {
+      out << "Flushing diagnostics\n";
+    } else {
+      out << "Pending diagnostics:\n";
+      stream_consumer.set_stream(&out);
+    }
+
+    for (auto& unit : units) {
+      unit->FlushForStackTrace();
+    }
+    stream_consumer.Flush();
+    stream_consumer.set_stream(&driver_env.error_stream);
+  });
+
   // Returns a DriverResult object. Called whenever Compile returns.
   auto make_result = [&]() {
     DriverResult result = {.success = true};