[SwiftSyntax] Ensure SyntaxRewriter respects nil children when rewriting (#15052)

Author: harlanhaskins

test/SwiftSyntax/Inputs/closure.swift

@@ -0,0 +1,3 @@
+// A closure without a signature. The test will ensure it stays the same after
+// applying a rewriting pass.
+let x: () -> Void = {}

test/SwiftSyntax/VisitorTest.swift

@@ -37,4 +37,20 @@ VisitorTests.test("Basic") {
   })
 }
 
+VisitorTests.test("RewritingNodeWithEmptyChild") {
+  class ClosureRewriter: SyntaxRewriter {
+    override func visit(_ node: ClosureExprSyntax) -> ExprSyntax {
+      // Perform a no-op transform that requires rebuilding the node.
+      return node.withSignature(node.signature)
+    }
+  }
+  expectDoesNotThrow({
+    let parsed = try SourceFileSyntax.decodeSourceFileSyntax(try
+      SwiftLang.parse(getInput("closure.swift")))
+    let rewriter = ClosureRewriter()
+    let rewritten = rewriter.visit(parsed)
+    expectEqual(parsed.description, rewritten.description)
+  })
+}
+
 runAllTests()

tools/SwiftSyntax/Syntax.swift

@@ -68,6 +68,12 @@ extension Syntax {
     return SyntaxChildren(node: self)
   }
 
+  /// The number of children, `present` or `missing`, in this node.
+  /// This value can be used safely with `child(at:)`.
+  public var numberOfChildren: Int {
+    return data.childCaches.count
+  }
+
   /// Whether or not this node it marked as `present`.
   public var isPresent: Bool {
     return raw.presence == .present

tools/SwiftSyntax/SyntaxRewriter.swift.gyb

@@ -64,7 +64,17 @@ open class SyntaxRewriter {
   }
 
   func visitChildren(_ node: Syntax) -> Syntax {
-    let newLayout = node.children.map { visit($0).raw }
+    // Visit all children of this node, returning `nil` if child is not
+    // present. This will ensure that there are always the same number
+    // of children after transforming.
+    let newLayout = (0..<node.numberOfChildren).map { (i: Int) -> RawSyntax? in
+      guard let child = node.child(at: i) else { return nil }
+      return visit(child).raw
+    }
+
+    // Sanity check, ensure the new children are the same length.
+    assert(newLayout.count == node.raw.layout.count)
+
     return makeSyntax(node.raw.replacingLayout(newLayout))
   }
 }