[ASTGen] TokenSyntax.keyworkKind to directly get Keyword kind

This should be faster than e.g. `token.tokenKind == .keyword(.true)`.

Author: rintaro

lib/ASTGen/Sources/ASTGen/ASTGen.swift

@@ -137,11 +137,10 @@ extension ASTGenVisitor {
   /// escaped identifier, backticks are stripped.
   @inline(__always)
   func generateIdentifier(_ token: TokenSyntax) -> BridgedIdentifier {
-    var text = token.rawText
-    // FIXME: Maybe `TokenSyntax.tokenView.rawKind == .wildcard`, or expose it as `.rawTokenKind`.
-    if text == "_" {
+    if token.rawTokenKind == .wildcard {
       return nil
     }
+    var text = token.rawText
     if text.count > 2 && text.hasPrefix("`") && text.hasSuffix("`") {
       text = .init(rebasing: text.dropFirst().dropLast())
     }
@@ -382,6 +381,24 @@ extension Optional where Wrapped: LazyCollectionProtocol {
   }
 }
 
+extension TokenSyntax {
+  /// Get `Keyword` kind if the token is a keyword.
+  var keywordKind: Keyword? {
+    // Performance note:
+    // This is faster than `token.tokenKind == .keyword(.true)` because
+    // `TokenKind.tokenKind` may instantiate `Swift.String`.
+    // That being said, `SwiftSyntax.Keyword` is a non-SPI public type, so it
+    // cannot be `@frozen`. Also `Keyword(_:SyntaxText)` itself is heavier than
+    // simple `token.rawText == "true"`.
+    // We should ensure `token.keywordKind == .true` is optimized out to
+    // a simple `cmp` instruction.
+    guard rawTokenKind == .keyword else {
+      return nil
+    }
+    return Keyword(self.rawText)
+  }
+}
+
 /// Generate AST nodes for all top-level entities in the given source file.
 @_cdecl("swift_ASTGen_buildTopLevelASTNodes")
 public func buildTopLevelASTNodes(

lib/ASTGen/Sources/ASTGen/Decls.swift

@@ -13,7 +13,7 @@
 import ASTBridging
 import BasicBridging
 import SwiftDiagnostics
-@_spi(ExperimentalLanguageFeatures) import SwiftSyntax
+@_spi(ExperimentalLanguageFeatures) @_spi(RawSyntax) import SwiftSyntax
 
 // MARK: - TypeDecl
 
@@ -293,7 +293,7 @@ extension ASTGenVisitor {
     let initializer = generate(initializerClause: node.bindings.first!.initializer!)
 
     let isStatic = false  // TODO: compute this
-    let isLet = node.bindingSpecifier.tokenKind == .keyword(.let)
+    let isLet = node.bindingSpecifier.keywordKind == .let
 
     return .createParsed(
       self.ctx,
@@ -347,7 +347,7 @@ extension ASTGenVisitor {
       declContext: self.declContext,
       initKeywordLoc: self.generateSourceLoc(node.initKeyword),
       failabilityMarkLoc: self.generateSourceLoc(node.optionalMark),
-      isIUO: node.optionalMark?.tokenKind == .exclamationMark,
+      isIUO: node.optionalMark?.rawTokenKind == .exclamationMark,
       genericParamList: self.generate(genericParameterClause: node.genericParameterClause),
       parameterList: self.generate(functionParameterClause: node.signature.parameterClause),
       asyncSpecifierLoc: self.generateSourceLoc(node.signature.effectSpecifiers?.asyncSpecifier),
@@ -385,11 +385,11 @@ extension ASTGenVisitor {
 // MARK: - OperatorDecl
 
 extension BridgedOperatorFixity {
-  fileprivate init?(from tokenKind: TokenKind) {
-    switch tokenKind {
-    case .keyword(.infix): self = .infix
-    case .keyword(.prefix): self = .prefix
-    case .keyword(.postfix): self = .postfix
+  fileprivate init?(from keyword: Keyword?) {
+    switch keyword {
+    case .infix: self = .infix
+    case .prefix: self = .prefix
+    case .postfix: self = .postfix
     default: return nil
     }
   }
@@ -402,7 +402,7 @@ extension ASTGenVisitor {
         self.generateIdentifierAndSourceLoc(node.operatorPrecedenceAndTypes?.precedenceGroup)
 
     let fixity: BridgedOperatorFixity
-    if let value = BridgedOperatorFixity(from: node.fixitySpecifier.tokenKind) {
+    if let value = BridgedOperatorFixity(from: node.fixitySpecifier.keywordKind) {
       fixity = value
     } else {
       fixity = .infix
@@ -428,11 +428,11 @@ extension ASTGenVisitor {
 // MARK: - PrecedenceGroupDecl
 
 extension BridgedAssociativity {
-  fileprivate init?(from tokenKind: TokenKind) {
-    switch tokenKind {
-    case .keyword(.none): self = .none
-    case .keyword(.left): self = .left
-    case .keyword(.right): self = .right
+  fileprivate init?(from keyword: Keyword?) {
+    switch keyword {
+    case .none?: self = .none
+    case .left?: self = .left
+    case .right?: self = .right
     default: return nil
     }
   }
@@ -459,14 +459,14 @@ extension ASTGenVisitor {
       switch element {
       case .precedenceGroupRelation(let relation):
         let keyword = relation.higherThanOrLowerThanLabel
-        switch keyword.tokenKind {
-        case .keyword(.higherThan):
+        switch keyword.keywordKind {
+        case .higherThan:
           if let current = body.higherThanRelation {
             diagnoseDuplicateSyntax(relation, original: current)
           } else {
             body.higherThanRelation = relation
           }
-        case .keyword(.lowerThan):
+        case .lowerThan:
           if let current = body.lowerThanRelation {
             diagnoseDuplicateSyntax(relation, original: current)
           } else {
@@ -492,7 +492,7 @@ extension ASTGenVisitor {
 
     let associativityValue: BridgedAssociativity
     if let token = body.associativity?.value {
-      if let value = BridgedAssociativity(from: token.tokenKind) {
+      if let value = BridgedAssociativity(from: token.keywordKind) {
         associativityValue = value
       } else {
         self.diagnose(Diagnostic(node: token, message: UnexpectedTokenKindError(token: token)))
@@ -504,7 +504,7 @@ extension ASTGenVisitor {
 
     let assignmentValue: Bool
     if let token = body.assignment?.value {
-      if token.tokenKind == .keyword(.true) {
+      if token.keywordKind == .true {
         assignmentValue = true
       } else {
         self.diagnose(Diagnostic(node: token, message: UnexpectedTokenKindError(token: token)))
@@ -538,15 +538,15 @@ extension ASTGenVisitor {
 // MARK: - ImportDecl
 
 extension BridgedImportKind {
-  fileprivate init?(from tokenKind: TokenKind) {
-    switch tokenKind {
-    case .keyword(.typealias): self = .type
-    case .keyword(.struct): self = .struct
-    case .keyword(.class): self = .class
-    case .keyword(.enum): self = .enum
-    case .keyword(.protocol): self = .protocol
-    case .keyword(.var), .keyword(.let): self = .var
-    case .keyword(.func): self = .func
+  fileprivate init?(from keyword: Keyword?) {
+    switch keyword {
+    case .typealias: self = .type
+    case .struct: self = .struct
+    case .class: self = .class
+    case .enum: self = .enum
+    case .protocol: self = .protocol
+    case .var, .let: self = .var
+    case .func: self = .func
     default: return nil
     }
   }
@@ -556,7 +556,7 @@ extension ASTGenVisitor {
   func generate(importDecl node: ImportDeclSyntax) -> BridgedImportDecl {
     let importKind: BridgedImportKind
     if let specifier = node.importKindSpecifier {
-      if let value = BridgedImportKind(from: specifier.tokenKind) {
+      if let value = BridgedImportKind(from: specifier.keywordKind) {
         importKind = value
       } else {
         self.diagnose(Diagnostic(node: specifier, message: UnexpectedTokenKindError(token: specifier)))

lib/ASTGen/Sources/ASTGen/Exprs.swift

@@ -12,7 +12,7 @@
 
 import ASTBridging
 import SwiftDiagnostics
-@_spi(ExperimentalLanguageFeatures) import SwiftSyntax
+@_spi(ExperimentalLanguageFeatures) @_spi(RawSyntax) import SwiftSyntax
 
 /// Check if an `ExprSyntax` can be generated using ASTGen.
 ///
@@ -292,12 +292,12 @@ extension ASTGenVisitor {
 
   private func createDeclNameRef(declReferenceExpr node: DeclReferenceExprSyntax) -> (name: BridgedDeclNameRef, loc: BridgedDeclNameLoc) {
     let baseName: BridgedDeclBaseName
-    switch node.baseName.tokenKind {
-    case .keyword(.`init`):
+    switch node.baseName.keywordKind {
+    case .`init`:
       baseName = .createConstructor()
-    case .keyword(.deinit):
+    case .deinit:
       baseName = .createDestructor()
-    case .keyword(.subscript):
+    case .subscript:
       baseName = .createSubscript()
     default:
       baseName = .createIdentifier(self.generateIdentifier(node.baseName))
@@ -480,14 +480,14 @@ extension ASTGenVisitor {
         asLoc: asLoc,
         type: type
       ).asExpr
-    case let question? where question.text == "?":
+    case let question? where question.rawTokenKind == .postfixQuestionMark:
       return BridgedConditionalCheckedCastExpr.createParsed(
         self.ctx,
         asLoc: asLoc,
         questionLoc: self.generateSourceLoc(question),
         type: type
       ).asExpr
-    case let exclaim? where exclaim.text == "!":
+    case let exclaim? where exclaim.rawTokenKind == .exclamationMark:
       return BridgedForcedCheckedCastExpr.createParsed(
         self.ctx,
         asLoc: asLoc,

lib/ASTGen/Sources/ASTGen/Literals.swift

@@ -39,7 +39,7 @@ extension ASTGenVisitor {
   }
 
   func generate(booleanLiteralExpr node: BooleanLiteralExprSyntax) -> BridgedBooleanLiteralExpr {
-    let value = node.literal.tokenKind == .keyword(.true)
+    let value = node.literal.keywordKind == .true
     return .createParsed(
       ctx,
       value: value,

lib/ASTGen/Sources/ASTGen/Types.swift

@@ -106,7 +106,7 @@ extension ASTGenVisitor {
     let loc = self.generateSourceLoc(node.name)
 
     // If this is the bare 'Any' keyword, produce an empty composition type.
-    if node.name.rawText == "Any" && node.genericArgumentClause == nil {
+    if node.name.keywordKind == .Any && node.genericArgumentClause == nil {
       return BridgedCompositionTypeRepr.createEmpty(self.ctx, anyKeywordLoc: loc)
     }
 
@@ -203,14 +203,14 @@ extension ASTGenVisitor {
   func generate(metatypeType node: MetatypeTypeSyntax) -> BridgedTypeRepr {
     let baseType = generate(type: node.baseType)
     let tyLoc = self.generateSourceLoc(node.metatypeSpecifier)
-    if node.metatypeSpecifier.rawText == "Type" {
+    if node.metatypeSpecifier.keywordKind == .Type {
       return BridgedMetatypeTypeRepr.createParsed(
         self.ctx,
         base: baseType,
         typeKeywordLoc: tyLoc
       )
     } else {
-      assert(node.metatypeSpecifier.rawText == "Protocol")
+      assert(node.metatypeSpecifier.keywordKind == .Protocol)
       return BridgedProtocolTypeRepr.createParsed(
         self.ctx,
         base: baseType,
@@ -298,14 +298,14 @@ extension ASTGenVisitor {
   func generate(someOrAnyType node: SomeOrAnyTypeSyntax) -> BridgedTypeRepr {
     let someOrAnyLoc = self.generateSourceLoc(node.someOrAnySpecifier)
     let baseTy = generate(type: node.constraint)
-    if node.someOrAnySpecifier.rawText == "some" {
+    if node.someOrAnySpecifier.keywordKind == .some {
       return BridgedOpaqueReturnTypeRepr.createParsed(
         self.ctx,
         someKeywordLoc: someOrAnyLoc,
         base: baseTy
       )
     } else {
-      assert(node.someOrAnySpecifier.text == "any")
+      assert(node.someOrAnySpecifier.keywordKind == .any)
       return BridgedExistentialTypeRepr.createParsed(
         self.ctx,
         anyKeywordLoc: someOrAnyLoc,
@@ -320,16 +320,16 @@ extension ASTGenVisitor {
 // MARK: - SpecifierTypeRepr/AttributedTypeRepr
 
 extension BridgedAttributedTypeSpecifier {
-  fileprivate init?(from tokenKind: TokenKind) {
-    switch tokenKind {
-    case .keyword(.inout): self = .inOut
-    case .keyword(.borrowing): self = .borrowing
-    case .keyword(.consuming): self = .consuming
-    case .keyword(.__shared): self = .legacyShared
-    case .keyword(.__owned): self = .legacyOwned
-    case .keyword(._const): self = .const
-    case .keyword(.isolated): self = .isolated
-    case .keyword(._resultDependsOn): self = .resultDependsOn
+  fileprivate init?(from keyword: Keyword?) {
+    switch keyword {
+    case .inout: self = .inOut
+    case .borrowing: self = .borrowing
+    case .consuming: self = .consuming
+    case .__shared: self = .legacyShared
+    case .__owned: self = .legacyOwned
+    case ._const: self = .const
+    case .isolated: self = .isolated
+    case ._resultDependsOn: self = .resultDependsOn
     default: return nil
     }
   }
@@ -341,7 +341,7 @@ extension ASTGenVisitor {
 
     // Handle specifiers.
     if let specifier = node.specifier {
-      if let kind = BridgedAttributedTypeSpecifier(from: specifier.tokenKind) {
+      if let kind = BridgedAttributedTypeSpecifier(from: specifier.keywordKind) {
         type = BridgedSpecifierTypeRepr.createParsed(
           self.ctx,
           base: type,