detect circular macro expansion

- `MacroApplication` now detects any freestanding macro that appears on an expansion path more than once and throws `MacroExpansionError.circularExpansion` - added a test case in `ExpressionMacroTests`
swiftlang · Jul 31, 2024 · c7c9704 · c7c9704
1 parent 24a2501
commit c7c9704
Show file tree

Hide file tree

Showing 3 changed files with 109 additions and 0 deletions.
diff --git a/Sources/SwiftSyntaxMacroExpansion/MacroExpansion.swift b/Sources/SwiftSyntaxMacroExpansion/MacroExpansion.swift
@@ -63,6 +63,7 @@ enum MacroExpansionError: Error, CustomStringConvertible {
   case noFreestandingMacroRoles(Macro.Type)
   case moreThanOneBodyMacro
   case preambleWithoutBody
+  case circularExpansion(Macro.Type, any FreestandingMacroExpansionSyntax)
 
   var description: String {
     switch self {
@@ -92,6 +93,9 @@ enum MacroExpansionError: Error, CustomStringConvertible {
 
     case .preambleWithoutBody:
       return "preamble macro cannot be applied to a function with no body"
+
+    case .circularExpansion(let type, let syntax):
+      return "circular expansion detected: '\(syntax)' with macro implementation type '\(type)'"
     }
   }
 }

diff --git a/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift b/Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift
@@ -667,6 +667,9 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
   /// added to top-level 'CodeBlockItemList'.
   var extensions: [CodeBlockItemSyntax] = []
 
+  /// Stores the type names of the freestanding macros ithat are currently expanding.
+  var expandingFreestandingMacros: [String] = []
+
   init(
     macroSystem: MacroSystem,
     contextGenerator: @escaping (Syntax) -> Context,
@@ -687,6 +690,11 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
       return nil
     }
 
+    let macroCount = expandingFreestandingMacros.count
+    defer {
+      expandingFreestandingMacros.removeLast(expandingFreestandingMacros.count - macroCount)
+    }
+
     // Expand 'MacroExpansionExpr'.
     // Note that 'MacroExpansionExpr'/'MacroExpansionExprDecl' at code item
     // position are handled by 'visit(_:CodeBlockItemListSyntax)'.
@@ -727,6 +735,11 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
   func visitBodyAndPreambleMacros<Node: DeclSyntaxProtocol & WithOptionalCodeBlockSyntax>(
     _ node: Node
   ) -> Node {
+    let macroCount = expandingFreestandingMacros.count
+    defer {
+      expandingFreestandingMacros.removeLast(expandingFreestandingMacros.count - macroCount)
+    }
+
     // Expand preamble macros into a set of code items.
     let preamble = expandMacros(
       attachedTo: DeclSyntax(node),
@@ -790,6 +803,11 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
   }
 
   override func visit(_ node: CodeBlockItemListSyntax) -> CodeBlockItemListSyntax {
+    let macroCount = expandingFreestandingMacros.count
+    defer {
+      expandingFreestandingMacros.removeLast(expandingFreestandingMacros.count - macroCount)
+    }
+
     var newItems: [CodeBlockItemSyntax] = []
     func addResult(_ node: CodeBlockItemSyntax) {
       // Expand freestanding macro.
@@ -831,6 +849,11 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
   }
 
   override func visit(_ node: MemberBlockSyntax) -> MemberBlockSyntax {
+    let macroCount = expandingFreestandingMacros.count
+    defer {
+      expandingFreestandingMacros.removeLast(expandingFreestandingMacros.count - macroCount)
+    }
+
     let parentDeclGroup = node
       .parent?
       .as(DeclSyntax.self)
@@ -917,6 +940,11 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
   }
 
   override func visit(_ node: VariableDeclSyntax) -> DeclSyntax {
+    let macroCount = expandingFreestandingMacros.count
+    defer {
+      expandingFreestandingMacros.removeLast(expandingFreestandingMacros.count - macroCount)
+    }
+
     var node = super.visit(node).cast(VariableDeclSyntax.self)
 
     guard !macroAttributes(attachedTo: DeclSyntax(node), ofType: AccessorMacro.Type.self).isEmpty else {
@@ -948,6 +976,11 @@ private class MacroApplication<Context: MacroExpansionContext>: SyntaxRewriter {
   }
 
   override func visit(_ node: SubscriptDeclSyntax) -> DeclSyntax {
+    let macroCount = expandingFreestandingMacros.count
+    defer {
+      expandingFreestandingMacros.removeLast(expandingFreestandingMacros.count - macroCount)
+    }
+
     var node = super.visit(node).cast(SubscriptDeclSyntax.self)
     node.accessorBlock = expandAccessors(of: node, existingAccessors: node.accessorBlock).accessors
     return DeclSyntax(node)
@@ -1226,7 +1259,14 @@ extension MacroApplication {
     else {
       return .notAMacro
     }
+
     do {
+      let macroTypeName = "\(macro)"
+      guard !expandingFreestandingMacros.contains(macroTypeName) else {
+        throw MacroExpansionError.circularExpansion(macro, node)
+      }
+      expandingFreestandingMacros.append(macroTypeName)
+
       if let expanded = try expandMacro(macro, node) {
         return .success(expanded)
       } else {

diff --git a/Tests/SwiftSyntaxMacroExpansionTest/ExpressionMacroTests.swift b/Tests/SwiftSyntaxMacroExpansionTest/ExpressionMacroTests.swift
@@ -37,6 +37,44 @@ fileprivate struct StringifyMacro: ExpressionMacro {
   }
 }
 
+private struct InfiniteRecursionMacro: ExpressionMacro {
+  static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  )
+    throws -> ExprSyntax
+  {
+    let i = node.arguments.first!.expression.cast(IntegerLiteralExprSyntax.self).representedLiteralValue!
+    if i > 0 {
+      return "1 + #infiniteRecursion(i: \(raw: i - 1))"
+    } else {
+      return "0"
+    }
+  }
+}
+
+private struct Nested1RecursionMacro: ExpressionMacro {
+  static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  )
+    throws -> ExprSyntax
+  {
+    "(#nested2, #nested2, #nested2)"
+  }
+}
+
+private struct Nested2RecursionMacro: ExpressionMacro {
+  static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  )
+    throws -> ExprSyntax
+  {
+    "0"
+  }
+}
+
 final class ExpressionMacroTests: XCTestCase {
   private let indentationWidth: Trivia = .spaces(2)
 
@@ -292,4 +330,31 @@ final class ExpressionMacroTests: XCTestCase {
       macros: ["test": DiagnoseFirstArgument.self]
     )
   }
+
+  func testDetectCircularExpansion() {
+    assertMacroExpansion(
+      "#infiniteRecursion(i: 1000)",
+      expandedSource: "1 + #infiniteRecursion(i: 999)",
+      diagnostics: [
+        DiagnosticSpec(
+          message:
+            "circular expansion detected: '#infiniteRecursion(i: 999)' with macro implementation type 'InfiniteRecursionMacro'",
+          line: 1,
+          column: 5
+        )
+      ],
+      macros: [
+        "infiniteRecursion": InfiniteRecursionMacro.self
+      ]
+    )
+
+    assertMacroExpansion(
+      "#nested1",
+      expandedSource: "(0, 0, 0)",
+      macros: [
+        "nested1": Nested1RecursionMacro.self,
+        "nested2": Nested2RecursionMacro.self,
+      ]
+    )
+  }
 }