RangeSpecificDiagnosticConsumer -> FileSpecificDiagnosticConsumer

It turns out that we need to have the diagnostic consumers set up
before we've actually opened the input files, which makes sense
because we might want to emit diagnostics about not being able to open
an input file. Switch to using file names instead, and mapping those
over to source ranges only once we actually need to handle a
diagnostic with a valid source location.

Author: jrose-apple

include/swift/AST/DiagnosticConsumer.h

@@ -113,24 +113,44 @@ class NullDiagnosticConsumer : public DiagnosticConsumer {
                         const DiagnosticInfo &Info) override;
 };
 
-/// \brief DiagnosticConsumer that funnels diagnostics in certain source ranges
-/// to particular sub-consumers.
+/// \brief DiagnosticConsumer that funnels diagnostics in certain files to
+/// particular sub-consumers.
 ///
-/// Diagnostics that are not in one of the special ranges are emitted into every
-/// sub-consumer.
-class RangeSpecificDiagnosticConsumer : public DiagnosticConsumer {
+/// 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<CharSourceRange, std::unique_ptr<DiagnosticConsumer>>;
+      std::pair<std::string, std::unique_ptr<DiagnosticConsumer>>;
 
 private:
-  SmallVector<std::unique_ptr<DiagnosticConsumer>, 4> SubConsumers;
+  /// All consumers owned by this FileSpecificDiagnosticConsumer.
+  const SmallVector<ConsumerPair, 4> SubConsumers;
+
+  using ConsumersOrderedByRangeEntry =
+    std::pair<CharSourceRange, DiagnosticConsumer *>;
 
-  /// A "map" sorted by the end locations of the ranges, so that a lookup by
-  /// position can be done using binary search.
+  /// 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<std::pair<CharSourceRange, unsigned>, 4> LocationToConsumerMap;
+  SmallVector<ConsumersOrderedByRangeEntry, 4> ConsumersOrderedByRange;
 
   /// Indicates which consumer to send Note diagnostics too.
   ///
@@ -141,12 +161,12 @@ class RangeSpecificDiagnosticConsumer : public DiagnosticConsumer {
   DiagnosticConsumer *ConsumerForSubsequentNotes = nullptr;
 
 public:
-  /// Takes ownership of the DiagnosticConsumers specified in \p consumers and
-  /// records their association with the CharSourceRanges.
+  /// Takes ownership of the DiagnosticConsumers specified in \p consumers.
   ///
-  /// The ranges must not be overlapping.
-  explicit RangeSpecificDiagnosticConsumer(
-      MutableArrayRef<ConsumerPair> 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,
@@ -157,7 +177,9 @@ class RangeSpecificDiagnosticConsumer : public DiagnosticConsumer {
    bool finishProcessing() override;
 
 private:
-  DiagnosticConsumer *consumerForLocation(SourceLoc loc) const;
+  void computeConsumersOrderedByRange(SourceManager &SM);
+  DiagnosticConsumer *consumerForLocation(SourceManager &SM,
+                                          SourceLoc loc) const;
 };
 
 } // end namespace swift

lib/AST/DiagnosticConsumer.cpp

@@ -14,10 +14,13 @@
 //
 //===----------------------------------------------------------------------===//
 
-#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;
@@ -28,62 +31,119 @@ llvm::SMLoc DiagnosticConsumer::getRawLoc(SourceLoc loc) {
   return loc.Value;
 }
 
-RangeSpecificDiagnosticConsumer::RangeSpecificDiagnosticConsumer(
-    MutableArrayRef<ConsumerPair> consumers) {
-  using MapEntry = std::pair<CharSourceRange, unsigned>;
+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;
+    }
 
-  // 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);
+    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(LocationToConsumerMap.begin(), LocationToConsumerMap.end(),
-            [](const MapEntry &left, const MapEntry &right) -> bool {
+  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.
-  assert(LocationToConsumerMap.end() ==
-           std::adjacent_find(LocationToConsumerMap.begin(),
-                              LocationToConsumerMap.end(),
-                              [](const MapEntry &left, const MapEntry &right) {
+  // 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 given to RangeSpecificDiagnosticConsumer");
+         "overlapping ranges despite having distinct files");
 }
 
 DiagnosticConsumer *
-RangeSpecificDiagnosticConsumer::consumerForLocation(SourceLoc loc) const {
-  // "Find the first range whose end location is greater than or equal to
-  // 'loc'."
+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(LocationToConsumerMap.begin(),
-                       LocationToConsumerMap.end(),
+      std::lower_bound(ConsumersOrderedByRange.begin(),
+                       ConsumersOrderedByRange.end(),
                        loc,
-                       [](const std::pair<CharSourceRange, unsigned> &entry,
+                       [](const ConsumersOrderedByRangeEntry &entry,
                           SourceLoc loc) -> bool {
     auto compare = std::less<const char *>();
     return compare(getRawLoc(entry.first.getEnd()).getPointer(),
                    getRawLoc(loc).getPointer());
   });
 
-  if (possiblyContainingRangeIter != LocationToConsumerMap.end() &&
+  if (possiblyContainingRangeIter != ConsumersOrderedByRange.end() &&
       possiblyContainingRangeIter->first.contains(loc)) {
-    return SubConsumers[possiblyContainingRangeIter->second].get();
+    return possiblyContainingRangeIter->second;
   }
 
   return nullptr;
 }
 
-void RangeSpecificDiagnosticConsumer::handleDiagnostic(
+void FileSpecificDiagnosticConsumer::handleDiagnostic(
     SourceManager &SM, SourceLoc Loc, DiagnosticKind Kind,
     StringRef FormatString, ArrayRef<DiagnosticArgument> FormatArgs,
     const DiagnosticInfo &Info) {
@@ -93,7 +153,7 @@ void RangeSpecificDiagnosticConsumer::handleDiagnostic(
   case DiagnosticKind::Error:
   case DiagnosticKind::Warning:
   case DiagnosticKind::Remark:
-    specificConsumer = consumerForLocation(Loc);
+    specificConsumer = consumerForLocation(SM, Loc);
     ConsumerForSubsequentNotes = specificConsumer;
     break;
   case DiagnosticKind::Note:
@@ -106,18 +166,18 @@ void RangeSpecificDiagnosticConsumer::handleDiagnostic(
                                        Info);
   } else {
     for (auto &subConsumer : SubConsumers) {
-      subConsumer->handleDiagnostic(SM, Loc, Kind, FormatString, FormatArgs,
-                                    Info);
+      subConsumer.second->handleDiagnostic(SM, Loc, Kind, FormatString,
+                                           FormatArgs, Info);
     }
   }
 }
 
-bool RangeSpecificDiagnosticConsumer::finishProcessing() {
+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->finishProcessing();
+    hadError |= subConsumer.second->finishProcessing();
   return hadError;
 }