Simplify SwiftBasicFormat by introducing structural information of the syntax tree as static properties

CodeGeneration/Sources/generate-swiftbasicformat/BasicFormatFile.swift

@@ -22,21 +22,151 @@ let basicFormatFile = SourceFile {
     path: [AccessPathComponent(name: "SwiftSyntax")]
   )
 
-  ClassDecl(modifiers: [DeclModifier(name: .open)], identifier: "BasicFormat", inheritanceClause: TypeInheritanceClause { InheritedType(typeName: Type("SyntaxRewriter")) }) {
+  ClassDecl("open class BasicFormat: SyntaxRewriter") {
     VariableDecl("public var indentationLevel: Int = 0")
     VariableDecl("open var indentation: TriviaPiece { .spaces(indentationLevel * 4) }")
     VariableDecl("public var indentedNewline: Trivia { Trivia(pieces: [.newlines(1), indentation]) }")
     VariableDecl("private var lastRewrittenToken: TokenSyntax?")
+    VariableDecl("private var putNextTokenOnNewLine: Bool = false")
 
-    for node in SYNTAX_NODES where !node.isBase {
-      if node.isSyntaxCollection {
-        makeSyntaxCollectionRewriteFunc(node: node)
-      } else {
-        makeLayoutNodeRewriteFunc(node: node)
+    FunctionDecl("""
+      open override func visitPre(_ node: Syntax) {
+        if let keyPath = getKeyPath(node), shouldIndent(keyPath) {
+          indentationLevel += 1
+        }
+        if let parent = node.parent, childrenSeparatedByNewline(parent) {
+          putNextTokenOnNewLine = true
+        }
+      }
+      """
+    )
+    FunctionDecl("""
+      open override func visitPost(_ node: Syntax) {
+        if let keyPath = getKeyPath(node), shouldIndent(keyPath) {
+          indentationLevel -= 1
+        }
+      }
+      """
+    )
+
+    FunctionDecl("""
+      open override func visit(_ node: TokenSyntax) -> TokenSyntax {
+        var leadingTrivia = node.leadingTrivia
+        var trailingTrivia = node.trailingTrivia
+        if requiresLeadingSpace(node.tokenKind) && leadingTrivia.isEmpty && lastRewrittenToken?.trailingTrivia.isEmpty != false {
+          leadingTrivia += .space
+        }
+        if requiresTrailingSpace(node.tokenKind) && trailingTrivia.isEmpty {
+          trailingTrivia += .space
+        }
+        if let keyPath = getKeyPath(Syntax(node)), requiresLeadingNewline(keyPath), !(leadingTrivia.first?.isNewline ?? false) {
+          leadingTrivia = .newline + leadingTrivia
+        }
+        leadingTrivia = leadingTrivia.indented(indentation: indentation)
+        trailingTrivia = trailingTrivia.indented(indentation: indentation)
+        let rewritten = TokenSyntax(
+          node.tokenKind,
+          leadingTrivia: leadingTrivia,
+          trailingTrivia: trailingTrivia,
+          presence: node.presence
+        )
+        lastRewrittenToken = rewritten
+        putNextTokenOnNewLine = false
+        return rewritten
+      }
+      """
+    )
+
+    FunctionDecl("open func shouldIndent(_ keyPath: AnyKeyPath) -> Bool") {
+      SwitchStmt(expression: Expr("keyPath")) {
+        for node in SYNTAX_NODES where !node.isBase {
+          for child in node.children where child.isIndented {
+            SwitchCase("case \\\(node.type.syntaxBaseName).\(child.swiftName):") {
+              ReturnStmt("return true")
+            }
+          }
+        }
+        SwitchCase("default:") {
+          ReturnStmt("return false")
+        }
+      }
+    }
+
+    FunctionDecl("open func requiresLeadingNewline(_ keyPath: AnyKeyPath) -> Bool") {
+      SwitchStmt(expression: Expr("keyPath")) {
+        for node in SYNTAX_NODES where !node.isBase {
+          for child in node.children where child.requiresLeadingNewline {
+            SwitchCase("case \\\(node.type.syntaxBaseName).\(child.swiftName):") {
+              ReturnStmt("return true")
+            }
+          }
+        }
+        SwitchCase("default:") {
+          ReturnStmt("return putNextTokenOnNewLine")
+        }
+      }
+    }
+
+    FunctionDecl("open func childrenSeparatedByNewline(_ node: Syntax) -> Bool") {
+      SwitchStmt(expression: Expr("node.as(SyntaxEnum.self)")) {
+        for node in SYNTAX_NODES where !node.isBase {
+          if node.elementsSeparatedByNewline {
+            SwitchCase("case .\(node.swiftSyntaxKind):") {
+              ReturnStmt("return true")
+            }
+          }
+        }
+        SwitchCase("default:") {
+          ReturnStmt("return false")
+        }
+      }
+    }
+
+    FunctionDecl("open func requiresLeadingSpace(_ tokenKind: TokenKind) -> Bool") {
+      SwitchStmt(expression: Expr("tokenKind")) {
+        for token in SYNTAX_TOKENS {
+          if token.requiresLeadingSpace {
+            SwitchCase("case .\(token.swiftKind):") {
+              ReturnStmt("return true")
+            }
+          }
+        }
+        SwitchCase("default:") {
+          ReturnStmt("return false")
+        }
+      }
+    }
+
+    FunctionDecl("open func requiresTrailingSpace(_ tokenKind: TokenKind) -> Bool") {
+      SwitchStmt(expression: Expr("tokenKind")) {
+        for token in SYNTAX_TOKENS {
+          if token.requiresTrailingSpace {
+            SwitchCase("case .\(token.swiftKind):") {
+              ReturnStmt("return true")
+            }
+          }
+        }
+        SwitchCase(#"case .contextualKeyword("async"):"#) {
+          ReturnStmt("return true")
+        }
+        SwitchCase("default:") {
+          ReturnStmt("return false")
+        }
       }
     }
 
-    createTokenFormatFunction()
+    FunctionDecl("""
+      private func getKeyPath(_ node: Syntax) -> AnyKeyPath? {
+        guard let parent = node.parent else {
+          return nil
+        }
+        guard case .layout(let childrenKeyPaths) = parent.kind.syntaxNodeType.structure else {
+          return nil
+        }
+        return childrenKeyPaths[node.indexInParent]
+      }
+      """
+    )
   }
 }
 
@@ -55,52 +185,6 @@ private func createChildVisitCall(childType: SyntaxBuildableType, rewrittenExpr:
   }
 }
 
-private func makeLayoutNodeRewriteFunc(node: Node) -> FunctionDecl {
-  let rewriteResultType: String
-  if node.type.baseType?.syntaxKind == "Syntax" && node.type.syntaxKind != "Missing" {
-    rewriteResultType = node.type.syntaxBaseName
-  } else {
-    rewriteResultType = node.type.baseType?.syntaxBaseName ?? node.type.syntaxBaseName
-  }
-  return FunctionDecl("""
-
-    open override func visit(_ node: \(node.type.syntaxBaseName)) -> \(rewriteResultType)
-    """) {
-    for child in node.children {
-      if child.isIndented {
-        SequenceExpr("indentationLevel += 1")
-      }
-      let variableLetVar = child.requiresLeadingNewline ? "var" : "let"
-      VariableDecl("\(variableLetVar) \(child.swiftName) = \(createChildVisitCall(childType: child.type, rewrittenExpr: MemberAccessExpr("node.\(child.swiftName)")))")
-      if child.requiresLeadingNewline {
-        IfStmt(
-          """
-          if \(child.swiftName).leadingTrivia.first?.isNewline != true {
-            \(child.swiftName).leadingTrivia = indentedNewline + \(child.swiftName).leadingTrivia
-          }
-          """
-        )
-      }
-      if child.isIndented {
-        SequenceExpr("indentationLevel -= 1")
-      }
-    }
-    let reconstructed = FunctionCallExpr(callee: node.type.syntaxBaseName) {
-      for child in node.children {
-        TupleExprElement(
-          label: child.isUnexpectedNodes ? nil : child.swiftName,
-          expression: Expr(child.swiftName)
-        )
-      }
-    }
-    if rewriteResultType != node.type.syntaxBaseName {
-      ReturnStmt("return \(rewriteResultType)(\(reconstructed))")
-    } else {
-      ReturnStmt("return \(reconstructed)")
-    }
-  }
-}
-
 private func makeSyntaxCollectionRewriteFunc(node: Node) -> FunctionDecl {
   let rewriteResultType = node.type.syntaxBaseName
   return FunctionDecl("""
@@ -132,70 +216,3 @@ private func makeSyntaxCollectionRewriteFunc(node: Node) -> FunctionDecl {
   }
 }
 
-private func createTokenFormatFunction() -> FunctionDecl {
-  return FunctionDecl("""
-
-    open override func visit(_ node: TokenSyntax) -> TokenSyntax
-    """) {
-    VariableDecl("var leadingTrivia = node.leadingTrivia")
-    VariableDecl("var trailingTrivia = node.trailingTrivia")
-    SwitchStmt(expression: MemberAccessExpr(base: "node", name: "tokenKind")) {
-      for token in SYNTAX_TOKENS where token.name != "ContextualKeyword" {
-        SwitchCase("case .\(token.swiftKind):") {
-          if token.requiresLeadingSpace {
-            IfStmt(
-              """
-              if leadingTrivia.isEmpty && lastRewrittenToken?.trailingTrivia.isEmpty != false {
-                leadingTrivia += .space
-              }
-              """
-            )
-          }
-          if token.requiresTrailingSpace {
-            IfStmt(
-              """
-              if trailingTrivia.isEmpty {
-                trailingTrivia += .space
-              }
-              """
-            )
-          }
-          if !token.requiresLeadingSpace && !token.requiresTrailingSpace {
-            BreakStmt("break")
-          }
-        }
-      }
-      SwitchCase("case .eof:") {
-        BreakStmt("break")
-      }
-      SwitchCase("case .contextualKeyword:") {
-        SwitchStmt(
-          """
-          switch node.text {
-          case "async":
-            if trailingTrivia.isEmpty {
-              trailingTrivia += .space
-            }
-          default:
-            break
-          }
-          """
-        )
-      }
-    }
-    SequenceExpr("leadingTrivia = leadingTrivia.indented(indentation: indentation)")
-    SequenceExpr("trailingTrivia = trailingTrivia.indented(indentation: indentation)")
-    VariableDecl(
-      """
-      let rewritten = TokenSyntax(
-        node.tokenKind,
-        leadingTrivia: leadingTrivia,
-        trailingTrivia: trailingTrivia,
-        presence: node.presence
-      )
-      """
-    )
-    SequenceExpr("lastRewrittenToken = rewritten")
-    ReturnStmt("return rewritten")
-  }
-}