Use two-step dispatch to SyntaxVisitor.vistImpl

`SyntaxRewriter` has had a problem with stack allocation space for a while because cases in a switch statement don’t share stack space in debug builds (rdar://55929175). We started seeing this issue with `SyntaxVisitor` as well now. Use the same approach to fix the problem:

To circumvent this problem, make calling the specific visitation function a two-step process: First determine the function to call in this function and return a reference to it, then call it. This way, the stack frame that determines the correct visitation function will be popped of the stack before the function is being called, making the switch's stack space transient instead of having it linger in the call stack.

And to make the entire approach look sane, allow 4 arguments on the same line for SyntaxVisitor instead of the 3 that we allow in all other files.

Author: ahoppen

CodeGeneration/Sources/Utils/CodeGenerationFormat.swift

@@ -15,7 +15,12 @@ import SwiftSyntax
 
 /// A format style for files generated by CodeGeneration.
 public class CodeGenerationFormat: BasicFormat {
-  public init() {
+  /// The maximum number of array/dictionary/function parameters/tuple elements
+  /// that should be put on the same line.
+  private let maxElementsOnSameLine: Int
+
+  public init(maxElementsOnSameLine: Int = 3) {
+    self.maxElementsOnSameLine = maxElementsOnSameLine
     super.init(indentationWidth: .spaces(2))
   }
 
@@ -26,7 +31,7 @@ public class CodeGenerationFormat: BasicFormat {
   public override func visit(_ node: ArrayElementListSyntax) -> ArrayElementListSyntax {
     let children = node.children(viewMode: .all)
     // Short array literals are presented on one line, list each element on a different line.
-    if children.count > 3 {
+    if children.count > maxElementsOnSameLine {
       return ArrayElementListSyntax(formatChildrenSeparatedByNewline(children: children, elementType: ArrayElementSyntax.self))
     } else {
       return super.visit(node)
@@ -45,7 +50,7 @@ public class CodeGenerationFormat: BasicFormat {
   public override func visit(_ node: DictionaryElementListSyntax) -> DictionaryElementListSyntax {
     let children = node.children(viewMode: .all)
     // Short dictionary literals are presented on one line, list each element on a different line.
-    if children.count > 3 {
+    if children.count > maxElementsOnSameLine {
       return DictionaryElementListSyntax(formatChildrenSeparatedByNewline(children: children, elementType: DictionaryElementSyntax.self))
     } else {
       return super.visit(node)
@@ -55,7 +60,7 @@ public class CodeGenerationFormat: BasicFormat {
   public override func visit(_ node: FunctionParameterListSyntax) -> FunctionParameterListSyntax {
     let children = node.children(viewMode: .all)
     // Short function parameter literals are presented on one line, list each element on a different line.
-    if children.count > 3 {
+    if children.count > maxElementsOnSameLine {
       return FunctionParameterListSyntax(formatChildrenSeparatedByNewline(children: children, elementType: FunctionParameterSyntax.self))
     } else {
       return super.visit(node)
@@ -82,7 +87,7 @@ public class CodeGenerationFormat: BasicFormat {
   public override func visit(_ node: LabeledExprListSyntax) -> LabeledExprListSyntax {
     let children = node.children(viewMode: .all)
     // Short tuple element list literals are presented on one line, list each element on a different line.
-    if children.count > 3 {
+    if children.count > maxElementsOnSameLine {
       return LabeledExprListSyntax(formatChildrenSeparatedByNewline(children: children, elementType: LabeledExprSyntax.self))
     } else {
       return super.visit(node)

CodeGeneration/Sources/generate-swift-syntax/GenerateSwiftSyntax.swift

@@ -47,8 +47,8 @@ struct GeneratedFileSpec {
     self.contentsGenerator = contents
   }
 
-  init(_ pathComponents: [String], _ contents: @escaping @autoclosure () -> SourceFileSyntax) {
-    self.init(pathComponents, "\(contents().formatted(using: CodeGenerationFormat()))\n")
+  init(_ pathComponents: [String], _ contents: @escaping @autoclosure () -> SourceFileSyntax, format: CodeGenerationFormat = CodeGenerationFormat()) {
+    self.init(pathComponents, "\(contents().formatted(using: format))\n")
   }
 }
 
@@ -112,7 +112,7 @@ struct GenerateSwiftSyntax: ParsableCommand {
       GeneratedFileSpec(swiftSyntaxGeneratedDir + ["SyntaxRewriter.swift"], syntaxRewriterFile),
       GeneratedFileSpec(swiftSyntaxGeneratedDir + ["SyntaxTraits.swift"], syntaxTraitsFile),
       GeneratedFileSpec(swiftSyntaxGeneratedDir + ["SyntaxTransform.swift"], syntaxTransformFile),
-      GeneratedFileSpec(swiftSyntaxGeneratedDir + ["SyntaxVisitor.swift"], syntaxVisitorFile),
+      GeneratedFileSpec(swiftSyntaxGeneratedDir + ["SyntaxVisitor.swift"], syntaxVisitorFile, format: CodeGenerationFormat(maxElementsOnSameLine: 4)),
       GeneratedFileSpec(swiftSyntaxGeneratedDir + ["TokenKind.swift"], tokenKindFile),
       GeneratedFileSpec(swiftSyntaxGeneratedDir + ["Tokens.swift"], tokensFile),
       GeneratedFileSpec(swiftSyntaxGeneratedDir + ["TriviaPieces.swift"], triviaPiecesFile),

CodeGeneration/Sources/generate-swift-syntax/templates/swiftsyntax/SyntaxRewriterFile.swift

@@ -210,7 +210,7 @@ let syntaxRewriterFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
                 /// we need to switch through a huge switch statement that covers all syntax
                 /// types. In debug builds, the cases of this switch statement do not share
                 /// stack space (rdar://55929175). Because of this, the switch statement
-                /// requires allocates about 15KB of stack space. In scenarios with reduced
+                /// requires about 15KB of stack space. In scenarios with reduced
                 /// stack size (in particular dispatch queues), this often results in a stack
                 /// overflow during syntax tree rewriting.
                 ///

CodeGeneration/Sources/generate-swift-syntax/templates/swiftsyntax/SyntaxVisitorFile.swift

@@ -113,27 +113,108 @@ let syntaxVisitorFile = SourceFileSyntax(leadingTrivia: copyrightHeader) {
       """
     )
 
-    try FunctionDeclSyntax("private func visit(_ data: SyntaxData)") {
-      try SwitchExprSyntax("switch data.raw.kind") {
-        SwitchCaseSyntax("case .token:") {
-          DeclSyntax("let node = TokenSyntax(data)")
-
-          ExprSyntax("_ = visit(node)")
-          ExprSyntax(
-            """
-            // No children to visit.
-            visitPost(node)
-            """
+    try IfConfigDeclSyntax(
+      leadingTrivia:
+        """
+        // SwiftSyntax requires a lot of stack space in debug builds for syntax tree
+        // visitation. In scenarios with reduced stack space (in particular dispatch
+        // queues), this easily results in a stack overflow. To work around this issue,
+        // use a less performant but also less stack-hungry version of SwiftSyntax's
+        // SyntaxVisitor in debug builds.
+
+        """,
+      clauses: IfConfigClauseListSyntax {
+        IfConfigClauseSyntax(
+          poundKeyword: .poundIfToken(),
+          condition: ExprSyntax("DEBUG"),
+          elements: .statements(
+            try CodeBlockItemListSyntax {
+              try FunctionDeclSyntax(
+                """
+                /// Implementation detail of visit(_:). Do not call directly.
+                ///
+                /// Returns the function that shall be called to visit a specific syntax node.
+                ///
+                /// To determine the correct specific visitation function for a syntax node,
+                /// we need to switch through a huge switch statement that covers all syntax
+                /// types. In debug builds, the cases of this switch statement do not share
+                /// stack space (rdar://55929175). Because of this, the switch statement
+                /// requires about 15KB of stack space. In scenarios with reduced
+                /// stack size (in particular dispatch queues), this often results in a stack
+                /// overflow during syntax tree rewriting.
+                ///
+                /// To circumvent this problem, make calling the specific visitation function
+                /// a two-step process: First determine the function to call in this function
+                /// and return a reference to it, then call it. This way, the stack frame
+                /// that determines the correct visitation function will be popped of the
+                /// stack before the function is being called, making the switch's stack
+                /// space transient instead of having it linger in the call stack.
+                private func visitationFunc(for data: SyntaxData) -> ((SyntaxData) -> Void)
+                """
+              ) {
+                try SwitchExprSyntax("switch data.raw.kind") {
+                  SwitchCaseSyntax("case .token:") {
+                    StmtSyntax(
+                      """
+                      return {
+                        let node = TokenSyntax($0)
+                        _ = self.visit(node)
+                        // No children to visit.
+                        self.visitPost(node)
+                      }
+                      """
+                    )
+                  }
+
+                  for node in NON_BASE_SYNTAX_NODES {
+                    SwitchCaseSyntax("case .\(node.varOrCaseName):") {
+                      StmtSyntax("return { self.visitImpl($0, \(node.kind.syntaxType).self, self.visit, self.visitPost) }")
+                    }
+                  }
+                }
+              }
+
+              DeclSyntax(
+                """
+                private func visit(_ data: SyntaxData) {
+                  return visitationFunc(for: data)(data)
+                }
+                """
+              )
+            }
           )
-        }
+        )
+        IfConfigClauseSyntax(
+          poundKeyword: .poundElseToken(),
+          elements: .statements(
+            CodeBlockItemListSyntax {
+              try! FunctionDeclSyntax("private func visit(_ data: SyntaxData)") {
+                try SwitchExprSyntax("switch data.raw.kind") {
+                  SwitchCaseSyntax("case .token:") {
+                    DeclSyntax("let node = TokenSyntax(data)")
 
-        for node in NON_BASE_SYNTAX_NODES {
-          SwitchCaseSyntax("case .\(node.varOrCaseName):") {
-            ExprSyntax("visitImpl(data, \(node.kind.syntaxType).self, visit, visitPost)")
-          }
-        }
+                    ExprSyntax("_ = visit(node)")
+                    ExprSyntax(
+                      """
+                      // No children to visit.
+                      visitPost(node)
+                      """
+                    )
+                  }
+
+                  for node in NON_BASE_SYNTAX_NODES {
+                    SwitchCaseSyntax("case .\(node.varOrCaseName):") {
+                      ExprSyntax("visitImpl(data, \(node.kind.syntaxType).self, visit, visitPost)")
+                    }
+                  }
+                }
+              }
+
+            }
+          )
+        )
       }
-    }
+    )
 
     DeclSyntax(
       """

Sources/SwiftSyntax/generated/SyntaxRewriter.swift

@@ -2073,7 +2073,7 @@ open class SyntaxRewriter {
   /// we need to switch through a huge switch statement that covers all syntax
   /// types. In debug builds, the cases of this switch statement do not share
   /// stack space (rdar://55929175). Because of this, the switch statement
-  /// requires allocates about 15KB of stack space. In scenarios with reduced
+  /// requires about 15KB of stack space. In scenarios with reduced
   /// stack size (in particular dispatch queues), this often results in a stack
   /// overflow during syntax tree rewriting.
   ///