[Macros] Prevent recursive expansion of macros.

Author: DougGregor

include/swift/AST/DiagnosticsSema.def

@@ -6799,6 +6799,8 @@ ERROR(macro_expansion_missing_arguments,none,
       "expansion of macro %0 requires arguments", (DeclName))
 ERROR(macro_unsupported,none,
       "macros are not supported in this compiler", ())
+ERROR(macro_recursive,none,
+      "recursive expansion of macro %0", (DeclName))
 
 //------------------------------------------------------------------------------
 // MARK: Move Only Errors

lib/Sema/TypeCheckMacros.cpp

@@ -121,6 +121,34 @@ MacroDefinition MacroDefinitionRequest::evaluate(
       ctx, macro->externalModuleName, macro->externalMacroTypeName);
 }
 
+/// Determine whether the given source file is from an expansion of the given
+/// macro.
+static bool isFromExpansionOfMacro(SourceFile *sourceFile, MacroDecl *macro) {
+  while (sourceFile) {
+    auto expansion = sourceFile->getMacroExpansion();
+    if (!expansion)
+      return false;
+
+    if (auto expansionExpr = dyn_cast_or_null<MacroExpansionExpr>(
+            expansion.dyn_cast<Expr *>())) {
+      if (expansionExpr->getMacroRef().getDecl() == macro)
+        return true;
+    } else if (auto expansionDecl = dyn_cast_or_null<MacroExpansionDecl>(
+            expansion.dyn_cast<Decl *>())) {
+      // FIXME: Update once MacroExpansionDecl has a proper macro reference
+      // in it.
+      if (expansionDecl->getMacro().getFullName() == macro->getName())
+        return true;
+    } else {
+      llvm_unreachable("Unknown macro expansion node kind");
+    }
+
+    sourceFile = sourceFile->getEnclosingSourceFile();
+  }
+
+  return false;
+}
+
 Expr *swift::expandMacroExpr(
     DeclContext *dc, Expr *expr, ConcreteDeclRef macroRef, Type expandedType
 ) {
@@ -137,6 +165,11 @@ Expr *swift::expandMacroExpr(
 
   MacroDecl *macro = cast<MacroDecl>(macroRef.getDecl());
 
+  if (isFromExpansionOfMacro(sourceFile, macro)) {
+    ctx.Diags.diagnose(expr->getLoc(), diag::macro_recursive, macro->getName());
+    return nullptr;
+  }
+
   auto macroDef = evaluateOrDefault(
       ctx.evaluator, MacroDefinitionRequest{macro},
       MacroDefinition::forUndefined());

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -154,3 +154,16 @@ public struct AddBlocker: ExpressionMacro {
     return result.as(MacroExpansionExprSyntax.self)!.argumentList.first!.expression
   }
 }
+
+public struct RecursiveMacro: ExpressionMacro {
+  public static func expansion(
+    of macro: MacroExpansionExprSyntax, in context: inout MacroExpansionContext
+  ) -> ExprSyntax {
+    guard let argument = macro.argumentList.first?.expression,
+          argument.description == "false" else {
+      return ExprSyntax(macro)
+    }
+
+    return "()"
+  }
+}

test/Macros/macro_expand.swift

@@ -11,7 +11,8 @@
 
 @expression macro customFileID: String = MacroDefinition.FileIDMacro
 @expression macro stringify<T>(_ value: T) -> (T, String) = MacroDefinition.StringifyMacro
-@expression macro fileID<T: _ExpressibleByStringLitera>: T = MacroDefinition.FileIDMacro
+@expression macro fileID<T: _ExpressibleByStringLiteral>: T = MacroDefinition.FileIDMacro
+@expression macro recurse(_: Bool) = MacroDefinition.RecursiveMacro
 
 func testFileID(a: Int, b: Int) {
   // CHECK: MacroUser/macro_expand.swift
@@ -73,5 +74,9 @@ func testAddBlocker(a: Int, b: Int, c: Int, oa: OnlyAdds) {
   _ = #addBlocker(oa + oa) // expected-error{{blocked an add; did you mean to subtract? (from macro 'addBlocker')}}
   // expected-note@-1{{in expansion of macro 'addBlocker' here}}
   // expected-note@-2{{use '-'}}{{22-23=-}}
+
+  // Check recursion.
+  #recurse(false) // okay
+  #recurse(true) // expected-note{{in expansion of macro 'recurse' here}}
 #endif
 }