[classification] Add API TokenSyntax.tokenClassification to get the classification for a specific token

Author: akyrtzi

Sources/SwiftSyntax/SyntaxClassification.swift.gyb

@@ -55,7 +55,7 @@ extension SyntaxClassification {
   /// - Returns: A pair of classification and whether it is "forced", or nil if
   ///   no classification is attached.
   internal static func classify(
-    parentKind: SyntaxKind, indexInParent: UInt32, childKind: SyntaxKind
+    parentKind: SyntaxKind, indexInParent: Int, childKind: SyntaxKind
   ) -> (SyntaxClassification, Bool)? {
     // Separate checks for token nodes (most common checks) versus checks for layout nodes.
     if childKind == .token {

Sources/SwiftSyntax/SyntaxClassifier.swift

@@ -10,6 +10,30 @@
 //
 //===----------------------------------------------------------------------===//
 
+extension TokenSyntax {
+  /// The `SyntaxClassifiedRange` for the token text, excluding trivia.
+  public var tokenClassification: SyntaxClassifiedRange {
+    var contextualClassification: (SyntaxClassification, Bool)? = nil
+    var curData = self.data
+    repeat {
+      guard let parent = curData.parent else { break }
+      contextualClassification = SyntaxClassification.classify(parentKind: parent.raw.kind,
+        indexInParent: curData.indexInParent, childKind: raw.kind)
+      curData = parent.data
+    } while contextualClassification == nil
+
+    let relativeOffset = raw.tokenLeadingTriviaLength.utf8Length
+    let absoluteOffset = position.utf8Offset + relativeOffset
+    let classifiedRange = raw.withUnsafeTokenText(
+      relativeOffset: relativeOffset
+    ) { (tokText: UnsafeTokenText?) -> SyntaxClassifiedRange in
+      return tokText!.classify(offset: absoluteOffset,
+        contextualClassification: contextualClassification)
+    }
+    return classifiedRange
+  }
+}
+
 extension UnsafeTriviaPiece {
   func classify(offset: Int) -> SyntaxClassifiedRange {
     let range = ByteSourceRange(offset: offset, length: length)
@@ -72,7 +96,7 @@ fileprivate struct AbsoluteNode {
     self.position = position
     // Check if this node has a classification otherwise inherit it from the parent.
     self.classification = SyntaxClassification.classify(parentKind: parent.raw.kind,
-      indexInParent: position.indexInParent, childKind: raw.kind) ?? parent.classification
+      indexInParent: Int(position.indexInParent), childKind: raw.kind) ?? parent.classification
   }
 
   var offset: Int { return Int(position.offset) }

Tests/SwiftSyntaxTest/ClassificationTests.swift

@@ -4,6 +4,7 @@ import SwiftSyntax
 public class ClassificationTests: XCTestCase {
   public static let allTests = [
     ("testClassification", testClassification),
+    ("testTokenClassification", testTokenClassification),
   ]
 
   public func testClassification() {
@@ -92,4 +93,25 @@ public class ClassificationTests: XCTestCase {
       XCTAssertEqual(classif.range, ByteSourceRange(offset: 11, length: 1))
     }
   }
+
+  public func testTokenClassification() {
+    let source = "let x: Int"
+    let tree = try! SyntaxParser.parse(source: source)
+    do {
+      let tokens = Array(tree.tokens)
+      XCTAssertEqual(tokens.count, 4)
+      guard tokens.count == 4 else {
+        return
+      }
+      let classif = tokens.map { $0.tokenClassification }
+      XCTAssertEqual(classif[0].kind, .keyword)
+      XCTAssertEqual(classif[0].range, ByteSourceRange(offset: 0, length: 3))
+      XCTAssertEqual(classif[1].kind, .none)
+      XCTAssertEqual(classif[1].range, ByteSourceRange(offset: 4, length: 1))
+      XCTAssertEqual(classif[2].kind, .none)
+      XCTAssertEqual(classif[2].range, ByteSourceRange(offset: 5, length: 1))
+      XCTAssertEqual(classif[3].kind, .typeIdentifier)
+      XCTAssertEqual(classif[3].range, ByteSourceRange(offset: 7, length: 3))
+    }
+  }
 }