From 69fe457cb540c1880e860e03a287f1313942f2cd Mon Sep 17 00:00:00 2001
From: davidsemakula <hello@davidsemakula.com>
Date: Mon, 15 Apr 2024 16:46:55 +0300
Subject: [PATCH] internal: simplify `TokenSet` implementation

---
 crates/parser/src/token_set.rs | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/crates/parser/src/token_set.rs b/crates/parser/src/token_set.rs
index 2d615b7a5d1..88a89a53a76 100644
--- a/crates/parser/src/token_set.rs
+++ b/crates/parser/src/token_set.rs
@@ -6,6 +6,9 @@
 #[derive(Clone, Copy)]
 pub(crate) struct TokenSet([u64; 3]);
 
+/// `TokenSet`s should only include token `SyntaxKind`s, so the discriminant of any passed/included
+/// `SyntaxKind` must *not* be greater than that of the last token `SyntaxKind`.
+/// See #17037.
 const LAST_TOKEN_KIND_DISCRIMINANT: usize = SyntaxKind::SHEBANG as usize;
 
 impl TokenSet {
@@ -15,13 +18,13 @@ pub(crate) const fn new(kinds: &[SyntaxKind]) -> TokenSet {
         let mut res = [0; 3];
         let mut i = 0;
         while i < kinds.len() {
-            let kind = kinds[i];
+            let discriminant = kinds[i] as usize;
             debug_assert!(
-                kind as usize <= LAST_TOKEN_KIND_DISCRIMINANT,
+                discriminant <= LAST_TOKEN_KIND_DISCRIMINANT,
                 "Expected a token `SyntaxKind`"
             );
-            let idx = kind as usize / 64;
-            res[idx] |= mask(kind);
+            let idx = discriminant / 64;
+            res[idx] |= 1 << (discriminant % 64);
             i += 1;
         }
         TokenSet(res)
@@ -32,20 +35,17 @@ pub(crate) const fn union(self, other: TokenSet) -> TokenSet {
     }
 
     pub(crate) const fn contains(&self, kind: SyntaxKind) -> bool {
+        let discriminant = kind as usize;
         debug_assert!(
-            kind as usize <= LAST_TOKEN_KIND_DISCRIMINANT,
+            discriminant <= LAST_TOKEN_KIND_DISCRIMINANT,
             "Expected a token `SyntaxKind`"
         );
-        let idx = kind as usize / 64;
-        self.0[idx] & mask(kind) != 0
+        let idx = discriminant / 64;
+        let mask = 1 << (discriminant % 64);
+        self.0[idx] & mask != 0
     }
 }
 
-const fn mask(kind: SyntaxKind) -> u64 {
-    debug_assert!(kind as usize <= LAST_TOKEN_KIND_DISCRIMINANT, "Expected a token `SyntaxKind`");
-    1 << (kind as usize % 64)
-}
-
 #[test]
 fn token_set_works_for_tokens() {
     use crate::SyntaxKind::*;