Merge pull request #63931 from DougGregor/swift-syntax-grouped-diagnostics

[Diagnostics] Switch swift-syntax diagnostic style to grouped diagnostics

include/swift/AST/CASTBridging.h

@@ -118,17 +118,18 @@ extern "C" {
 /// information and then must be finished via \c SwiftDiagnostic_finish.
 BridgedDiagnostic SwiftDiagnostic_create(
     void *diagnosticEngine, BridgedDiagnosticSeverity severity,
-    void *_Nullable sourceLoc,
+    const void *_Nullable sourceLoc,
     const uint8_t *_Nullable text, long textLen);
 
 /// Highlight a source range as part of the diagnostic.
 void SwiftDiagnostic_highlight(
-    BridgedDiagnostic diag, void *_Nullable startLoc, void *_Nullable endLoc);
+    BridgedDiagnostic diag, const void *_Nullable startLoc, const void *_Nullable endLoc);
 
 /// Add a Fix-It to replace a source range as part of the diagnostic.
 void SwiftDiagnostic_fixItReplace(
     BridgedDiagnostic diag,
-    void *_Nullable replaceStartLoc, void *_Nullable replaceEndLoc,
+    const void *_Nullable replaceStartLoc,
+    const void *_Nullable replaceEndLoc,
     const uint8_t *_Nullable newText, long newTextLen);
 
 /// Finish the given diagnostic and emit it.

include/swift/AST/DiagnosticEngine.h

@@ -41,6 +41,7 @@ namespace swift {
   class Decl;
   class DeclAttribute;
   class DiagnosticEngine;
+  class GeneratedSourceInfo;
   class SourceManager;
   class ValueDecl;
   class SourceFile;
@@ -1478,6 +1479,10 @@ namespace swift {
 std::pair<unsigned, DeclName>
 getAccessorKindAndNameForDiagnostics(const ValueDecl *D);
 
+/// Retrieve the macro name for a generated source info that represents
+/// a macro expansion.
+DeclName getGeneratedSourceInfoMacroName(const GeneratedSourceInfo &info);
+
 } // end namespace swift
 
 #endif

include/swift/Frontend/PrintingDiagnosticConsumer.h

@@ -22,6 +22,7 @@
 #include "swift/Basic/DiagnosticOptions.h"
 #include "swift/Basic/LLVM.h"
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Support/Process.h"
 
@@ -47,10 +48,14 @@ class PrintingDiagnosticConsumer : public DiagnosticConsumer {
   bool SuppressOutput = false;
 
   /// swift-syntax rendering
+
+  /// A queued up source file known to the queued diagnostics.
+  using QueuedBuffer = void *;
+
+  /// The queued diagnostics structure.
   void *queuedDiagnostics = nullptr;
-  void *queuedSourceFile = nullptr;
-  unsigned queuedDiagnosticsBufferID;
-  StringRef queuedBufferName;
+  llvm::DenseMap<unsigned, QueuedBuffer> queuedBuffers;
+  unsigned queuedDiagnosticsOutermostBufferID;
 
 public:
   PrintingDiagnosticConsumer(llvm::raw_ostream &stream = llvm::errs());
@@ -91,6 +96,7 @@ public:
   }
 
 private:
+  void queueBuffer(SourceManager &sourceMgr, unsigned bufferID);
   void printDiagnostic(SourceManager &SM, const DiagnosticInfo &Info);
 };
   

lib/AST/CASTBridging.cpp

@@ -20,7 +20,7 @@ inline llvm::ArrayRef<T> getArrayRef(BridgedArrayRef bridged) {
   return {static_cast<const T *>(bridged.data), size_t(bridged.numElements)};
 }
 
-static SourceLoc getSourceLocFromPointer(void *loc) {
+static SourceLoc getSourceLocFromPointer(const void *loc) {
   auto smLoc = llvm::SMLoc::getFromPointer((const char *)loc);
   return SourceLoc(smLoc);
 }
@@ -46,7 +46,7 @@ namespace {
 
 BridgedDiagnostic SwiftDiagnostic_create(
     void *diagnosticEngine, BridgedDiagnosticSeverity severity,
-    void *sourceLocPtr,
+    const void *sourceLocPtr,
     const uint8_t *textPtr, long textLen
 ) {
   StringRef origText{
@@ -81,7 +81,7 @@ BridgedDiagnostic SwiftDiagnostic_create(
 
 /// Highlight a source range as part of the diagnostic.
 void SwiftDiagnostic_highlight(
-    BridgedDiagnostic diagPtr, void *startLocPtr, void *endLocPtr
+    BridgedDiagnostic diagPtr, const void *startLocPtr, const void *endLocPtr
 ) {
   SourceLoc startLoc = getSourceLocFromPointer(startLocPtr);
   SourceLoc endLoc = getSourceLocFromPointer(endLocPtr);
@@ -92,7 +92,8 @@ void SwiftDiagnostic_highlight(
 
 /// Add a Fix-It to replace a source range as part of the diagnostic.
 void SwiftDiagnostic_fixItReplace(
-    BridgedDiagnostic diagPtr, void *replaceStartLocPtr, void *replaceEndLocPtr,
+    BridgedDiagnostic diagPtr,
+    const void *replaceStartLocPtr, const void *replaceEndLocPtr,
     const uint8_t *newTextPtr, long newTextLen) {
 
   SourceLoc startLoc = getSourceLocFromPointer(replaceStartLocPtr);

lib/AST/DiagnosticEngine.cpp

@@ -1316,19 +1316,7 @@ std::vector<Diagnostic> DiagnosticEngine::getGeneratedSourceBufferNotes(
     case GeneratedSourceInfo::PeerMacroExpansion:
     case GeneratedSourceInfo::ConformanceMacroExpansion: {
       SourceRange origRange = expansionNode.getSourceRange();
-      DeclName macroName;
-      if (auto customAttr = generatedInfo->attachedMacroCustomAttr) {
-        // FIXME: How will we handle deserialized custom attributes like this?
-        auto declRefType = dyn_cast<DeclRefTypeRepr>(customAttr->getTypeRepr());
-        macroName = declRefType->getNameRef().getFullName();
-      } else if (auto expansionExpr = dyn_cast_or_null<MacroExpansionExpr>(
-              expansionNode.dyn_cast<Expr *>())) {
-        macroName = expansionExpr->getMacroName().getFullName();
-      } else {
-        auto expansionDecl =
-            cast<MacroExpansionDecl>(expansionNode.get<Decl *>());
-        macroName = expansionDecl->getMacroName().getFullName();
-      }
+      DeclName macroName = getGeneratedSourceInfoMacroName(*generatedInfo);
 
       Diagnostic expansionNote(diag::in_macro_expansion, macroName);
       expansionNote.setLoc(origRange.Start);
@@ -1510,3 +1498,37 @@ swift::getAccessorKindAndNameForDiagnostics(const ValueDecl *D) {
 
   return {NOT_ACCESSOR_INDEX, D->getName()};
 }
+
+DeclName
+swift::getGeneratedSourceInfoMacroName(const GeneratedSourceInfo &info) {
+  ASTNode expansionNode = ASTNode::getFromOpaqueValue(info.astNode);
+  switch (info.kind) {
+  case GeneratedSourceInfo::ExpressionMacroExpansion:
+  case GeneratedSourceInfo::FreestandingDeclMacroExpansion:
+  case GeneratedSourceInfo::AccessorMacroExpansion:
+  case GeneratedSourceInfo::MemberAttributeMacroExpansion:
+  case GeneratedSourceInfo::MemberMacroExpansion:
+  case GeneratedSourceInfo::PeerMacroExpansion:
+  case GeneratedSourceInfo::ConformanceMacroExpansion: {
+    DeclName macroName;
+    if (auto customAttr = info.attachedMacroCustomAttr) {
+      // FIXME: How will we handle deserialized custom attributes like this?
+      auto declRefType = cast<DeclRefTypeRepr>(customAttr->getTypeRepr());
+      return declRefType->getNameRef().getFullName();
+    }
+
+    if (auto expansionExpr = dyn_cast_or_null<MacroExpansionExpr>(
+            expansionNode.dyn_cast<Expr *>())) {
+      return expansionExpr->getMacroName().getFullName();
+    }
+
+    auto expansionDecl =
+        cast<MacroExpansionDecl>(expansionNode.get<Decl *>());
+      return expansionDecl->getMacroName().getFullName();
+  }
+
+  case GeneratedSourceInfo::PrettyPrinted:
+  case GeneratedSourceInfo::ReplacedFunctionBody:
+    return DeclName();
+  }
+}

lib/ASTGen/Sources/ASTGen/Diagnostics.swift

@@ -44,8 +44,8 @@ fileprivate func emitDiagnosticParts(
   // Emit highlights
   for highlight in highlights {
     SwiftDiagnostic_highlight(
-      diag, sourceLoc(at: highlight.position),
-      sourceLoc(at: highlight.endPosition)
+      diag, sourceLoc(at: highlight.positionAfterSkippingLeadingTrivia),
+      sourceLoc(at: highlight.endPositionBeforeTrailingTrivia)
     )
   }
 
@@ -154,8 +154,8 @@ extension SourceManager {
     for highlight in highlights {
       SwiftDiagnostic_highlight(
         diag,
-        cxxSourceLocation(for: highlight),
-        cxxSourceLocation(for: highlight, at: highlight.endPosition)
+        cxxSourceLocation(for: highlight, at: highlight.positionAfterSkippingLeadingTrivia),
+        cxxSourceLocation(for: highlight, at: highlight.endPositionBeforeTrailingTrivia)
       )
     }
 
@@ -241,47 +241,23 @@ extension SourceManager {
   }
 }
 
-/// A set of queued diagnostics created by the C++ compiler and rendered
-/// via the swift-syntax renderer.
 struct QueuedDiagnostics {
-  /// The source file in which all of the diagnostics occur.
-  let sourceFile: SourceFileSyntax
+  var grouped: GroupedDiagnostics = GroupedDiagnostics()
 
-  /// The underlying buffer within the C++ SourceManager, which is used
-  /// for computations of source locations.
-  let buffer: UnsafeBufferPointer<UInt8>
+  /// The source file IDs we allocated, mapped from the buffer IDs used
+  /// by the C++ source manager.
+  var sourceFileIDs: [Int: UnsafeMutablePointer<GroupedDiagnostics.SourceFileID>] = [:]
 
-  /// The set of diagnostics.
-  fileprivate var diagnostics: [Diagnostic] = []
-
-  mutating func diagnose(_ diagnostic: Diagnostic) {
-    assert(diagnostic.node.root == sourceFile.root)
-    diagnostics.append(diagnostic)
-  }
-
-  func render() -> String {
-    return DiagnosticsFormatter.annotatedSource(
-      tree: sourceFile,
-      diags: diagnostics
-    )
-  }
+  /// The known source files
+  var sourceFiles: [ExportedSourceFile] = []
 }
 
-/// Create a queued diagnostics structure in which we can
+/// Create a grouped diagnostics structure in which we can add osou
 @_cdecl("swift_ASTGen_createQueuedDiagnostics")
-public func createQueuedDiagnostics(
-  sourceFilePtr: UnsafeMutablePointer<UInt8>
-) -> UnsafeRawPointer {
-  return sourceFilePtr.withMemoryRebound(
-    to: ExportedSourceFile.self, capacity: 1
-  ) { sourceFile in
-    let ptr = UnsafeMutablePointer<QueuedDiagnostics>.allocate(capacity: 1)
-    ptr.initialize(to: .init(
-      sourceFile: sourceFile.pointee.syntax,
-      buffer: sourceFile.pointee.buffer)
-    )
-    return UnsafeRawPointer(ptr)
-  }
+public func createQueuedDiagnostics() -> UnsafeRawPointer {
+  let ptr = UnsafeMutablePointer<QueuedDiagnostics>.allocate(capacity: 1)
+  ptr.initialize(to: .init())
+  return UnsafeRawPointer(ptr)
 }
 
 /// Destroy the queued diagnostics.
@@ -290,6 +266,11 @@ public func destroyQueuedDiagnostics(
   queuedDiagnosticsPtr: UnsafeMutablePointer<UInt8>
 ) {
   queuedDiagnosticsPtr.withMemoryRebound(to: QueuedDiagnostics.self, capacity: 1) { queuedDiagnostics in
+    for (_, sourceFileID) in queuedDiagnostics.pointee.sourceFileIDs {
+      sourceFileID.deinitialize(count: 1)
+      sourceFileID.deallocate()
+    }
+
     queuedDiagnostics.deinitialize(count: 1)
     queuedDiagnostics.deallocate()
   }
@@ -319,6 +300,50 @@ extension BridgedDiagnosticSeverity {
   }
 }
 
+/// Register a source file wih the queued diagnostics.
+@_cdecl("swift_ASTGen_addQueuedSourceFile")
+public func addQueuedSourceFile(
+  queuedDiagnosticsPtr: UnsafeMutableRawPointer,
+  bufferID: Int,
+  sourceFilePtr: UnsafeRawPointer,
+  displayNamePtr: UnsafePointer<UInt8>,
+  displayNameLength: Int,
+  parentID: Int,
+  positionInParent: Int
+) {
+  let queuedDiagnostics = queuedDiagnosticsPtr.assumingMemoryBound(to: QueuedDiagnostics.self)
+  // Determine the parent link, for a child buffer.
+  let parent: (GroupedDiagnostics.SourceFileID, AbsolutePosition)?
+  if parentID >= 0,
+      let parentSourceFileID = queuedDiagnostics.pointee.sourceFileIDs[parentID] {
+    parent = (parentSourceFileID.pointee, AbsolutePosition(utf8Offset: positionInParent))
+  } else {
+    parent = nil
+  }
+
+  let displayName = String(
+    decoding: UnsafeBufferPointer(
+      start: displayNamePtr,
+      count: displayNameLength
+    ),
+    as: UTF8.self
+  )
+
+  // Add the source file.
+  let sourceFile = sourceFilePtr.assumingMemoryBound(to: ExportedSourceFile.self)
+  let sourceFileID = queuedDiagnostics.pointee.grouped.addSourceFile(
+    tree: sourceFile.pointee.syntax,
+    displayName: displayName,
+    parent: parent
+  )
+  queuedDiagnostics.pointee.sourceFiles.append(sourceFile.pointee)
+
+  // Record the buffer ID.
+  let allocatedSourceFileID = UnsafeMutablePointer<GroupedDiagnostics.SourceFileID>.allocate(capacity: 1)
+  allocatedSourceFileID.initialize(to: sourceFileID)
+  queuedDiagnostics.pointee.sourceFileIDs[bufferID] = allocatedSourceFileID
+}
+
 /// Add a new diagnostic to the queue.
 @_cdecl("swift_ASTGen_addQueuedDiagnostic")
 public func addQueuedDiagnostic(
@@ -326,38 +351,91 @@ public func addQueuedDiagnostic(
   text: UnsafePointer<UInt8>,
   textLength: Int,
   severity: BridgedDiagnosticSeverity,
-  position: UnsafePointer<UInt8>
+  position: CxxSourceLoc,
+  highlightRangesPtr: UnsafePointer<CxxSourceLoc>?,
+  numHighlightRanges: Int
 ) {
-  let queuedDiagnostics = queuedDiagnosticsPtr.bindMemory(
-    to: QueuedDiagnostics.self, capacity: 1
+  let queuedDiagnostics = queuedDiagnosticsPtr.assumingMemoryBound(
+    to: QueuedDiagnostics.self
   )
 
-  // Find the offset.
-  let buffer = queuedDiagnostics.pointee.buffer
-  let offset = position - buffer.baseAddress!
-  if offset < 0 || offset >= buffer.count {
+  // Find the source file that contains this location.
+  let sourceFile = queuedDiagnostics.pointee.sourceFiles.first { sf in
+    guard let baseAddress = sf.buffer.baseAddress else {
+      return false
+    }
+
+    return position >= baseAddress && position < baseAddress + sf.buffer.count
+  }
+  guard let sourceFile = sourceFile else {
+    // FIXME: Hard to report an error here...
     return
   }
 
   // Find the token at that offset.
-  let sf = queuedDiagnostics.pointee.sourceFile
-  guard let token = sf.token(at: AbsolutePosition(utf8Offset: offset)) else {
+  let sourceFileBaseAddress = sourceFile.buffer.baseAddress!
+  let sourceFileEndAddress = sourceFileBaseAddress + sourceFile.buffer.count
+  let offset = position - sourceFileBaseAddress
+  guard let token = sourceFile.syntax.token(at: AbsolutePosition(utf8Offset: offset)) else {
     return
   }
 
+  // Map the highlights.
+  var highlights: [Syntax] = []
+  let highlightRanges = UnsafeBufferPointer<CxxSourceLoc>(
+    start: highlightRangesPtr, count: numHighlightRanges * 2
+  )
+  for index in 0..<numHighlightRanges {
+    // Make sure both the start and the end land within this source file.
+    let start = highlightRanges[index * 2]
+    let end = highlightRanges[index * 2 + 1]
+
+    guard start >= sourceFileBaseAddress && start < sourceFileEndAddress,
+          end >= sourceFileBaseAddress && end <= sourceFileEndAddress else {
+      continue
+    }
+
+    // Find start tokens in the source file.
+    let startPos = AbsolutePosition(utf8Offset: start - sourceFileBaseAddress)
+    guard let startToken = sourceFile.syntax.token(at: startPos) else {
+      continue
+    }
+
+    // Walk up from the start token until we find a syntax node that matches
+    // the highlight range.
+    let endPos = AbsolutePosition(utf8Offset: end - sourceFileBaseAddress)
+    var highlightSyntax = Syntax(startToken)
+    while true {
+      // If this syntax matches our starting/ending positions, add the
+      // highlight and we're done.
+      if highlightSyntax.positionAfterSkippingLeadingTrivia == startPos &&
+          highlightSyntax.endPositionBeforeTrailingTrivia == endPos {
+        highlights.append(highlightSyntax)
+        break
+      }
+
+      // Go up to the parent.
+      guard let parent = highlightSyntax.parent else {
+        break
+      }
+
+      highlightSyntax = parent
+    }
+  }
+
   let textBuffer = UnsafeBufferPointer(start: text, count: textLength)
   let diagnostic = Diagnostic(
     node: Syntax(token),
     message: SimpleDiagnostic(
       message: String(decoding: textBuffer, as: UTF8.self),
       severity: severity.asSeverity
-    )
+    ),
+    highlights: highlights
   )
 
-  queuedDiagnostics.pointee.diagnose(diagnostic)
+  queuedDiagnostics.pointee.grouped.addDiagnostic(diagnostic)
 }
 
-
 /// Render the queued diagnostics into a UTF-8 string.
 @_cdecl("swift_ASTGen_renderQueuedDiagnostics")
 public func renterQueuedDiagnostics(
@@ -368,12 +446,8 @@ public func renterQueuedDiagnostics(
   renderedLength: UnsafeMutablePointer<Int>
 ) {
   queuedDiagnosticsPtr.withMemoryRebound(to: QueuedDiagnostics.self, capacity: 1) { queuedDiagnostics in
-    let renderedStr = DiagnosticsFormatter.annotatedSource(
-      tree: queuedDiagnostics.pointee.sourceFile,
-      diags: queuedDiagnostics.pointee.diagnostics,
-      contextSize: contextSize,
-      colorize: colorize != 0
-    )
+    let formatter = DiagnosticsFormatter(contextSize: contextSize, colorize: colorize != 0)
+    let renderedStr = formatter.annotateSources(in: queuedDiagnostics.pointee.grouped)
 
     (renderedPointer.pointee, renderedLength.pointee) =
         allocateUTF8String(renderedStr)

lib/ASTGen/Sources/ASTGen/PluginHost.swift

@@ -225,8 +225,7 @@ class PluginDiagnosticsEngine {
     else {
       return nil
     }
-    let address = bufferBaseAddress.advanced(by: offset)
-    return CxxSourceLoc(mutating: address)
+    return bufferBaseAddress.advanced(by: offset)
   }
 
   /// C++ source location from a position value from a plugin.

lib/ASTGen/Sources/ASTGen/SourceManager.swift

@@ -3,7 +3,7 @@ import SwiftSyntax
 import SwiftSyntaxMacros
 
 /// A C++ source location.
-typealias CxxSourceLoc = UnsafeMutablePointer<UInt8>
+public typealias CxxSourceLoc = UnsafePointer<UInt8>
 
 /// A source manager that keeps track of the source files in the program.
 class SourceManager {
@@ -128,6 +128,6 @@ extension SourceManager {
       return nil
     }
     let address = bufferBaseAddress.advanced(by: offsetFromSourceFile)
-    return CxxSourceLoc(mutating: address)
+    return address
   }
 }

lib/Frontend/PrintingDiagnosticConsumer.cpp

@@ -33,13 +33,23 @@
 using namespace swift;
 using namespace swift::markup;
 
-extern "C" void *swift_ASTGen_createQueuedDiagnostics(void *sourceFile);
+extern "C" void *swift_ASTGen_createQueuedDiagnostics();
 extern "C" void swift_ASTGen_destroyQueuedDiagnostics(void *queued);
+extern "C" void swift_ASTGen_addQueuedSourceFile(
+      void *queuedDiagnostics,
+      int bufferID,
+      void *sourceFile,
+      const uint8_t *displayNamePtr,
+      intptr_t displayNameLength,
+      int parentID,
+      int positionInParent);
 extern "C" void swift_ASTGen_addQueuedDiagnostic(
     void *queued,
     const char* text, ptrdiff_t textLength,
     BridgedDiagnosticSeverity severity,
-    const void *sourceLoc
+    const void *sourceLoc,
+    const void **highlightRanges,
+    ptrdiff_t numHighlightRanges
 );
 extern "C" void swift_ASTGen_renderQueuedDiagnostics(
     void *queued, ptrdiff_t contextSize, ptrdiff_t colorize,
@@ -984,11 +994,96 @@ static void enqueueDiagnostic(
     break;
   }
 
+  // Map the highlight ranges.
+  SmallVector<const void *, 2> highlightRanges;
+  for (const auto &range : info.Ranges) {
+    if (range.isInvalid())
+      continue;
+
+    highlightRanges.push_back(range.getStart().getOpaquePointerValue());
+    highlightRanges.push_back(range.getEnd().getOpaquePointerValue());
+  }
+
+  // FIXME: Translate Fix-Its.
+
   swift_ASTGen_addQueuedDiagnostic(
       queuedDiagnostics, text.data(), text.size(), severity,
-      info.Loc.getOpaquePointerValue());
+      info.Loc.getOpaquePointerValue(),
+      highlightRanges.data(), highlightRanges.size() / 2);
+}
+#endif
 
-  // FIXME: Need a way to add highlights, Fix-Its, and so on.
+/// Retrieve the stack of source buffers from the provided location out to
+/// a physical source file, with source buffer IDs for each step along the way
+/// due to (e.g.) macro expansions or generated code.
+///
+/// The resulting vector will always contain valid source locations. If the
+/// initial location is invalid, the result will be empty.
+static SmallVector<unsigned, 1> getSourceBufferStack(
+    SourceManager &sourceMgr, SourceLoc loc) {
+  SmallVector<unsigned, 1> stack;
+  while (true) {
+    if (loc.isInvalid())
+      return stack;
+
+    unsigned bufferID = sourceMgr.findBufferContainingLoc(loc);
+    stack.push_back(bufferID);
+
+    auto generatedSourceInfo = sourceMgr.getGeneratedSourceInfo(bufferID);
+    if (!generatedSourceInfo)
+      return stack;
+
+    loc = generatedSourceInfo->originalSourceRange.getStart();
+  }
+}
+
+#if SWIFT_SWIFT_PARSER
+void PrintingDiagnosticConsumer::queueBuffer(
+    SourceManager &sourceMgr, unsigned bufferID) {
+  QueuedBuffer knownSourceFile = queuedBuffers[bufferID];
+  if (knownSourceFile)
+    return;
+
+  auto bufferContents = sourceMgr.getEntireTextForBuffer(bufferID);
+  auto sourceFile = swift_ASTGen_parseSourceFile(
+      bufferContents.data(), bufferContents.size(),
+      "module", "file.swift");
+
+  // Find the parent and position in parent, if there is one.
+  int parentID = -1;
+  int positionInParent = 0;
+  std::string displayName;
+  auto generatedSourceInfo = sourceMgr.getGeneratedSourceInfo(bufferID);
+  if (generatedSourceInfo) {
+    SourceLoc parentLoc = generatedSourceInfo->originalSourceRange.getEnd();
+    if (parentLoc.isValid()) {
+      parentID = sourceMgr.findBufferContainingLoc(parentLoc);
+      positionInParent = sourceMgr.getLocOffsetInBuffer(parentLoc, parentID);
+
+      // Queue the parent buffer.
+      queueBuffer(sourceMgr, parentID);
+    }
+
+    if (DeclName macroName =
+            getGeneratedSourceInfoMacroName(*generatedSourceInfo)) {
+      SmallString<64> buffer;
+      if (generatedSourceInfo->attachedMacroCustomAttr)
+        displayName = ("macro expansion @" + macroName.getString(buffer)).str();
+      else
+        displayName = ("macro expansion #" + macroName.getString(buffer)).str();
+    }
+  }
+
+  if (displayName.empty()) {
+    displayName = sourceMgr.getDisplayNameForLoc(
+        sourceMgr.getLocForBufferStart(bufferID)).str();
+  }
+
+  swift_ASTGen_addQueuedSourceFile(
+      queuedDiagnostics, bufferID, sourceFile,
+      (const uint8_t*)displayName.data(), displayName.size(),
+      parentID, positionInParent);
+  queuedBuffers[bufferID] = sourceFile;
 }
 #endif
 
@@ -1008,32 +1103,21 @@ void PrintingDiagnosticConsumer::handleDiagnostic(SourceManager &SM,
   switch (FormattingStyle) {
   case DiagnosticOptions::FormattingStyle::SwiftSyntax: {
 #if SWIFT_SWIFT_PARSER
-    if (Info.Loc.isValid()) {
-      // Ignore "in macro expansion" diagnostics; we want to put them
-      // elsewhere.
-      // FIXME: We should render the "in macro expansion" information in
-      // some other manner. Not quite sure how at this point, though.
-      if (Info.ID == diag::in_macro_expansion.ID)
-        break;
-
-      // If there are no enqueued diagnostics, they are from a different
-      // buffer, flush any enqueued diagnostics and create a new set.
-      unsigned bufferID = SM.findBufferContainingLoc(Info.Loc);
-      if (!queuedDiagnostics || bufferID != queuedDiagnosticsBufferID) {
+    auto bufferStack = getSourceBufferStack(SM, Info.Loc);
+    if (!bufferStack.empty()) {
+      // If there are no enqueued diagnostics, or they are from a different
+      // outermost buffer, flush any enqueued diagnostics and start fresh.
+      unsigned outermostBufferID = bufferStack.back();
+      if (!queuedDiagnostics ||
+          outermostBufferID != queuedDiagnosticsOutermostBufferID) {
         flush(/*includeTrailingBreak*/ true);
 
-        // FIXME: Go parse the source file again. This is an awful hack.
-        auto bufferContents = SM.getEntireTextForBuffer(bufferID);
-        queuedSourceFile = swift_ASTGen_parseSourceFile(
-            bufferContents.data(), bufferContents.size(),
-            "module", "file.swift");
-
-        queuedBufferName = SM.getDisplayNameForLoc(Info.Loc);
-        queuedDiagnostics =
-            swift_ASTGen_createQueuedDiagnostics(queuedSourceFile);
-        queuedDiagnosticsBufferID = bufferID;
+        queuedDiagnosticsOutermostBufferID = outermostBufferID;
+        queuedDiagnostics = swift_ASTGen_createQueuedDiagnostics();
       }
 
+      unsigned innermostBufferID = bufferStack.front();
+      queueBuffer(SM, innermostBufferID);
       enqueueDiagnostic(queuedDiagnostics, Info, SM);
       break;
     }
@@ -1102,8 +1186,6 @@ void PrintingDiagnosticConsumer::flush(bool includeTrailingBreak) {
 
 #if SWIFT_SWIFT_PARSER
   if (queuedDiagnostics) {
-    Stream << "=== " << queuedBufferName << " ===\n";
-
     char *renderedString = nullptr;
     ptrdiff_t renderedStringLen = 0;
     swift_ASTGen_renderQueuedDiagnostics(
@@ -1113,9 +1195,11 @@ void PrintingDiagnosticConsumer::flush(bool includeTrailingBreak) {
       Stream.write(renderedString, renderedStringLen);
     }
     swift_ASTGen_destroyQueuedDiagnostics(queuedDiagnostics);
-    swift_ASTGen_destroySourceFile(queuedSourceFile);
     queuedDiagnostics = nullptr;
-    queuedSourceFile = nullptr;
+    for (const auto &buffer : queuedBuffers) {
+      swift_ASTGen_destroySourceFile(buffer.second);
+    }
+    queuedBuffers.clear();
 
     if (includeTrailingBreak)
       Stream << "\n";