Emit "in expansion of macro" notes for diagnostics within macro expansion buffers

Author: DougGregor

include/swift/AST/DiagnosticEngine.h

@@ -502,6 +502,7 @@ namespace swift {
     }
 
     void addChildNote(Diagnostic &&D);
+    void insertChildNote(unsigned beforeIndex, Diagnostic &&D);
   };
 
   /// Describes an in-flight diagnostic, which is currently active
@@ -1164,6 +1165,11 @@ namespace swift {
     /// Send \c diag to all diagnostic consumers.
     void emitDiagnostic(const Diagnostic &diag);
 
+    /// Retrieve the set of child notes that describe how the generated
+    /// source buffer was derived, e.g., a macro expansion backtrace.
+    std::vector<Diagnostic> getGeneratedSourceBufferNotes(
+        SourceLoc loc, Optional<unsigned> &lastBufferID);
+
     /// Handle a new diagnostic, which will either be emitted, or added to an
     /// active transaction.
     void handleDiagnostic(Diagnostic &&diag);

include/swift/AST/DiagnosticsCommon.def

@@ -214,6 +214,12 @@ REMARK(remark_save_cache, none,
 ERROR(unknown_attribute,none,
       "unknown attribute '%0'", (StringRef))
 
+//------------------------------------------------------------------------------
+// MARK: macro diagnostics
+//------------------------------------------------------------------------------
+NOTE(in_macro_expansion,none,
+     "in expansion of macro %0 here", (DeclName))
+
 //------------------------------------------------------------------------------
 // MARK: bridged diagnostics
 //------------------------------------------------------------------------------

lib/AST/DiagnosticEngine.cpp

@@ -21,6 +21,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/DiagnosticSuppression.h"
 #include "swift/AST/DiagnosticsCommon.h"
+#include "swift/AST/Expr.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/Pattern.h"
 #include "swift/AST/PrintOptions.h"
@@ -1267,10 +1268,82 @@ DiagnosticEngine::diagnosticInfoForDiagnostic(const Diagnostic &diagnostic) {
       diagnostic.isChildNote());
 }
 
+std::vector<Diagnostic> DiagnosticEngine::getGeneratedSourceBufferNotes(
+    SourceLoc loc, Optional<unsigned> &lastBufferID
+) {
+  // The set of child notes we're building up.
+  std::vector<Diagnostic> childNotes;
+
+  // If the location is invalid, there's nothing to do.
+  if (loc.isInvalid())
+    return childNotes;
+
+  // If we already emitted these notes for a prior part of the diagnostic,
+  // don't do so again.
+  auto currentBufferID = SourceMgr.findBufferContainingLoc(loc);
+  if (currentBufferID == lastBufferID)
+    return childNotes;
+
+  // Keep track of the last buffer ID we considered.
+  lastBufferID = currentBufferID;
+
+  SourceLoc currentLoc = loc;
+  do {
+    auto generatedInfo = SourceMgr.getGeneratedSourceInfo(currentBufferID);
+    if (!generatedInfo)
+      return childNotes;
+
+    ASTNode expansionNode =
+        ASTNode::getFromOpaqueValue(generatedInfo->astNode);
+
+    switch (generatedInfo->kind) {
+    case GeneratedSourceInfo::MacroExpansion: {
+      SourceRange origRange = expansionNode.getSourceRange();
+      DeclName macroName;
+      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->getMacro().getFullName();
+      }
+
+      Diagnostic expansionNote(diag::in_macro_expansion, macroName);
+      expansionNote.setLoc(origRange.Start);
+      expansionNote.addRange(
+          Lexer::getCharSourceRangeFromSourceRange(SourceMgr, origRange));
+      childNotes.push_back(std::move(expansionNote));
+      break;
+    }
+
+    case GeneratedSourceInfo::ReplacedFunctionBody:
+      return childNotes;
+    }
+
+    // Walk up the stack.
+    currentLoc = expansionNode.getStartLoc();
+    currentBufferID = SourceMgr.findBufferContainingLoc(currentLoc);
+  } while (true);
+}
+
 void DiagnosticEngine::emitDiagnostic(const Diagnostic &diagnostic) {
+  Optional<unsigned> lastBufferID;
+
+  ArrayRef<Diagnostic> childNotes = diagnostic.getChildNotes();
+  std::vector<Diagnostic> extendedChildNotes;
+
   if (auto info = diagnosticInfoForDiagnostic(diagnostic)) {
+    // If the diagnostic location is within a buffer containing generated
+    // source code, add child notes showing where the generation occurred.
+    extendedChildNotes = getGeneratedSourceBufferNotes(info->Loc, lastBufferID);
+    if (!extendedChildNotes.empty()) {
+      extendedChildNotes.insert(extendedChildNotes.end(),
+                                childNotes.begin(), childNotes.end());
+      childNotes = extendedChildNotes;
+    }
+
     SmallVector<DiagnosticInfo, 1> childInfo;
-    auto childNotes = diagnostic.getChildNotes();
     for (unsigned i : indices(childNotes)) {
       auto child = diagnosticInfoForDiagnostic(childNotes[i]);
       assert(child);
@@ -1302,7 +1375,7 @@ void DiagnosticEngine::emitDiagnostic(const Diagnostic &diagnostic) {
   // For compatibility with DiagnosticConsumers which don't know about child
   // notes. These can be ignored by consumers which do take advantage of the
   // grouping.
-  for (auto &childNote : diagnostic.getChildNotes())
+  for (auto &childNote : childNotes)
     emitDiagnostic(childNote);
 }
 

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -84,16 +84,12 @@ public struct AddBlocker: ExpressionMacro {
     // Link the folded argument back into the tree.
     var node = node.withArgumentList(node.argumentList.replacing(childAt: 0, with: node.argumentList.first!.withExpression(foldedArgument.as(ExprSyntax.self)!)))
 
-   class AddVisitor: SyntaxVisitor {
+   class AddVisitor: SyntaxRewriter {
       var diagnostics: [Diagnostic] = []
 
-      init() {
-        super.init(viewMode: .sourceAccurate)
-      }
-
       override func visit(
         _ node: InfixOperatorExprSyntax
-      ) -> SyntaxVisitorContinueKind {
+      ) -> ExprSyntax {
         if let binOp = node.operatorOperand.as(BinaryOperatorExprSyntax.self) {
           if binOp.operatorToken.text == "+" {
             let messageID = MessageID(domain: "silly", id: "addblock")
@@ -131,20 +127,30 @@ public struct AddBlocker: ExpressionMacro {
                 ]
               )
             )
+
+            return ExprSyntax(
+              node.withOperatorOperand(
+                ExprSyntax(
+                  binOp.withOperatorToken(
+                    binOp.operatorToken.withKind(.spacedBinaryOperator("-"))
+                  )
+                )
+              )
+            )
           }
         }
 
-        return .visitChildren
+        return ExprSyntax(node)
       }
     }
 
     let visitor = AddVisitor()
-    visitor.walk(Syntax(node))
+    let result = visitor.visit(Syntax(node))
 
     for diag in visitor.diagnostics {
       context.diagnose(diag)
     }
 
-    return argument
+    return result.as(MacroExpansionExprSyntax.self)!.argumentList.first!.expression
   }
 }

test/Macros/macro_expand.swift

@@ -50,9 +50,15 @@ testStringify(a: 1, b: 1)
 
 macro addBlocker<T>(_ value: T) -> T = MacroDefinition.AddBlocker
 
-func testAddBlocker(a: Int, b: Int, c: Int) {
+struct OnlyAdds {
+  static func +(lhs: OnlyAdds, rhs: OnlyAdds) -> OnlyAdds { lhs }
+}
+
+func testAddBlocker(a: Int, b: Int, c: Int, oa: OnlyAdds) {
   _ = #addBlocker(a * b * c)
 #if TEST_DIAGNOSTICS
   _ = #addBlocker(a + b * c) // expected-error{{blocked an add; did you mean to subtract? (from macro 'addBlocker')}}{{21-22=-}}
+  _ = #addBlocker(oa + oa) // expected-error{{blocked an add; did you mean to subtract? (from macro 'addBlocker')}}{{22-23=-}}
+      // expected-note@-1{{in expansion of macro 'addBlocker' here}}
 #endif
 }