Implement depth-limited rewriting

Author: grynspan

Sources/TestingMacros/Support/AttributeDiscovery.swift

@@ -126,7 +126,7 @@ struct AttributeInfo {
       var expr = argument.expression
       if argument.expression.tokens(viewMode: .sourceAccurate).map(\.tokenKind).contains(.keyword(.Self)) {
         let selfRemover = _SelfRemover(in: context)
-        expr = selfRemover.rewrite(Syntax(argument.expression)).cast(ExprSyntax.self)
+        expr = selfRemover.rewrite(Syntax(argument.expression), detach: true).cast(ExprSyntax.self)
       }
       return Argument(label: argument.label, expression: expr)
     }

Sources/TestingMacros/Support/ConditionArgumentParsing.swift

@@ -91,6 +91,12 @@ func removeParentheses(from expr: ExprSyntax) -> ExprSyntax? {
 
 // MARK: - Inserting expression context callouts
 
+/// The maximum value of `_rewriteDepth` allowed by `_rewrite()` before it will
+/// start bailing early.
+private let _maximumRewriteDepth = {
+  Int.max // disable rewrite-limiting (need to evaluate possible heuristics)
+}()
+
 /// A type that inserts calls to an `__ExpectationContext` instance into an
 /// expression's syntax tree.
 private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansionContext, M: FreestandingMacroExpansionSyntax {
@@ -123,6 +129,12 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
     super.init()
   }
 
+  /// The number of calls to `_rewrite()` made along the current node hierarchy.
+  ///
+  /// This value is incremented with each call to `_rewrite()` and managed by
+  /// `_visitChild()`.
+  private var _rewriteDepth = 0
+
   /// Rewrite a given syntax node by inserting a call to the expression context
   /// (or rather, its `callAsFunction(_:_:)` member).
   ///
@@ -137,14 +149,27 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
   ///
   /// - Returns: A rewritten copy of `node` that calls into the expression
   ///   context when it is evaluated at runtime.
-  private func _rewrite(_ node: some ExprSyntaxProtocol, originalWas originalNode: some ExprSyntaxProtocol, calling functionName: TokenSyntax? = nil, passing additionalArguments: [Argument] = []) -> ExprSyntax {
-    guard rewrittenNodes.insert(Syntax(originalNode)).inserted else {
-      // If this node has already been rewritten, we don't need to rewrite it
-      // again. (Currently, this can only happen when expanding binary operators
-      // which need a bit of extra help.)
-      return ExprSyntax(node)
+  private func _rewrite(
+    _ node: @autoclosure () -> some ExprSyntaxProtocol,
+    originalWas originalNode: @autoclosure () -> some ExprSyntaxProtocol,
+    calling functionName: @autoclosure () -> TokenSyntax? = nil,
+    passing additionalArguments: @autoclosure () -> [Argument] = []
+  ) -> ExprSyntax {
+    _rewriteDepth += 1
+    if _rewriteDepth > _maximumRewriteDepth {
+      // At least 2 ancestors of this node have already been rewritten, so do
+      // not recursively rewrite further. This is necessary to limit the added
+      // exponentional complexity we're throwing at the type checker.
+      return ExprSyntax(originalNode())
     }
 
+    // We're going to rewrite the node, so we'll evaluate the arguments now.
+    let node = node()
+    let originalNode = originalNode()
+    let functionName = functionName()
+    let additionalArguments = additionalArguments()
+    rewrittenNodes.insert(Syntax(originalNode))
+
     let calledExpr: ExprSyntax = if let functionName {
       ExprSyntax(MemberAccessExprSyntax(base: expressionContextNameExpr, name: functionName))
     } else {
@@ -200,6 +225,43 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
     _rewrite(node, originalWas: node, calling: functionName, passing: additionalArguments)
   }
 
+  /// Visit `node` as a child of another previously-visited node.
+  ///
+  /// - Parameters:
+  ///   - node: The node to visit.
+  ///
+  /// - Returns: `node`, or a modified copy thereof if `node` or a child node
+  ///   was rewritten.
+  ///
+  /// Use this function instead of calling `visit(_:)` or `rewrite(_:detach:)`
+  /// recursively.
+  ///
+  /// This overload simply visits `node` and is used for nodes that cannot be
+  /// rewritten directly (because they are not expressions.)
+  @_disfavoredOverload
+  private func _visitChild<S>(_ node: S) -> S where S: SyntaxProtocol {
+    rewrite(node, detach: true).cast(S.self)
+  }
+
+  /// Visit `node` as a child of another previously-visited node.
+  ///
+  /// - Parameters:
+  ///   - node: The node to visit.
+  ///
+  /// - Returns: `node`, or a modified copy thereof if `node` or a child node
+  ///   was rewritten.
+  ///
+  /// Use this function instead of calling `visit(_:)` or `rewrite(_:detach:)`
+  /// recursively.
+  private func _visitChild(_ node: some ExprSyntaxProtocol) -> ExprSyntax {
+    let oldRewriteDepth = _rewriteDepth
+    defer {
+      _rewriteDepth = oldRewriteDepth
+    }
+
+    return rewrite(node, detach: true).cast(ExprSyntax.self)
+  }
+
   /// Whether or not the parent node of the given node is capable of containing
   /// a rewritten `DeclReferenceExprSyntax` instance.
   ///
@@ -281,7 +343,7 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
       return _rewrite(
         TupleExprSyntax {
           for element in node.elements {
-            visit(element).trimmed
+            _visitChild(element).trimmed
           }
         },
         originalWas: node
@@ -302,28 +364,28 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
     // expressions can be directly extracted out.
     if _isParentOfDeclReferenceExprValidForRewriting(node) {
       return _rewrite(
-        node.with(\.base, node.base.map(visit)),
+        node.with(\.base, node.base.map(_visitChild)),
         originalWas: node
       )
     }
 
-    return ExprSyntax(node.with(\.base, node.base.map(visit)))
+    return ExprSyntax(node.with(\.base, node.base.map(_visitChild)))
   }
 
   override func visit(_ node: FunctionCallExprSyntax) -> ExprSyntax {
     _rewrite(
       node
-        .with(\.calledExpression, visit(node.calledExpression))
-        .with(\.arguments, visit(node.arguments)),
+        .with(\.calledExpression, _visitChild(node.calledExpression))
+        .with(\.arguments, _visitChild(node.arguments)),
       originalWas: node
     )
   }
 
   override func visit(_ node: SubscriptCallExprSyntax) -> ExprSyntax {
     _rewrite(
       node
-        .with(\.calledExpression, visit(node.calledExpression))
-        .with(\.arguments, visit(node.arguments)),
+        .with(\.calledExpression, _visitChild(node.calledExpression))
+        .with(\.arguments, _visitChild(node.arguments)),
       originalWas: node
     )
   }
@@ -355,7 +417,7 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
 
     return _rewrite(
       node
-        .with(\.expression, visit(node.expression)),
+        .with(\.expression, _visitChild(node.expression)),
       originalWas: node
     )
   }
@@ -377,18 +439,18 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
         originalWas: node,
         calling: .identifier("__cmp"),
         passing: [
-          Argument(expression: visit(node.leftOperand)),
+          Argument(expression: _visitChild(node.leftOperand)),
           Argument(expression: node.leftOperand.expressionID(rootedAt: effectiveRootNode)),
-          Argument(expression: visit(node.rightOperand)),
+          Argument(expression: _visitChild(node.rightOperand)),
           Argument(expression: node.rightOperand.expressionID(rootedAt: effectiveRootNode))
         ]
       )
     }
 
     return _rewrite(
       node
-        .with(\.leftOperand, visit(node.leftOperand))
-        .with(\.rightOperand, visit(node.rightOperand)),
+        .with(\.leftOperand, _visitChild(node.leftOperand))
+        .with(\.rightOperand, _visitChild(node.rightOperand)),
       originalWas: node
     )
   }
@@ -399,12 +461,11 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
     // `inout`, so it should be sufficient to capture it in a `defer` statement
     // that runs after the expression is evaluated.
 
-    let teardownItem = CodeBlockItemSyntax(
-      item: .expr(
-        _rewrite(node.expression, calling: .identifier("__inoutAfter"))
-      )
-    )
-    teardownItems.append(teardownItem)
+    let rewrittenExpr = _rewrite(node.expression, calling: .identifier("__inoutAfter"))
+    if rewrittenExpr != ExprSyntax(node.expression) {
+      let teardownItem = CodeBlockItemSyntax(item: .expr(rewrittenExpr))
+      teardownItems.append(teardownItem)
+    }
 
     // The argument should not be expanded in-place as we can't return an
     // argument passed `inout` and expect it to remain semantically correct.
@@ -427,7 +488,7 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
     rewrittenNodes.insert(Syntax(type))
 
     return _rewrite(
-      visit(valueExpr).trimmed,
+      _visitChild(valueExpr).trimmed,
       originalWas: originalNode,
       calling: .identifier("__\(isAsKeyword)"),
       passing: [
@@ -503,7 +564,7 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
       node.with(
         \.elements, ArrayElementListSyntax {
           for element in node.elements {
-            ArrayElementSyntax(expression: visit(element.expression).trimmed)
+            ArrayElementSyntax(expression: _visitChild(element.expression).trimmed)
           }
         }
       ),
@@ -520,7 +581,7 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
         \.content, .elements(
           DictionaryElementListSyntax {
             for element in elements {
-              DictionaryElementSyntax(key: visit(element.key).trimmed, value: visit(element.value).trimmed)
+              DictionaryElementSyntax(key: _visitChild(element.key).trimmed, value: _visitChild(element.value).trimmed)
             }
           }
         )
@@ -570,7 +631,7 @@ extension ConditionMacro {
     _diagnoseTrivialBooleanValue(from: ExprSyntax(node), for: macro, in: context)
 
     let contextInserter = _ContextInserter(in: context, for: macro, rootedAt: Syntax(effectiveRootNode), expressionContextName: expressionContextName)
-    var expandedExpr = contextInserter.rewrite(node).cast(ExprSyntax.self)
+    var expandedExpr = contextInserter.rewrite(node, detach: true).cast(ExprSyntax.self)
     let rewrittenNodes = contextInserter.rewrittenNodes
 
     // Insert additional effect keywords/thunks as needed.
@@ -606,7 +667,7 @@ extension ConditionMacro {
     var captureList: ClosureCaptureClauseSyntax?
     do {
       let dollarIDReplacer = _DollarIdentifierReplacer()
-      codeBlockItems = dollarIDReplacer.rewrite(codeBlockItems).cast(CodeBlockItemListSyntax.self)
+      codeBlockItems = dollarIDReplacer.rewrite(codeBlockItems, detach: true).cast(CodeBlockItemListSyntax.self)
       if !dollarIDReplacer.dollarIdentifierTokenKinds.isEmpty {
         let dollarIdentifierTokens = dollarIDReplacer.dollarIdentifierTokenKinds.map { tokenKind in
           TokenSyntax(tokenKind, presence: .present)

Tests/TestingTests/Support/CartesianProductTests.swift

@@ -33,7 +33,7 @@ struct CartesianProductTests {
     #expect(product.underestimatedCount == c1.underestimatedCount * c2.underestimatedCount)
     let productCount = Array(product).count
     #expect(productCount == c1.count * c2.count)
-    #expect(productCount == (26 * 100) as Int)
+    #expect(productCount == 26 * 100)
   }
 
   @Test("First element is correct")