Add a RangeSpecificDiagnosticConsumer

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