Simplify fast path lookup

Author: onevcat

Sources/Chroma/LanguageDefinition+Helpers.swift

@@ -5,11 +5,11 @@ public extension LanguageDefinition {
         switch kind {
         case .keyword:
             var updated = ensureFastPath()
-            updated.keywords.formUnion(words)
+            updated.appendWords(words, kind: .keyword)
             fastPath = updated
         case .type:
             var updated = ensureFastPath()
-            updated.types.formUnion(words)
+            updated.appendWords(words, kind: .type)
             fastPath = updated
         default:
             break

Sources/Chroma/LanguageDefinition.swift

@@ -18,15 +18,40 @@ public struct LanguageDefinition {
 }
 
 public struct LanguageFastPath: Equatable {
-    public var keywords: Set<String>
-    public var types: Set<String>
+    private var lookup: [String: TokenKind]
 
     public init(keywords: [String] = [], types: [String] = []) {
-        self.keywords = Set(keywords)
-        self.types = Set(types)
+        var lookup: [String: TokenKind] = [:]
+        lookup.reserveCapacity(keywords.count + types.count)
+        for word in types {
+            lookup[word] = .type
+        }
+        for word in keywords {
+            lookup[word] = .keyword
+        }
+        self.lookup = lookup
     }
 
     public var isEmpty: Bool {
-        keywords.isEmpty && types.isEmpty
+        lookup.isEmpty
+    }
+
+    func kind(for word: String) -> TokenKind? {
+        lookup[word]
+    }
+
+    mutating func appendWords(_ words: [String], kind: TokenKind) {
+        switch kind {
+        case .keyword:
+            for word in words {
+                lookup[word] = .keyword
+            }
+        case .type:
+            for word in words where lookup[word] == nil {
+                lookup[word] = .type
+            }
+        default:
+            break
+        }
     }
 }

Sources/Chroma/Tokenizer.swift

@@ -211,13 +211,8 @@ final class RegexTokenizer {
             }
             let wordRange = NSRange(location: location, length: index - location)
             let word = ns.substring(with: wordRange)
-            if fastPath.keywords.contains(word) {
-                return (.keyword, wordRange)
-            }
-            if fastPath.types.contains(word) {
-                return (.type, wordRange)
-            }
-            return nil
+            guard let kind = fastPath.kind(for: word) else { return nil }
+            return (kind, wordRange)
         }
 
         func appendPlainWithFastPath(_ range: NSRange, fastPath: LanguageFastPath) {
@@ -236,10 +231,8 @@ final class RegexTokenizer {
                     }
                     let wordRange = NSRange(location: start, length: index - start)
                     let word = ns.substring(with: wordRange)
-                    if fastPath.keywords.contains(word) {
-                        appendToken(Token(kind: .keyword, range: wordRange))
-                    } else if fastPath.types.contains(word) {
-                        appendToken(Token(kind: .type, range: wordRange))
+                    if let kind = fastPath.kind(for: word) {
+                        appendToken(Token(kind: kind, range: wordRange))
                     } else {
                         appendPlainASCII(wordRange)
                     }