internal: update TokenSet implementation

Author: davidsemakula

crates/parser/src/token_set.rs

@@ -8,9 +8,8 @@ pub(crate) use self::generated::TokenSet;
 #[test]
 fn token_set_works_for_tokens() {
     use crate::SyntaxKind::*;
-    let ts = TokenSet::new(&[EOF, SHEBANG, __LAST]);
+    let ts = TokenSet::new(&[EOF, SHEBANG]);
     assert!(ts.contains(EOF));
     assert!(ts.contains(SHEBANG));
-    assert!(ts.contains(__LAST));
     assert!(!ts.contains(PLUS));
 }

crates/parser/src/token_set/generated.rs

@@ -1,46 +1,37 @@
 //! Generated by `sourcegen_ast`, do not edit by hand.
 
 use crate::SyntaxKind;
-#[doc = r" A bit-set of `SyntaxKind`s"]
+#[doc = r" A bit-set of token `SyntaxKind`s"]
 #[derive(Clone, Copy)]
-pub(crate) struct TokenSet([u16; 17usize]);
+pub(crate) struct TokenSet([u64; 2usize]);
+const LAST_TOKEN_DISCRIMINANT: usize = 121usize;
 impl TokenSet {
-    pub(crate) const EMPTY: TokenSet = TokenSet([0; 17usize]);
+    pub(crate) const EMPTY: TokenSet = TokenSet([0; 2usize]);
     pub(crate) const fn new(kinds: &[SyntaxKind]) -> TokenSet {
-        let mut res = [0; 17usize];
+        let mut res = [0; 2usize];
         let mut i = 0;
         while i < kinds.len() {
             let kind = kinds[i];
-            let idx = kind as usize / 16usize;
+            debug_assert!(
+                (kind as usize) <= LAST_TOKEN_DISCRIMINANT,
+                "Expected a token `SyntaxKind`"
+            );
+            let idx = kind as usize / 64usize;
             res[idx] |= mask(kind);
             i += 1;
         }
         TokenSet(res)
     }
     pub(crate) const fn union(self, other: TokenSet) -> TokenSet {
-        TokenSet([
-            self.0[0usize] | other.0[0usize],
-            self.0[1usize] | other.0[1usize],
-            self.0[2usize] | other.0[2usize],
-            self.0[3usize] | other.0[3usize],
-            self.0[4usize] | other.0[4usize],
-            self.0[5usize] | other.0[5usize],
-            self.0[6usize] | other.0[6usize],
-            self.0[7usize] | other.0[7usize],
-            self.0[8usize] | other.0[8usize],
-            self.0[9usize] | other.0[9usize],
-            self.0[10usize] | other.0[10usize],
-            self.0[11usize] | other.0[11usize],
-            self.0[12usize] | other.0[12usize],
-            self.0[13usize] | other.0[13usize],
-            self.0[14usize] | other.0[14usize],
-            self.0[15usize] | other.0[15usize],
-            self.0[16usize] | other.0[16usize],
-        ])
+        TokenSet([self.0[0usize] | other.0[0usize], self.0[1usize] | other.0[1usize]])
     }
     pub(crate) const fn contains(&self, kind: SyntaxKind) -> bool {
-        let idx = kind as usize / 16usize;
+        debug_assert!((kind as usize) <= LAST_TOKEN_DISCRIMINANT, "Expected a token `SyntaxKind`");
+        let idx = kind as usize / 64usize;
         self.0[idx] & mask(kind) != 0
     }
 }
-const fn mask(kind: SyntaxKind) -> u16 { 1 << (kind as usize % 16usize) }
+const fn mask(kind: SyntaxKind) -> u64 {
+    debug_assert!((kind as usize) <= LAST_TOKEN_DISCRIMINANT, "Expected a token `SyntaxKind`");
+    1 << (kind as usize % 64usize)
+}

xtask/src/codegen/grammar.rs

@@ -466,28 +466,36 @@ fn generate_syntax_kinds(grammar: KindsSrc<'_>) -> String {
 }
 
 fn generate_token_set(grammar: KindsSrc<'_>) -> String {
-    // +3 for technical syntax kinds (i.e. `TOMBSTONE`, `EOF` and `__LAST`).
-    let n_syntax_kinds = 3
+    // +2 for `TOMBSTONE` and `EOF`.
+    let n_token_kinds = 2
         + grammar.punct.len()
         + grammar.keywords.len()
         + grammar.contextual_keywords.len()
         + grammar.literals.len()
-        + grammar.tokens.len()
-        + grammar.nodes.len();
+        + grammar.tokens.len();
 
-    let base_type_size = 16usize;
+    let base_type_size = 64usize;
     let base_type = format_ident!("u{}", base_type_size);
-    let n_chunks = (n_syntax_kinds as f32 / base_type_size as f32).ceil() as usize;
+    let n_chunks = (n_token_kinds as f32 / base_type_size as f32).ceil() as usize;
+
+    // 0 and 1 discriminants are assigned to `TOMBSTONE` and `EOF`.
+    let last_token_discriminant = n_token_kinds - 1;
 
     let union_items = (0..n_chunks).map(|idx| quote! { self.0[#idx] | other.0[#idx] });
 
+    let assert_token_kind = quote! {
+        debug_assert!((kind as usize) <= LAST_TOKEN_DISCRIMINANT, "Expected a token `SyntaxKind`");
+    };
+
     let ast = quote! {
         use crate::SyntaxKind;
 
-        /// A bit-set of `SyntaxKind`s
+        /// A bit-set of token `SyntaxKind`s
         #[derive(Clone, Copy)]
         pub(crate) struct TokenSet([#base_type; #n_chunks]);
 
+        const LAST_TOKEN_DISCRIMINANT: usize = #last_token_discriminant;
+
         impl TokenSet {
             pub(crate) const EMPTY: TokenSet = TokenSet([0; #n_chunks]);
 
@@ -496,6 +504,7 @@ fn generate_token_set(grammar: KindsSrc<'_>) -> String {
                 let mut i = 0;
                 while i < kinds.len() {
                     let kind = kinds[i];
+                    #assert_token_kind
                     let idx = kind as usize / #base_type_size;
                     res[idx] |= mask(kind);
                     i += 1;
@@ -510,12 +519,14 @@ fn generate_token_set(grammar: KindsSrc<'_>) -> String {
             }
 
             pub(crate) const fn contains(&self, kind: SyntaxKind) -> bool {
+                #assert_token_kind
                 let idx = kind as usize / #base_type_size;
                 self.0[idx] & mask(kind) != 0
             }
         }
 
         const fn mask(kind: SyntaxKind) -> #base_type {
+            #assert_token_kind
             1 << (kind as usize % #base_type_size)
         }
     };