Fix a bug where allTokens would not include the last token if it is missing

Author: ahoppen

Sources/SwiftSyntax/Syntax.swift

@@ -660,21 +660,22 @@ public protocol SyntaxChildChoices: SyntaxProtocol {}
 public struct TokenSequence: Sequence {
   public struct Iterator: IteratorProtocol {
     var nextToken: TokenSyntax?
-    let endPosition: AbsolutePosition
+    /// The last token to iterate (inclusive).
+    let endToken: TokenSyntax?
     let viewMode: SyntaxTreeViewMode
 
-    init(_ token: TokenSyntax?, endPosition: AbsolutePosition, viewMode: SyntaxTreeViewMode) {
+    init(_ token: TokenSyntax?, endToken: TokenSyntax?, viewMode: SyntaxTreeViewMode) {
       self.nextToken = token
-      self.endPosition = endPosition
+      self.endToken = endToken
       self.viewMode = viewMode
     }
 
     public mutating func next() -> TokenSyntax? {
       guard let token = self.nextToken else { return nil }
-      self.nextToken = token.nextToken(viewMode: viewMode)
-      // Make sure we stop once we reach the end of the containing node.
-      if let nextTok = self.nextToken, nextTok.position >= self.endPosition {
+      if nextToken == endToken {
         self.nextToken = nil
+      } else {
+        self.nextToken = token.nextToken(viewMode: viewMode)
       }
       return token
     }
@@ -689,7 +690,7 @@ public struct TokenSequence: Sequence {
   }
 
   public func makeIterator() -> Iterator {
-    return Iterator(node.firstToken(viewMode: viewMode), endPosition: node.endPosition, viewMode: viewMode)
+    return Iterator(node.firstToken(viewMode: viewMode), endToken: node.lastToken(viewMode: viewMode), viewMode: viewMode)
   }
 
   public func reversed() -> ReversedTokenSequence {

Tests/SwiftIDEUtilsTest/ClassificationTests.swift

@@ -117,8 +117,8 @@ public class ClassificationTests: XCTestCase {
     let tree = Parser.parse(source: source)
     do {
       let tokens = Array(tree.tokens(viewMode: .sourceAccurate))
-      XCTAssertEqual(tokens.count, 4)
-      guard tokens.count == 4 else {
+      XCTAssertEqual(tokens.count, 5)
+      guard tokens.count == 5 else {
         return
       }
       let classif = tokens.map { $0.tokenClassification }
@@ -130,6 +130,8 @@ public class ClassificationTests: XCTestCase {
       XCTAssertEqual(classif[2].range, ByteSourceRange(offset: 5, length: 1))
       XCTAssertEqual(classif[3].kind, .typeIdentifier)
       XCTAssertEqual(classif[3].range, ByteSourceRange(offset: 7, length: 3))
+      XCTAssertEqual(classif[4].kind, .none)
+      XCTAssertEqual(classif[4].range, ByteSourceRange(offset: 10, length: 0))
     }
     do {
       let tok = tree.lastToken(viewMode: .sourceAccurate)!.previousToken(viewMode: .sourceAccurate)!
@@ -145,8 +147,8 @@ public class ClassificationTests: XCTestCase {
       let tree = Parser.parse(source: source)
 
       let tokens = Array(tree.tokens(viewMode: .sourceAccurate))
-      XCTAssertEqual(tokens.count, 10)
-      guard tokens.count == 10 else {
+      XCTAssertEqual(tokens.count, 11)
+      guard tokens.count == 11 else {
         return
       }
       let classif = tokens.map { $0.tokenClassification }
@@ -170,15 +172,17 @@ public class ClassificationTests: XCTestCase {
       XCTAssertEqual(classif[8].range, ByteSourceRange(offset: 19, length: 1))
       XCTAssertEqual(classif[9].kind, .integerLiteral)
       XCTAssertEqual(classif[9].range, ByteSourceRange(offset: 21, length: 1))
+      XCTAssertEqual(classif[10].kind, .none)
+      XCTAssertEqual(classif[10].range, ByteSourceRange(offset: 22, length: 0))
     }
 
     do {
       let source = "infix operator *--*"
       let tree = Parser.parse(source: source)
 
       let tokens = Array(tree.tokens(viewMode: .sourceAccurate))
-      XCTAssertEqual(tokens.count, 3)
-      guard tokens.count == 3 else {
+      XCTAssertEqual(tokens.count, 4)
+      guard tokens.count == 4 else {
         return
       }
       let classif = tokens.map { $0.tokenClassification }
@@ -188,6 +192,8 @@ public class ClassificationTests: XCTestCase {
       XCTAssertEqual(classif[1].range, ByteSourceRange(offset: 6, length: 8))
       XCTAssertEqual(classif[2].kind, .operatorIdentifier)
       XCTAssertEqual(classif[2].range, ByteSourceRange(offset: 15, length: 4))
+      XCTAssertEqual(classif[3].kind, .none)
+      XCTAssertEqual(classif[3].range, ByteSourceRange(offset: 19, length: 0))
     }
   }
 }

Tests/SwiftSyntaxTest/SyntaxChildrenTests.swift

@@ -71,4 +71,14 @@ public class SyntaxChildrenTests: XCTestCase {
     try XCTAssertNext(&fixedUpIt) { $0.is(TokenSyntax.self) }
     try XCTAssertNext(&fixedUpIt) { $0.is(MissingExprSyntax.self) }
   }
+
+  public func testTokenSequencesWithMissingChild() {
+    let codeBlock = CodeBlockSyntax(
+      leftBrace: .leftBraceToken(presence: .missing),
+      statements: [],
+      rightBrace: .rightBraceToken(presence: .missing)
+    )
+
+    XCTAssertEqual(Array(codeBlock.tokens(viewMode: .all)).count, 2)
+  }
 }