[SwiftSyntax] Fix SyntaxChildren iteration (swiftlang#12319)

* [SwiftSyntax] Fix SyntaxChildren iteration

SyntaxChildren.Iterator stops the first time it sees a nil child;
since it traverses all indexes without considering presence, this
causes it to terminate iteration early if there is a missing child
in the list.

* [SwiftSyntax] Add tests for SyntaxChildren

Author: allevato

test/SwiftSyntax/SyntaxChildren.swift

@@ -0,0 +1,60 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+// REQUIRES: OS=macosx
+// REQUIRES: sdk_overlay
+
+import Foundation
+import StdlibUnittest
+import SwiftSyntax
+
+/// Verifies that there is a next item returned by the iterator and that it
+/// satisfies the given predicate.
+func expectNext<Iterator: IteratorProtocol>(
+  _ iterator: inout Iterator,
+  satisfies predicate: (Iterator.Element) throws -> Bool
+) rethrows {
+  let next = iterator.next()
+  expectNotNil(next)
+  expectTrue(try predicate(next!))
+}
+
+/// Verifies that the iterator is exhausted.
+func expectNextIsNil<Iterator: IteratorProtocol>(_ iterator: inout Iterator) {
+  expectNil(iterator.next())
+}
+
+var SyntaxChildrenAPI = TestSuite("SyntaxChildrenAPI")
+
+SyntaxChildrenAPI.test("IterateWithAllPresent") {
+  let returnStmt = SyntaxFactory.makeReturnStmt(
+    returnKeyword: SyntaxFactory.makeReturnKeyword(),
+    expression: SyntaxFactory.makeBlankExpr(),
+    semicolon: SyntaxFactory.makeSemicolonToken())
+
+  var iterator = returnStmt.children.makeIterator()
+  expectNext(&iterator) { ($0 as? TokenSyntax)?.tokenKind == .returnKeyword }
+  expectNext(&iterator) { $0 is ExprSyntax }
+  expectNext(&iterator) { ($0 as? TokenSyntax)?.tokenKind == .semicolon }
+  expectNextIsNil(&iterator)
+}
+
+SyntaxChildrenAPI.test("IterateWithSomeMissing") {
+  let returnStmt = SyntaxFactory.makeReturnStmt(
+    returnKeyword: SyntaxFactory.makeReturnKeyword(),
+    expression: nil,
+    semicolon: SyntaxFactory.makeSemicolonToken())
+
+  var iterator = returnStmt.children.makeIterator()
+  expectNext(&iterator) { ($0 as? TokenSyntax)?.tokenKind == .returnKeyword }
+  expectNext(&iterator) { ($0 as? TokenSyntax)?.tokenKind == .semicolon }
+  expectNextIsNil(&iterator)
+}
+
+SyntaxChildrenAPI.test("IterateWithAllMissing") {
+  let returnStmt = SyntaxFactory.makeBlankReturnStmt()
+
+  var iterator = returnStmt.children.makeIterator()
+  expectNextIsNil(&iterator)
+}
+
+runAllTests()

tools/SwiftSyntax/Syntax.swift

@@ -16,6 +16,10 @@ import Foundation
 /// Each node has accessors for its known children, and allows efficient
 /// iteration over the children through its `children` property.
 public class Syntax: CustomStringConvertible {
+  /// The type of sequence containing the indices of present children.
+  internal typealias PresentChildIndicesSequence =
+    LazyFilterSequence<CountableRange<Int>>
+
   /// The root of the tree this node is currently in.
   internal let _root: SyntaxData
 
@@ -102,6 +106,14 @@ public class Syntax: CustomStringConvertible {
     return Syntax.make(root: _root,  data: _root)
   }
 
+  /// The sequence of indices that correspond to child nodes that are not
+  /// missing.
+  ///
+  /// This property is an implementation detail of `SyntaxChildren`.
+  internal var presentChildIndices: PresentChildIndicesSequence {
+    return raw.layout.indices.lazy.filter { self.raw.layout[$0].isPresent }
+  }
+
   /// Gets the child at the provided index in this node's children.
   /// - Parameter index: The index of the child node you're looking for.
   /// - Returns: A Syntax node for the provided child, or `nil` if there

tools/SwiftSyntax/SyntaxChildren.swift

@@ -14,17 +14,23 @@ import Foundation
 
 public struct SyntaxChildren: Sequence {
   public struct Iterator: IteratorProtocol {
-    var index: Int = 0
     let node: Syntax
+    var indexIterator: Syntax.PresentChildIndicesSequence.Iterator
+
+    init(node: Syntax) {
+      self.indexIterator = node.presentChildIndices.makeIterator()
+      self.node = node
+    }
 
     public mutating func next() -> Syntax? {
-      defer { index += 1 }
+      guard let index = indexIterator.next() else { return nil }
       return node.child(at: index)
     }
   }
+
   let node: Syntax
 
   public func makeIterator() -> SyntaxChildren.Iterator {
-    return Iterator(index: 0, node: node)
+    return Iterator(node: node)
   }
 }