Add a RangeSpecificDiagnosticConsumer

If a top-level diagnostic is in a particular source range, the
RangeSpecificDiagnosticConsumer will funnel it and any attached notes
to a particular "sub-consumer" designated for that range (intended to
be used with whole files). If it's not in a range designated for any
sub-consumer, the diagnostic is passed to all registered
sub-consumers.

This is intended to be used for batch mode, so that diagnostics that
are definitely associated with a particular file can be emitted in
that file's .dia output, while diagnostics that may be associated with
other files (such as those that come from Clang) will still get
presented to the user.

Author: jrose-apple

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,53 @@ class NullDiagnosticConsumer : public DiagnosticConsumer {
                         ArrayRef<DiagnosticArgument> FormatArgs,
                         const DiagnosticInfo &Info) override;
 };
+
+/// \brief DiagnosticConsumer that funnels diagnostics in certain source ranges
+/// to particular sub-consumers.
+///
+/// Diagnostics that are not in one of the special ranges are emitted into every
+/// sub-consumer.
+class RangeSpecificDiagnosticConsumer : public DiagnosticConsumer {
+public:
+  using ConsumerPair =
+      std::pair<CharSourceRange, std::unique_ptr<DiagnosticConsumer>>;
+
+private:
+  SmallVector<std::unique_ptr<DiagnosticConsumer>, 4> SubConsumers;
+
+  /// A "map" sorted by the end locations of the ranges, so that a lookup by
+  /// position can be done using binary search.
+  ///
+  /// \see #consumerForLocation
+  SmallVector<std::pair<CharSourceRange, unsigned>, 4> LocationToConsumerMap;
+
+  /// 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 and
+  /// records their association with the CharSourceRanges.
+  ///
+  /// The ranges must not be overlapping.
+  explicit RangeSpecificDiagnosticConsumer(
+      MutableArrayRef<ConsumerPair> consumers);
+
+  void handleDiagnostic(SourceManager &SM, SourceLoc Loc,
+                        DiagnosticKind Kind,
+                        StringRef FormatString,
+                        ArrayRef<DiagnosticArgument> FormatArgs,
+                        const DiagnosticInfo &Info) override;
+
+   bool finishProcessing() override;
+
+private:
+  DiagnosticConsumer *consumerForLocation(SourceLoc loc) const;
+};
 
 } // end namespace swift
 

lib/AST/DiagnosticConsumer.cpp

@@ -22,7 +22,104 @@
 #include "llvm/Support/raw_ostream.h"
 using namespace swift;
 
-DiagnosticConsumer::~DiagnosticConsumer() { }
+DiagnosticConsumer::~DiagnosticConsumer() = default;
+
+llvm::SMLoc DiagnosticConsumer::getRawLoc(SourceLoc loc) {
+  return loc.Value;
+}
+
+RangeSpecificDiagnosticConsumer::RangeSpecificDiagnosticConsumer(
+    MutableArrayRef<ConsumerPair> consumers) {
+  using MapEntry = std::pair<CharSourceRange, unsigned>;
+
+  // Split up the ConsumerPairs into the "map" (to be sorted) and the actual
+  // owning vector (preserving order).
+  for (ConsumerPair &pair : consumers) {
+    LocationToConsumerMap.emplace_back(pair.first, SubConsumers.size());
+    SubConsumers.emplace_back(std::move(pair).second);
+  }
+
+  // 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(LocationToConsumerMap.begin(), LocationToConsumerMap.end(),
+            [](const MapEntry &left, const MapEntry &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.
+  assert(LocationToConsumerMap.end() ==
+           std::adjacent_find(LocationToConsumerMap.begin(),
+                              LocationToConsumerMap.end(),
+                              [](const MapEntry &left, const MapEntry &right) {
+                                return left.first.overlaps(right.first);
+                              }) &&
+         "overlapping ranges given to RangeSpecificDiagnosticConsumer");
+}
+
+DiagnosticConsumer *
+RangeSpecificDiagnosticConsumer::consumerForLocation(SourceLoc loc) const {
+  // "Find the first range whose end location is greater than or equal to
+  // 'loc'."
+  auto possiblyContainingRangeIter =
+      std::lower_bound(LocationToConsumerMap.begin(),
+                       LocationToConsumerMap.end(),
+                       loc,
+                       [](const std::pair<CharSourceRange, unsigned> &entry,
+                          SourceLoc loc) -> bool {
+    auto compare = std::less<const char *>();
+    return compare(getRawLoc(entry.first.getEnd()).getPointer(),
+                   getRawLoc(loc).getPointer());
+  });
+
+  if (possiblyContainingRangeIter != LocationToConsumerMap.end() &&
+      possiblyContainingRangeIter->first.contains(loc)) {
+    return SubConsumers[possiblyContainingRangeIter->second].get();
+  }
+
+  return nullptr;
+}
+
+void RangeSpecificDiagnosticConsumer::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(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->handleDiagnostic(SM, Loc, Kind, FormatString, FormatArgs,
+                                    Info);
+    }
+  }
+}
+
+bool RangeSpecificDiagnosticConsumer::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->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,

unittests/AST/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_swift_unittest(SwiftASTTests
+  DiagnosticConsumerTests.cpp
   SourceLocTests.cpp
   TestContext.cpp
   TypeMatchTests.cpp