Merge pull request #1257 from ahoppen/ahoppen/performance-improvements

Fix performance regressions introduced by my recent changes

Author: ahoppen

CodeGeneration/Sources/generate-swiftsyntax/KeywordFile.swift

@@ -15,6 +15,12 @@ import SwiftSyntaxBuilder
 import SyntaxSupport
 import Utils
 
+let lookupTable = ArrayExprSyntax {
+  for keyword in KEYWORDS {
+    ArrayElementSyntax(expression: ExprSyntax("\(literal: keyword.name)"))
+  }
+}
+
 let keywordFile = SourceFileSyntax {
   ExtensionDeclSyntax(
     """
@@ -31,10 +37,10 @@ let keywordFile = SourceFileSyntax {
 
   EnumDeclSyntax("""
     @frozen  // FIXME: Not actually stable, works around a miscompile
-    public enum Keyword: StaticString
+    public enum Keyword: UInt8, Hashable
     """) {
-    for keyword in KEYWORDS {
-      EnumCaseDeclSyntax("case \(raw: keyword.escapedName)")
+    for (index, keyword) in KEYWORDS.enumerated() {
+      EnumCaseDeclSyntax("case \(raw: keyword.escapedName) = \(literal: index)")
     }
 
     InitializerDeclSyntax("@_spi(RawSyntax) public init?(_ text: SyntaxText)") {
@@ -75,10 +81,20 @@ let keywordFile = SourceFileSyntax {
       }
     }
 
+    VariableDeclSyntax("""
+    /// This is really unfortunate. Really, we should have a `switch` in
+    /// `Keyword.defaultText` to return the keyword's kind but the constant lookup
+    /// table is significantly faster. Ideally, we could also get the compiler to
+    /// constant-evaluate `Keyword.spi.defaultText` to a `SyntaxText` but I don't
+    /// see how that's possible right now.
+    private static let keywordTextLookupTable: [SyntaxText] = \(lookupTable)
+    """)
+
     VariableDeclSyntax(
       """
-      var defaultText: SyntaxText {
-        return SyntaxText(self.rawValue)
+      @_spi(RawSyntax)
+      public var defaultText: SyntaxText {
+        return Keyword.keywordTextLookupTable[Int(self.rawValue)]
       }
       """
     )

CodeGeneration/Sources/generate-swiftsyntax/templates/ideutils/SyntaxClassificationFile.swift

@@ -94,7 +94,7 @@ let syntaxClassificationFile = SourceFileSyntax {
       modifiers: [DeclModifierSyntax(name: .keyword(.internal))],
       name: IdentifierPatternSyntax("classification"),
       type: TypeAnnotationSyntax(type: TypeSyntax("SyntaxClassification"))) {
-        SwitchStmtSyntax(expression: ExprSyntax("self")) {
+        SwitchStmtSyntax(expression: ExprSyntax("self.base")) {
           for token in SYNTAX_TOKENS {
             SwitchCaseSyntax("case .\(raw: token.swiftKind):") {
               if let classification = token.classification {

CodeGeneration/Sources/generate-swiftsyntax/templates/swiftsyntax/TokenKindFile.swift

@@ -203,35 +203,98 @@ let tokenKindFile = SourceFileSyntax {
   }
 
   EnumDeclSyntax("""
-    /// Plain token kind value, without an associated `String` value.
+    // Note: It's important that this enum is marked as having a raw base kind
+    // because it significantly improves performance when comparing two
+    // `RawTokenBaseKind` for equality. With the raw value, it compiles down to
+    // a primitive integer compare, without, it calls into `__derived_enum_equals`.
     @frozen // FIXME: Not actually stable, works around a miscompile
-    public enum RawTokenKind: Equatable, Hashable
+    public enum RawTokenBaseKind: UInt8, Equatable, Hashable
     """) {
     EnumCaseDeclSyntax("case eof")
 
     for token in SYNTAX_TOKENS {
-      if let associatedValueClass = token.associatedValueClass {
-        EnumCaseDeclSyntax("case \(raw: token.swiftKind)(\(raw: associatedValueClass))")
-      } else {
-        EnumCaseDeclSyntax("case \(raw: token.swiftKind)")
+      EnumCaseDeclSyntax("case \(raw: token.swiftKind)")
+    }
+  }
+
+  DeclSyntax("""
+  fileprivate extension Keyword {
+    static var rawValueZero: Keyword {
+      return Keyword(rawValue: 0)!
+    }
+  }
+  """)
+
+  StructDeclSyntax("""
+    /// Similar to `TokenKind` but without a `String` associated value.
+    /// Technically, this should be an enum like
+    /// ```
+    /// enum RawTokenKind {
+    ///   case eof
+    ///   case associatedtypeKeyword
+    ///   // remaining case from `RawTokenBaseKind`...
+    ///   case keyword(Keyword)
+    /// }
+    /// ```
+    ///
+    /// But modelling it this way has significant performance implications since
+    /// comparing two `RawTokenKind` calls into `__derived_enum_equals`. It's more
+    /// effient to model the base kind as an enum with a raw value and store the
+    /// keyword separately.
+    ///
+    /// Whenever `base` is not `keyword`, `keyword` should have a raw value
+    /// of `0`.
+    @frozen // FIXME: Not actually stable, works around a miscompile
+    public struct RawTokenKind: Equatable, Hashable
+    """) {
+    DeclSyntax("public let base: RawTokenBaseKind")
+    DeclSyntax("public let keyword: Keyword")
+
+    DeclSyntax("""
+      public init(base: RawTokenBaseKind, keyword: Keyword) {
+        assert(base == .keyword || keyword.rawValue == 0)
+        self.base = base
+        self.keyword = keyword
       }
+      """
+    )
+
+    DeclSyntax("""
+    public static var eof: RawTokenKind {
+      return RawTokenKind(base: .eof, keyword: .rawValueZero)
+    }
+    """)
+    for token in SYNTAX_TOKENS where token.swiftKind != "keyword" {
+      VariableDeclSyntax(
+        modifiers: [DeclModifierSyntax(leadingTrivia: .newline, name: .keyword(.public)), DeclModifierSyntax(name: .keyword(.static))],
+        name: IdentifierPatternSyntax("\(raw: token.swiftKind)"),
+        type: TypeAnnotationSyntax(type: TypeSyntax("RawTokenKind"))
+      ) {
+        StmtSyntax("return RawTokenKind(base: .\(raw: token.swiftKind), keyword: .rawValueZero)")
+      }
+    }
+
+    DeclSyntax("""
+    public static func keyword(_ keyword: Keyword) -> RawTokenKind {
+      return RawTokenKind(base: .keyword, keyword: keyword)
     }
+    """)
 
     VariableDeclSyntax(
       attributes: [.attribute(AttributeSyntax(attributeName: TypeSyntax("_spi"), leftParen: .leftParenToken(), argument: .token(.identifier("RawSyntax")), rightParen: .rightParenToken()))],
       modifiers: [DeclModifierSyntax(leadingTrivia: .newline, name: .keyword(.public))],
       name: IdentifierPatternSyntax("defaultText"),
       type: TypeAnnotationSyntax(type: OptionalTypeSyntax("SyntaxText?"))
     ) {
-      SwitchStmtSyntax(expression: ExprSyntax("self")) {
+      SwitchStmtSyntax(expression: ExprSyntax("self.base")) {
         SwitchCaseSyntax("case .eof:") {
           ReturnStmtSyntax(#"return """#)
         }
 
         for token in SYNTAX_TOKENS {
-          if token.associatedValueClass != nil {
-            SwitchCaseSyntax("case .\(raw: token.swiftKind)(let assoc):") {
-              ReturnStmtSyntax("return assoc.defaultText")
+          if token.swiftKind == "keyword" {
+            SwitchCaseSyntax("case .\(raw: token.swiftKind):") {
+              ReturnStmtSyntax("return self.keyword.defaultText")
             }
           } else if let text = token.text {
             SwitchCaseSyntax("case .\(raw: token.swiftKind):") {
@@ -251,7 +314,7 @@ let tokenKindFile = SourceFileSyntax {
       name: IdentifierPatternSyntax("nameForDiagnostics"),
       type: TypeAnnotationSyntax(type: TypeSyntax("String"))
     ) {
-      SwitchStmtSyntax(expression: ExprSyntax("self")) {
+      SwitchStmtSyntax(expression: ExprSyntax("self.base")) {
         SwitchCaseSyntax("case .eof:") {
           ReturnStmtSyntax(#"return "end of file""#)
         }
@@ -261,12 +324,12 @@ let tokenKindFile = SourceFileSyntax {
             ReturnStmtSyntax("return #\"\(raw: token.nameForDiagnostics)\"#")
           }
         }
-        SwitchCaseSyntax("case .keyword(let keyword):") {
-          ReturnStmtSyntax("return String(syntaxText: keyword.defaultText)")
+        SwitchCaseSyntax("case .keyword:") {
+          ReturnStmtSyntax("return String(syntaxText: self.keyword.defaultText)")
         }
       }
     }
-    
+
     VariableDeclSyntax(
       leadingTrivia: [
         .docBlockComment("/// Returns `true` if the token is a Swift keyword."),
@@ -284,7 +347,7 @@ let tokenKindFile = SourceFileSyntax {
       name: IdentifierPatternSyntax("isLexerClassifiedKeyword"),
       type: TypeAnnotationSyntax(type: TypeSyntax("Bool"))
     ) {
-      SwitchStmtSyntax(expression: ExprSyntax("self")) {
+      SwitchStmtSyntax(expression: ExprSyntax("self.base")) {
         SwitchCaseSyntax("case .eof:") {
           ReturnStmtSyntax("return false")
         }
@@ -294,8 +357,8 @@ let tokenKindFile = SourceFileSyntax {
             ReturnStmtSyntax("return \(raw: token.isKeyword)")
           }
         }
-        SwitchCaseSyntax("case .keyword(let keyword):") {
-          ReturnStmtSyntax("return keyword.isLexerClassified")
+        SwitchCaseSyntax("case .keyword:") {
+          ReturnStmtSyntax("return self.keyword.isLexerClassified")
         }
       }
     }
@@ -317,7 +380,7 @@ let tokenKindFile = SourceFileSyntax {
       name: IdentifierPatternSyntax("isPunctuation"),
       type: TypeAnnotationSyntax(type: TypeSyntax("Bool"))
     ) {
-      SwitchStmtSyntax(expression: ExprSyntax("self")) {
+      SwitchStmtSyntax(expression: ExprSyntax("self.base")) {
         SwitchCaseSyntax("case .eof:") {
           ReturnStmtSyntax("return false")
         }
@@ -370,16 +433,16 @@ let tokenKindFile = SourceFileSyntax {
       @_spi(RawSyntax)
       public static func fromRaw(kind rawKind: RawTokenKind, text: String) -> TokenKind
       """) {
-      SwitchStmtSyntax(expression: ExprSyntax("rawKind")) {
+      SwitchStmtSyntax(expression: ExprSyntax("rawKind.base")) {
         SwitchCaseSyntax("case .eof:") {
           ReturnStmtSyntax("return .eof")
         }
 
         for token in SYNTAX_TOKENS {
-          if token.associatedValueClass != nil {
-            SwitchCaseSyntax("case .\(raw: token.swiftKind)(let assoc):") {
-              FunctionCallExprSyntax("assert(text.isEmpty || String(syntaxText: assoc.defaultText) == text)")
-              ReturnStmtSyntax("return .\(raw: token.swiftKind)(assoc)")
+          if token.swiftKind == "keyword" {
+            SwitchCaseSyntax("case .\(raw: token.swiftKind):") {
+              FunctionCallExprSyntax("assert(text.isEmpty || String(syntaxText: rawKind.keyword.defaultText) == text)")
+              ReturnStmtSyntax("return .keyword(rawKind.keyword)")
             }
           } else if token.text != nil {
             SwitchCaseSyntax("case .\(raw: token.swiftKind):") {
@@ -407,9 +470,9 @@ let tokenKindFile = SourceFileSyntax {
         }
 
         for token in SYNTAX_TOKENS {
-          if token.associatedValueClass != nil {
-            SwitchCaseSyntax("case .\(raw: token.swiftKind)(let assoc):") {
-              ReturnStmtSyntax("return (.\(raw: token.swiftKind)(assoc), nil)")
+          if token.swiftKind == "keyword" {
+            SwitchCaseSyntax("case .\(raw: token.swiftKind)(let keyword):") {
+              ReturnStmtSyntax("return (.\(raw: token.swiftKind)(keyword), nil)")
             }
           } else if token.text != nil {
             SwitchCaseSyntax("case .\(raw: token.swiftKind):") {

Sources/IDEUtils/generated/SyntaxClassification.swift

@@ -137,7 +137,7 @@ extension SyntaxClassification {
 
 extension RawTokenKind {
   internal var classification: SyntaxClassification {
-    switch self {
+    switch self.base {
     case .wildcard: 
       return .none
     case .leftParen: 

Sources/SwiftParser/Attributes.swift

@@ -170,7 +170,7 @@ extension Parser {
     case .required:
       shouldParseArgument = true
     case .customAttribute:
-      shouldParseArgument = self.lookahead().isCustomAttributeArgument() && self.at(.leftParen, allowTokenAtStartOfLine: false)
+      shouldParseArgument = self.withLookahead { $0.isCustomAttributeArgument() } && self.at(.leftParen, allowTokenAtStartOfLine: false)
     case .optional:
       shouldParseArgument = self.at(.leftParen)
     }
@@ -1113,7 +1113,7 @@ extension Parser {
 // MARK: Lookahead
 
 extension Parser.Lookahead {
-  func isCustomAttributeArgument() -> Bool {
+  mutating func isCustomAttributeArgument() -> Bool {
     var lookahead = self.lookahead()
     lookahead.skipSingle()
 
@@ -1146,7 +1146,7 @@ extension Parser.Lookahead {
       return false
     }
 
-    if self.at(.leftParen, allowTokenAtStartOfLine: false) && self.lookahead().isCustomAttributeArgument() {
+    if self.at(.leftParen, allowTokenAtStartOfLine: false) && self.withLookahead({ $0.isCustomAttributeArgument() }) {
       self.skipSingle()
     }
 

Sources/SwiftParser/Declarations.swift

@@ -29,7 +29,7 @@ extension DeclarationModifier {
 }
 
 extension TokenConsumer {
-  func atStartOfDeclaration(
+  mutating func atStartOfDeclaration(
     isAtTopLevel: Bool = false,
     allowInitDecl: Bool = true,
     allowRecovery: Bool = false
@@ -1242,7 +1242,7 @@ extension Parser {
 
 extension Parser {
   /// Are we at a regular expression literal that could act as an operator?
-  private func atRegexLiteralThatCouldBeAnOperator() -> Bool {
+  private mutating func atRegexLiteralThatCouldBeAnOperator() -> Bool {
     guard self.at(.regexLiteral) else {
       return false
     }
@@ -2258,7 +2258,7 @@ extension Parser {
 
     // Parse the optional generic argument list.
     let generics: RawGenericArgumentClauseSyntax?
-    if self.lookahead().canParseAsGenericArgumentList() {
+    if self.withLookahead({ $0.canParseAsGenericArgumentList() }) {
       generics = self.parseGenericArguments()
     } else {
       generics = nil
@@ -2282,7 +2282,7 @@ extension Parser {
     let trailingClosure: RawClosureExprSyntax?
     let additionalTrailingClosures: RawMultipleTrailingClosureElementListSyntax?
     if self.at(.leftBrace),
-      self.lookahead().isValidTrailingClosure(.trailingClosure)
+      self.withLookahead({ $0.isValidTrailingClosure(.trailingClosure) })
     {
       (trailingClosure, additionalTrailingClosures) =
         self.parseTrailingClosures(.trailingClosure)