Wire up FileSpecificDiagnosticConsumer for serialized diagnostics

include/swift/AST/DiagnosticConsumer.h

@@ -19,6 +19,7 @@
 #ifndef SWIFT_BASIC_DIAGNOSTICCONSUMER_H
 #define SWIFT_BASIC_DIAGNOSTICCONSUMER_H
 
+#include "swift/Basic/LLVM.h"
 #include "swift/Basic/SourceLoc.h"
 #include "llvm/Support/SourceMgr.h"
 
@@ -111,6 +112,75 @@ class NullDiagnosticConsumer : public DiagnosticConsumer {
                         ArrayRef<DiagnosticArgument> FormatArgs,
                         const DiagnosticInfo &Info) override;
 };
+
+/// \brief DiagnosticConsumer that funnels diagnostics in certain files to
+/// particular sub-consumers.
+///
+/// The intended use case for such a consumer is "batch mode" compilations,
+/// where we want to record diagnostics for each file as if they were compiled
+/// separately. This is important for incremental builds, so that if a file has
+/// warnings but doesn't get recompiled in the next build, the warnings persist.
+///
+/// Diagnostics that are not in one of the special files are emitted into every
+/// sub-consumer. This is necessary to deal with, for example, diagnostics in a
+/// bridging header imported from Objective-C, which isn't really about the
+/// current file.
+class FileSpecificDiagnosticConsumer : public DiagnosticConsumer {
+public:
+  /// A diagnostic consumer, along with the name of the buffer that it should
+  /// be associated with. An empty string means that a consumer is not
+  /// associated with any particular buffer, and should only receive diagnostics
+  /// that are not in any of the other consumers' files.
+  using ConsumerPair =
+      std::pair<std::string, std::unique_ptr<DiagnosticConsumer>>;
+
+private:
+  /// All consumers owned by this FileSpecificDiagnosticConsumer.
+  const SmallVector<ConsumerPair, 4> SubConsumers;
+
+  using ConsumersOrderedByRangeEntry =
+    std::pair<CharSourceRange, DiagnosticConsumer *>;
+
+  /// The consumers owned by this FileSpecificDiagnosticConsumer, sorted by
+  /// the end locations of each file so that a lookup by position can be done
+  /// using binary search.
+  ///
+  /// Generated and cached when the first diagnostic with a location is emitted.
+  /// This allows diagnostics to be emitted before files are actually opened,
+  /// as long as they don't have source locations.
+  ///
+  /// \see #consumerForLocation
+  SmallVector<ConsumersOrderedByRangeEntry, 4> ConsumersOrderedByRange;
+
+  /// Indicates which consumer to send Note diagnostics too.
+  ///
+  /// Notes are always considered attached to the error, warning, or remark
+  /// that was most recently emitted.
+  ///
+  /// If null, Note diagnostics are sent to every consumer.
+  DiagnosticConsumer *ConsumerForSubsequentNotes = nullptr;
+
+public:
+  /// Takes ownership of the DiagnosticConsumers specified in \p consumers.
+  ///
+  /// There must not be two consumers for the same file (i.e., having the same
+  /// buffer name).
+  explicit FileSpecificDiagnosticConsumer(
+      SmallVectorImpl<ConsumerPair> &consumers);
+
+  void handleDiagnostic(SourceManager &SM, SourceLoc Loc,
+                        DiagnosticKind Kind,
+                        StringRef FormatString,
+                        ArrayRef<DiagnosticArgument> FormatArgs,
+                        const DiagnosticInfo &Info) override;
+
+   bool finishProcessing() override;
+
+private:
+  void computeConsumersOrderedByRange(SourceManager &SM);
+  DiagnosticConsumer *consumerForLocation(SourceManager &SM,
+                                          SourceLoc loc) const;
+};
 
 } // end namespace swift
 

include/swift/Frontend/InputFile.h

@@ -91,6 +91,10 @@ class InputFile {
   std::string loadedModuleTracePath() const {
     return getPrimarySpecificPaths().SupplementaryOutputs.LoadedModuleTracePath;
   }
+  std::string serializedDiagnosticsPath() const {
+    return getPrimarySpecificPaths().SupplementaryOutputs
+        .SerializedDiagnosticsPath;
+  }
 };
 } // namespace swift
 

include/swift/Frontend/SerializedDiagnosticConsumer.h

@@ -29,13 +29,14 @@ namespace swift {
   class DiagnosticConsumer;
 
   namespace serialized_diagnostics {
-    /// \brief Create a DiagnosticConsumer that serializes diagnostics to a
-    ///        file.
+    /// Create a DiagnosticConsumer that serializes diagnostics to a file, using
+    /// Clang's serialized diagnostics format.
     ///
-    /// \param serializedDiagnosticsPath the file path to write the diagnostics.
+    /// \param outputPath the file path to write the diagnostics to.
     ///
     /// \returns A new diagnostic consumer that serializes diagnostics.
-    DiagnosticConsumer *createConsumer(llvm::StringRef serializedDiagnosticsPath);
+    std::unique_ptr<DiagnosticConsumer>
+    createConsumer(llvm::StringRef outputPath);
   }
 }
 

lib/AST/DiagnosticConsumer.cpp

@@ -14,15 +14,172 @@
 //
 //===----------------------------------------------------------------------===//
 
-#define DEBUG_TYPE "swift-basic"
+#define DEBUG_TYPE "swift-ast"
 #include "swift/AST/DiagnosticConsumer.h"
 #include "swift/AST/DiagnosticEngine.h"
+#include "swift/Basic/Defer.h"
+#include "swift/Basic/SourceManager.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 using namespace swift;
 
-DiagnosticConsumer::~DiagnosticConsumer() { }
+DiagnosticConsumer::~DiagnosticConsumer() = default;
+
+llvm::SMLoc DiagnosticConsumer::getRawLoc(SourceLoc loc) {
+  return loc.Value;
+}
+
+LLVM_ATTRIBUTE_UNUSED
+static bool hasDuplicateFileNames(
+    ArrayRef<FileSpecificDiagnosticConsumer::ConsumerPair> consumers) {
+  llvm::StringSet<> seenFiles;
+  for (const auto &consumerPair : consumers) {
+    if (consumerPair.first.empty()) {
+      // We can handle multiple consumers that aren't associated with any file,
+      // because they only collect diagnostics that aren't in any of the special
+      // files. This isn't an important use case to support, but also SmallSet
+      // doesn't handle empty strings anyway!
+      continue;
+    }
+
+    bool isUnique = seenFiles.insert(consumerPair.first).second;
+    if (!isUnique)
+      return true;
+  }
+  return false;
+}
+
+FileSpecificDiagnosticConsumer::FileSpecificDiagnosticConsumer(
+    SmallVectorImpl<ConsumerPair> &consumers)
+  : SubConsumers(std::move(consumers)) {
+  assert(!SubConsumers.empty() &&
+         "don't waste time handling diagnostics that will never get emitted");
+  assert(!hasDuplicateFileNames(SubConsumers) &&
+         "having multiple consumers for the same file is not implemented");
+}
+
+void FileSpecificDiagnosticConsumer::computeConsumersOrderedByRange(
+    SourceManager &SM) {
+  // Look up each file's source range and add it to the "map" (to be sorted).
+  for (const ConsumerPair &pair : SubConsumers) {
+    if (pair.first.empty())
+      continue;
+
+    Optional<unsigned> bufferID = SM.getIDForBufferIdentifier(pair.first);
+    assert(bufferID.hasValue() && "consumer registered for unknown file");
+    CharSourceRange range = SM.getRangeForBuffer(bufferID.getValue());
+    ConsumersOrderedByRange.emplace_back(range, pair.second.get());
+  }
+
+  // Sort the "map" by buffer /end/ location, for use with std::lower_bound
+  // later. (Sorting by start location would produce the same sort, since the
+  // ranges must not be overlapping, but since we need to check end locations
+  // later it's consistent to sort by that here.)
+  std::sort(ConsumersOrderedByRange.begin(), ConsumersOrderedByRange.end(),
+            [](const ConsumersOrderedByRangeEntry &left,
+               const ConsumersOrderedByRangeEntry &right) -> bool {
+    auto compare = std::less<const char *>();
+    return compare(getRawLoc(left.first.getEnd()).getPointer(),
+                   getRawLoc(right.first.getEnd()).getPointer());
+  });
+
+  // Check that the ranges are non-overlapping. If the files really are all
+  // distinct, this should be trivially true, but if it's ever not we might end
+  // up mis-filing diagnostics.
+  assert(ConsumersOrderedByRange.end() ==
+           std::adjacent_find(ConsumersOrderedByRange.begin(),
+                              ConsumersOrderedByRange.end(),
+                              [](const ConsumersOrderedByRangeEntry &left,
+                                 const ConsumersOrderedByRangeEntry &right) {
+                                return left.first.overlaps(right.first);
+                              }) &&
+         "overlapping ranges despite having distinct files");
+}
+
+DiagnosticConsumer *
+FileSpecificDiagnosticConsumer::consumerForLocation(SourceManager &SM,
+                                                    SourceLoc loc) const {
+  // If there's only one consumer, we'll use it no matter what, because...
+  // - ...all diagnostics within the file will go to that consumer.
+  // - ...all diagnostics not within the file will not be claimed by any
+  //   consumer, and so will go to all (one) consumers.
+  if (SubConsumers.size() == 1)
+    return SubConsumers.front().second.get();
+
+  // Diagnostics with invalid locations always go to every consumer.
+  if (loc.isInvalid())
+    return nullptr;
+
+  // This map is generated on first use and cached, to allow the
+  // FileSpecificDiagnosticConsumer to be set up before the source files are
+  // actually loaded.
+  if (ConsumersOrderedByRange.empty()) {
+    auto *mutableThis = const_cast<FileSpecificDiagnosticConsumer*>(this);
+    mutableThis->computeConsumersOrderedByRange(SM);
+  }
+
+  // This std::lower_bound call is doing a binary search for the first range
+  // that /might/ contain 'loc'. Specifically, since the ranges are sorted
+  // by end location, it's looking for the first range where the end location
+  // is greater than or equal to 'loc'.
+  auto possiblyContainingRangeIter =
+      std::lower_bound(ConsumersOrderedByRange.begin(),
+                       ConsumersOrderedByRange.end(),
+                       loc,
+                       [](const ConsumersOrderedByRangeEntry &entry,
+                          SourceLoc loc) -> bool {
+    auto compare = std::less<const char *>();
+    return compare(getRawLoc(entry.first.getEnd()).getPointer(),
+                   getRawLoc(loc).getPointer());
+  });
+
+  if (possiblyContainingRangeIter != ConsumersOrderedByRange.end() &&
+      possiblyContainingRangeIter->first.contains(loc)) {
+    return possiblyContainingRangeIter->second;
+  }
+
+  return nullptr;
+}
+
+void FileSpecificDiagnosticConsumer::handleDiagnostic(
+    SourceManager &SM, SourceLoc Loc, DiagnosticKind Kind,
+    StringRef FormatString, ArrayRef<DiagnosticArgument> FormatArgs,
+    const DiagnosticInfo &Info) {
+
+  DiagnosticConsumer *specificConsumer;
+  switch (Kind) {
+  case DiagnosticKind::Error:
+  case DiagnosticKind::Warning:
+  case DiagnosticKind::Remark:
+    specificConsumer = consumerForLocation(SM, Loc);
+    ConsumerForSubsequentNotes = specificConsumer;
+    break;
+  case DiagnosticKind::Note:
+    specificConsumer = ConsumerForSubsequentNotes;
+    break;
+  }
+
+  if (specificConsumer) {
+    specificConsumer->handleDiagnostic(SM, Loc, Kind, FormatString, FormatArgs,
+                                       Info);
+  } else {
+    for (auto &subConsumer : SubConsumers) {
+      subConsumer.second->handleDiagnostic(SM, Loc, Kind, FormatString,
+                                           FormatArgs, Info);
+    }
+  }
+}
+
+bool FileSpecificDiagnosticConsumer::finishProcessing() {
+  // Deliberately don't use std::any_of here because we don't want early-exit
+  // behavior.
+  bool hadError = false;
+  for (auto &subConsumer : SubConsumers)
+    hadError |= subConsumer.second->finishProcessing();
+  return hadError;
+}
 
 void NullDiagnosticConsumer::handleDiagnostic(
     SourceManager &SM, SourceLoc Loc, DiagnosticKind Kind,

lib/Frontend/PrintingDiagnosticConsumer.cpp

@@ -63,10 +63,6 @@ namespace {
   };
 } // end anonymous namespace
 
-llvm::SMLoc DiagnosticConsumer::getRawLoc(SourceLoc loc) {
-  return loc.Value;
-}
-
 void PrintingDiagnosticConsumer::handleDiagnostic(
     SourceManager &SM, SourceLoc Loc, DiagnosticKind Kind,
     StringRef FormatString, ArrayRef<DiagnosticArgument> FormatArgs,

lib/Frontend/SerializedDiagnosticConsumer.cpp

@@ -222,9 +222,10 @@ class SerializedDiagnosticConsumer : public DiagnosticConsumer {
 };
 } // end anonymous namespace
 
-namespace swift { namespace serialized_diagnostics {
-  DiagnosticConsumer *createConsumer(StringRef serializedDiagnosticsPath) {
-    return new SerializedDiagnosticConsumer(serializedDiagnosticsPath);
+namespace swift {
+namespace serialized_diagnostics {
+  std::unique_ptr<DiagnosticConsumer> createConsumer(StringRef outputPath) {
+    return llvm::make_unique<SerializedDiagnosticConsumer>(outputPath);
   }
 } // namespace serialized_diagnostics
 } // namespace swift