Skip to main content

wgsl_parse/
lexer.rs

1//! Prefer using [`crate::parse_str`]. You shouldn't need to manipulate the lexer.
2
3use crate::error::ParseError;
4use itertools::Itertools;
5use logos::{Logos, SpannedIter};
6use std::{fmt::Display, num::NonZeroU8, sync::LazyLock};
7use wgsl_types::idents::RESERVED_WORDS;
8
9type Span = std::ops::Range<usize>;
10
11fn maybe_template_end(
12    lex: &mut logos::Lexer<Token>,
13    current: Token,
14    lookahead: Option<Token>,
15) -> Token {
16    if let Some(depth) = lex.extras.template_depths.last() {
17        // if found a ">" on the same nesting level as the opening "<", it is a template end.
18        if lex.extras.depth == *depth {
19            lex.extras.template_depths.pop();
20            // if lookahead is GreaterThan, we may have a second closing template.
21            // note that >>= can never be (TemplateEnd, TemplateEnd, Equal).
22            if let Some(depth) = lex.extras.template_depths.last() {
23                if lex.extras.depth == *depth && lookahead == Some(Token::SymGreaterThan) {
24                    lex.extras.template_depths.pop();
25                    lex.extras.lookahead = Some(Token::TemplateArgsEnd);
26                } else {
27                    lex.extras.lookahead = lookahead;
28                }
29            } else {
30                lex.extras.lookahead = lookahead;
31            }
32            return Token::TemplateArgsEnd;
33        }
34    }
35
36    current
37}
38
39// operators && and || have lower precedence than < and >.
40// therefore, this is not a template: a < b || c > d
41fn maybe_fail_template(lex: &mut logos::Lexer<Token>) -> bool {
42    if let Some(depth) = lex.extras.template_depths.last()
43        && lex.extras.depth == *depth
44    {
45        return false;
46    }
47    true
48}
49
50fn incr_depth(lex: &mut logos::Lexer<Token>) {
51    lex.extras.depth += 1;
52}
53
54fn decr_depth(lex: &mut logos::Lexer<Token>) {
55    lex.extras.depth -= 1;
56}
57
58// TODO: get rid of crate `lexical`
59
60// don't have to be super strict, the lexer regex already did the heavy lifting
61const DEC_FORMAT: u128 = lexical::NumberFormatBuilder::new().build_unchecked();
62
63// don't have to be super strict, the lexer regex already did the heavy lifting
64const HEX_FORMAT: u128 = lexical::NumberFormatBuilder::new()
65    .mantissa_radix(16)
66    .base_prefix(NonZeroU8::new(b'x'))
67    .exponent_base(NonZeroU8::new(16))
68    .exponent_radix(NonZeroU8::new(10))
69    .build_unchecked();
70
71static FLOAT_HEX_OPTIONS: LazyLock<lexical::parse_float_options::Options> = LazyLock::new(|| {
72    lexical::parse_float_options::OptionsBuilder::new()
73        .exponent(b'p')
74        .decimal_point(b'.')
75        .build()
76        .unwrap()
77});
78
79fn parse_dec_abstract_int(lex: &mut logos::Lexer<Token>) -> Option<i64> {
80    let options = &lexical::parse_integer_options::STANDARD;
81    let str = lex.slice();
82    lexical::parse_with_options::<i64, _, DEC_FORMAT>(str, options).ok()
83}
84
85fn parse_hex_abstract_int(lex: &mut logos::Lexer<Token>) -> Option<i64> {
86    let options = &lexical::parse_integer_options::STANDARD;
87    let str = lex.slice();
88    lexical::parse_with_options::<i64, _, HEX_FORMAT>(str, options).ok()
89}
90
91fn parse_dec_i32(lex: &mut logos::Lexer<Token>) -> Option<i32> {
92    let options = &lexical::parse_integer_options::STANDARD;
93    let str = lex.slice();
94    let str = &str[..str.len() - 1];
95    lexical::parse_with_options::<i32, _, DEC_FORMAT>(str, options).ok()
96}
97
98fn parse_hex_i32(lex: &mut logos::Lexer<Token>) -> Option<i32> {
99    let options = &lexical::parse_integer_options::STANDARD;
100    let str = lex.slice();
101    let str = &str[..str.len() - 1];
102    lexical::parse_with_options::<i32, _, HEX_FORMAT>(str, options).ok()
103}
104
105fn parse_dec_u32(lex: &mut logos::Lexer<Token>) -> Option<u32> {
106    let options = &lexical::parse_integer_options::STANDARD;
107    let str = lex.slice();
108    let str = &str[..str.len() - 1];
109    lexical::parse_with_options::<u32, _, DEC_FORMAT>(str, options).ok()
110}
111
112fn parse_hex_u32(lex: &mut logos::Lexer<Token>) -> Option<u32> {
113    let options = &lexical::parse_integer_options::STANDARD;
114    let str = lex.slice();
115    let str = &str[..str.len() - 1];
116    lexical::parse_with_options::<u32, _, HEX_FORMAT>(str, options).ok()
117}
118
119fn parse_dec_abs_float(lex: &mut logos::Lexer<Token>) -> Option<f64> {
120    let options = &lexical::parse_float_options::STANDARD;
121    let str = lex.slice();
122    lexical::parse_with_options::<f64, _, DEC_FORMAT>(str, options).ok()
123}
124
125fn parse_hex_abs_float(lex: &mut logos::Lexer<Token>) -> Option<f64> {
126    let str = lex.slice();
127    lexical::parse_with_options::<f64, _, HEX_FORMAT>(str, &FLOAT_HEX_OPTIONS).ok()
128}
129
130fn parse_dec_f32(lex: &mut logos::Lexer<Token>) -> Option<f32> {
131    let options = &lexical::parse_float_options::STANDARD;
132    let str = lex.slice();
133    let str = &str[..str.len() - 1];
134    lexical::parse_with_options::<f32, _, DEC_FORMAT>(str, options).ok()
135}
136
137fn parse_hex_f32(lex: &mut logos::Lexer<Token>) -> Option<f32> {
138    let str = lex.slice();
139    let str = &str[..str.len() - 1];
140    lexical::parse_with_options::<f32, _, HEX_FORMAT>(str, &FLOAT_HEX_OPTIONS).ok()
141}
142
143fn parse_dec_f16(lex: &mut logos::Lexer<Token>) -> Option<f32> {
144    let options = &lexical::parse_float_options::STANDARD;
145    let str = lex.slice();
146    let str = &str[..str.len() - 1];
147    lexical::parse_with_options::<f32, _, DEC_FORMAT>(str, options).ok()
148}
149
150fn parse_hex_f16(lex: &mut logos::Lexer<Token>) -> Option<f32> {
151    let str = lex.slice();
152    let str = &str[..str.len() - 1];
153    lexical::parse_with_options::<f32, _, HEX_FORMAT>(str, &FLOAT_HEX_OPTIONS).ok()
154}
155
156// naga extension
157fn parse_dec_i64(lex: &mut logos::Lexer<Token>) -> Option<i64> {
158    let options = &lexical::parse_integer_options::STANDARD;
159    let str = lex.slice();
160    let str = &str[..str.len() - 2];
161    lexical::parse_with_options::<i64, _, DEC_FORMAT>(str, options).ok()
162}
163
164// naga extension
165fn parse_hex_i64(lex: &mut logos::Lexer<Token>) -> Option<i64> {
166    let options = &lexical::parse_integer_options::STANDARD;
167    let str = lex.slice();
168    let str = &str[..str.len() - 2];
169    lexical::parse_with_options::<i64, _, HEX_FORMAT>(str, options).ok()
170}
171
172// naga extension
173fn parse_dec_u64(lex: &mut logos::Lexer<Token>) -> Option<u64> {
174    let options = &lexical::parse_integer_options::STANDARD;
175    let str = lex.slice();
176    let str = &str[..str.len() - 2];
177    lexical::parse_with_options::<u64, _, DEC_FORMAT>(str, options).ok()
178}
179
180// naga extension
181fn parse_hex_u64(lex: &mut logos::Lexer<Token>) -> Option<u64> {
182    let options = &lexical::parse_integer_options::STANDARD;
183    let str = lex.slice();
184    let str = &str[..str.len() - 2];
185    lexical::parse_with_options::<u64, _, HEX_FORMAT>(str, options).ok()
186}
187
188// naga extension
189fn parse_dec_f64(lex: &mut logos::Lexer<Token>) -> Option<f64> {
190    let options = &lexical::parse_float_options::STANDARD;
191    let str = lex.slice();
192    let str = &str[..str.len() - 2];
193    lexical::parse_with_options::<f64, _, DEC_FORMAT>(str, options).ok()
194}
195
196// naga extension
197fn parse_hex_f64(lex: &mut logos::Lexer<Token>) -> Option<f64> {
198    let str = lex.slice();
199    let str = &str[..str.len() - 2];
200    lexical::parse_with_options::<f64, _, HEX_FORMAT>(str, &FLOAT_HEX_OPTIONS).ok()
201}
202
203fn parse_line_comment(lex: &mut logos::Lexer<Token>) {
204    let rem = lex.remainder();
205    // see blankspace and line breaks: https://www.w3.org/TR/WGSL/#blankspace-and-line-breaks
206    let line_end = rem
207        .char_indices()
208        .find(|(_, c)| "\n\u{000B}\u{000C}\r\u{0085}\u{2028}\u{2029}".contains(*c))
209        .map(|(i, _)| i)
210        .unwrap_or(rem.len());
211    lex.bump(line_end);
212}
213
214fn parse_block_comment(lex: &mut logos::Lexer<Token>) {
215    let mut depth = 1;
216    while depth > 0 {
217        let rem = lex.remainder();
218        if rem.is_empty() {
219            break;
220        } else if rem.starts_with("/*") {
221            lex.bump(2);
222            depth += 1;
223        } else if rem.starts_with("*/") {
224            lex.bump(2);
225            depth -= 1;
226        } else {
227            let mut next_char = 1;
228            while !rem.is_char_boundary(next_char) {
229                next_char += 1;
230            }
231            lex.bump(next_char);
232        }
233    }
234}
235
236fn parse_ident(lex: &mut logos::Lexer<Token>) -> Token {
237    let ident = lex.slice().to_string();
238    if RESERVED_WORDS.iter().contains(&ident.as_str()) {
239        Token::ReservedWord(ident)
240    } else {
241        Token::Ident(ident)
242    }
243}
244
245#[derive(Default, Clone, Debug, PartialEq)]
246pub struct LexerState {
247    depth: i32,
248    template_depths: Vec<i32>,
249    lookahead: Option<Token>,
250}
251
252// following the spec at this date: https://www.w3.org/TR/2024/WD-WGSL-20240731/
253#[derive(Logos, Clone, Debug, PartialEq)]
254#[logos(
255    // see blankspace and line breaks: https://www.w3.org/TR/WGSL/#blankspace-and-line-breaks
256    skip r"[\s\u0085\u200e\u200f\u2028\u2029]+", // blankspace
257    extras = LexerState,
258    error = ParseError)]
259pub enum Token {
260    // HACK: parsing entrypoints.
261    // See: https://github.com/lalrpop/lalrpop/issues/65#issuecomment-516769995
262    // The first token provided by the lexer to the parser must be one of these tokens.
263    // They tell the parser which syntax node to parse. The parser returns a union type
264    // containing the corresponding syntax node type.
265    EntryPointTryTemplateList,
266    EntryPointTranslationUnit,
267    EntryPointGlobalDecl,
268    EntryPointLiteral,
269    EntryPointGlobalDirective,
270    EntryPointExpression,
271    EntryPointStatement,
272
273    // wesl import extension
274    EntryPointImportStatement,
275
276    #[token("//", parse_line_comment)]
277    LineComment,
278    #[token("/*", parse_block_comment, priority = 2)]
279    BlockComment,
280    // the parse_ident function can return either Token::Ident or Token::ReservedWord.
281    // Token::Ignored variant is never produced.
282    // It serves as a placeholder for running parse_ident.
283    #[regex(
284        r#"([_\p{XID_Start}][\p{XID_Continue}]+)|([\p{XID_Start}])"#,
285        parse_ident,
286        priority = 1
287    )]
288    Ignored,
289    // syntactic tokens
290    // https://www.w3.org/TR/WGSL/#syntactic-tokens
291    #[token("&")]
292    SymAnd,
293    #[token("&&", maybe_fail_template)]
294    SymAndAnd,
295    #[token("->")]
296    SymArrow,
297    #[token("@")]
298    SymAttr,
299    #[token("/")]
300    SymForwardSlash,
301    #[token("!")]
302    SymBang,
303    #[token("[", incr_depth)]
304    SymBracketLeft,
305    #[token("]", decr_depth)]
306    SymBracketRight,
307    #[token("{")]
308    SymBraceLeft,
309    #[token("}")]
310    SymBraceRight,
311    #[token(":")]
312    SymColon,
313    #[token(",")]
314    SymComma,
315    #[token("=")]
316    SymEqual,
317    #[token("==")]
318    SymEqualEqual,
319    #[token("!=")]
320    SymNotEqual,
321    #[token(">", |lex| maybe_template_end(lex, Token::SymGreaterThan, None))]
322    SymGreaterThan,
323    #[token(">=", |lex| maybe_template_end(lex, Token::SymGreaterThanEqual, Some(Token::SymEqual)))]
324    SymGreaterThanEqual,
325    #[token(">>", |lex| maybe_template_end(lex, Token::SymShiftRight, Some(Token::SymGreaterThan)))]
326    SymShiftRight,
327    #[token("<")]
328    SymLessThan,
329    #[token("<=")]
330    SymLessThanEqual,
331    #[token("<<")]
332    SymShiftLeft,
333    #[token("%")]
334    SymModulo,
335    #[token("-")]
336    SymMinus,
337    #[token("--")]
338    SymMinusMinus,
339    #[token(".")]
340    SymPeriod,
341    #[token("+")]
342    SymPlus,
343    #[token("++")]
344    SymPlusPlus,
345    #[token("|")]
346    SymOr,
347    #[token("||", maybe_fail_template)]
348    SymOrOr,
349    #[token("(", incr_depth)]
350    SymParenLeft,
351    #[token(")", decr_depth)]
352    SymParenRight,
353    #[token(";")]
354    SymSemicolon,
355    #[token("*")]
356    SymStar,
357    #[token("~")]
358    SymTilde,
359    #[token("_")]
360    SymUnderscore,
361    #[token("^")]
362    SymXor,
363    #[token("+=")]
364    SymPlusEqual,
365    #[token("-=")]
366    SymMinusEqual,
367    #[token("*=")]
368    SymTimesEqual,
369    #[token("/=")]
370    SymDivisionEqual,
371    #[token("%=")]
372    SymModuloEqual,
373    #[token("&=")]
374    SymAndEqual,
375    #[token("|=")]
376    SymOrEqual,
377    #[token("^=")]
378    SymXorEqual,
379    #[token(">>=", |lex| maybe_template_end(lex, Token::SymShiftRightAssign, Some(Token::SymGreaterThanEqual)))]
380    SymShiftRightAssign,
381    #[token("<<=")]
382    SymShiftLeftAssign,
383
384    // keywords
385    // https://www.w3.org/TR/WGSL/#keyword-summary
386    #[token("alias")]
387    KwAlias,
388    #[token("break")]
389    KwBreak,
390    #[token("case")]
391    KwCase,
392    #[token("const", priority = 2)]
393    KwConst,
394    #[token("const_assert")]
395    KwConstAssert,
396    #[token("continue")]
397    KwContinue,
398    #[token("continuing")]
399    KwContinuing,
400    #[token("default")]
401    KwDefault,
402    #[token("diagnostic")]
403    KwDiagnostic,
404    #[token("discard")]
405    KwDiscard,
406    #[token("else")]
407    KwElse,
408    #[token("enable")]
409    KwEnable,
410    #[token("false")]
411    KwFalse,
412    #[token("fn")]
413    KwFn,
414    #[token("for")]
415    KwFor,
416    #[token("if")]
417    KwIf,
418    #[token("let")]
419    KwLet,
420    #[token("loop")]
421    KwLoop,
422    #[token("override")]
423    KwOverride,
424    #[token("requires")]
425    KwRequires,
426    #[token("return")]
427    KwReturn,
428    #[token("struct")]
429    KwStruct,
430    #[token("switch")]
431    KwSwitch,
432    #[token("true")]
433    KwTrue,
434    #[token("var")]
435    KwVar,
436    #[token("while")]
437    KwWhile,
438
439    // Idents and ReservedWord tokens are parsed on the Ignored variant, because of a current
440    // limitation of logos. See logos#295.
441    Ident(String),
442    // variant produced by parse_ident for reserved words.
443    // Reserved words can be used in context-dependent words, e.g. attribute names and module names.
444    ReservedWord(String),
445
446    #[regex(r#"0|[1-9]\d*"#, parse_dec_abstract_int)]
447    #[regex(r#"0[xX][\da-fA-F]+"#, parse_hex_abstract_int)]
448    AbstractInt(i64),
449    #[regex(r#"(\d+\.\d*|\.\d+)([eE][+-]?\d+)?"#, parse_dec_abs_float)]
450    #[regex(r#"\d+[eE][+-]?\d+"#, parse_dec_abs_float)]
451    #[regex(r#"0[xX][\da-fA-F]+\.[\da-fA-F]*([pP][+-]?\d+)?"#, parse_hex_abs_float)]
452    #[regex(r#"0[xX]\.[\da-fA-F]+([pP][+-]?\d+)?"#, parse_hex_abs_float)]
453    #[regex(r#"0[xX][\da-fA-F]+[pP][+-]?\d+"#, parse_hex_abs_float)]
454    AbstractFloat(f64),
455    #[regex(r#"(0|[1-9]\d*)i"#, parse_dec_i32)]
456    #[regex(r#"0[xX][\da-fA-F]+i"#, parse_hex_i32)]
457    I32(i32),
458    #[regex(r#"(0|[1-9]\d*)u"#, parse_dec_u32)]
459    #[regex(r#"0[xX][\da-fA-F]+u"#, parse_hex_u32)]
460    U32(u32),
461    #[regex(r#"(\d+\.\d*|\.\d+)([eE][+-]?\d+)?f"#, parse_dec_f32)]
462    #[regex(r#"\d+([eE][+-]?\d+)?f"#, parse_dec_f32)]
463    #[regex(r#"0[xX][\da-fA-F]+\.[\da-fA-F]*[pP][+-]?\d+f"#, parse_hex_f32)]
464    #[regex(r#"0[xX]\.[\da-fA-F]+[pP][+-]?\d+f"#, parse_hex_f32)]
465    #[regex(r#"0[xX][\da-fA-F]+[pP][+-]?\d+f"#, parse_hex_f32)]
466    F32(f32),
467    #[regex(r#"(\d+\.\d*|\.\d+)([eE][+-]?\d+)?h"#, parse_dec_f16)]
468    #[regex(r#"\d+([eE][+-]?\d+)?h"#, parse_dec_f16)]
469    #[regex(r#"0[xX][\da-fA-F]+\.[\da-fA-F]*[pP][+-]?\d+h"#, parse_hex_f16)]
470    #[regex(r#"0[xX]\.[\da-fA-F]+[pP][+-]?\d+h"#, parse_hex_f16)]
471    #[regex(r#"0[xX][\da-fA-F]+[pP][+-]?\d+h"#, parse_hex_f16)]
472    F16(f32),
473    TemplateArgsStart,
474    TemplateArgsEnd,
475
476    // naga extension
477    #[regex(r#"(0|[1-9]\d*)li"#, parse_dec_i64)]
478    #[regex(r#"0[xX][\da-fA-F]+li"#, parse_hex_i64)]
479    I64(i64),
480    #[regex(r#"(0|[1-9]\d*)lu"#, parse_dec_u64)]
481    #[regex(r#"0[xX][\da-fA-F]+lu"#, parse_hex_u64)]
482    U64(u64),
483    #[regex(r#"(\d+\.\d*|\.\d+)([eE][+-]?\d+)?lf"#, parse_dec_f64)]
484    #[regex(r#"\d+([eE][+-]?\d+)?lf"#, parse_dec_f64)]
485    #[regex(r#"0[xX][\da-fA-F]+\.[\da-fA-F]*[pP][+-]?\d+lf"#, parse_hex_f64)]
486    #[regex(r#"0[xX]\.[\da-fA-F]+[pP][+-]?\d+lf"#, parse_hex_f64)]
487    #[regex(r#"0[xX][\da-fA-F]+[pP][+-]?\d+lf"#, parse_hex_f64)]
488    F64(f64),
489
490    // extension: wesl-visibility
491    // https://github.com/webgpu-tools/wesl-spec/blob/main/visibility.md
492    // date: 2026-09-06, hash: af8236a875332f672a0fc5bcc4f9a034aec4f1ee
493    #[token("public")]
494    KwPublic,
495    #[token("private")]
496    KwPrivate,
497
498    // extension: wesl-imports
499    // https://github.com/webgpu-tools/wesl-spec/blob/imports-update/Imports.md
500    // date: 2025-01-18, hash: 2db8e7f681087db6bdcd4a254963deb5c0159775
501    #[token("::")]
502    SymColonColon,
503    #[token("self")]
504    KwSelf,
505    #[token("super")]
506    KwSuper,
507    #[token("package")]
508    KwPackage,
509    #[token("as")]
510    KwAs,
511    #[token("import")]
512    KwImport,
513}
514
515impl Token {
516    pub fn is_trivia(&self) -> bool {
517        matches!(
518            self,
519            Token::LineComment | Token::BlockComment | Token::Ignored
520        )
521    }
522
523    #[allow(unused)]
524    pub fn is_symbol(&self) -> bool {
525        matches!(
526            self,
527            Token::SymAnd
528                | Token::SymAndAnd
529                | Token::SymArrow
530                | Token::SymAttr
531                | Token::SymForwardSlash
532                | Token::SymBang
533                | Token::SymBracketLeft
534                | Token::SymBracketRight
535                | Token::SymBraceLeft
536                | Token::SymBraceRight
537                | Token::SymColon
538                | Token::SymComma
539                | Token::SymEqual
540                | Token::SymEqualEqual
541                | Token::SymNotEqual
542                | Token::SymGreaterThan
543                | Token::SymGreaterThanEqual
544                | Token::SymShiftRight
545                | Token::SymLessThan
546                | Token::SymLessThanEqual
547                | Token::SymShiftLeft
548                | Token::SymModulo
549                | Token::SymMinus
550                | Token::SymMinusMinus
551                | Token::SymPeriod
552                | Token::SymPlus
553                | Token::SymPlusPlus
554                | Token::SymOr
555                | Token::SymOrOr
556                | Token::SymParenLeft
557                | Token::SymParenRight
558                | Token::SymSemicolon
559                | Token::SymStar
560                | Token::SymTilde
561                | Token::SymUnderscore
562                | Token::SymXor
563                | Token::SymPlusEqual
564                | Token::SymMinusEqual
565                | Token::SymTimesEqual
566                | Token::SymDivisionEqual
567                | Token::SymModuloEqual
568                | Token::SymAndEqual
569                | Token::SymOrEqual
570                | Token::SymXorEqual
571                | Token::SymShiftRightAssign
572                | Token::SymShiftLeftAssign
573        )
574    }
575
576    #[allow(unused)]
577    pub fn is_keyword(&self) -> bool {
578        matches!(
579            self,
580            Token::KwAlias
581                | Token::KwBreak
582                | Token::KwCase
583                | Token::KwConst
584                | Token::KwConstAssert
585                | Token::KwContinue
586                | Token::KwContinuing
587                | Token::KwDefault
588                | Token::KwDiagnostic
589                | Token::KwDiscard
590                | Token::KwElse
591                | Token::KwEnable
592                | Token::KwFalse
593                | Token::KwFn
594                | Token::KwFor
595                | Token::KwIf
596                | Token::KwLet
597                | Token::KwLoop
598                | Token::KwOverride
599                | Token::KwRequires
600                | Token::KwReturn
601                | Token::KwStruct
602                | Token::KwSwitch
603                | Token::KwTrue
604                | Token::KwVar
605                | Token::KwWhile
606        )
607    }
608
609    #[allow(unused)]
610    pub fn is_numeric_literal(&self) -> bool {
611        matches!(
612            self,
613            Token::AbstractInt(_)
614                | Token::AbstractFloat(_)
615                | Token::I32(_)
616                | Token::U32(_)
617                | Token::F32(_)
618                | Token::F16(_)
619        )
620    }
621}
622
623impl Display for Token {
624    /// This display implementation is used for error messages.
625    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
626        match self {
627            Token::EntryPointTryTemplateList => f.write_str("EntryPointTryTemplateList"),
628            Token::EntryPointTranslationUnit => f.write_str("EntryPointTranslationUnit"),
629            Token::EntryPointGlobalDecl => f.write_str("EntryPointGlobalDecl"),
630            Token::EntryPointLiteral => f.write_str("EntryPointLiteral"),
631            Token::EntryPointGlobalDirective => f.write_str("EntryPointGlobalDirective"),
632            Token::EntryPointExpression => f.write_str("EntryPointExpression"),
633            Token::EntryPointStatement => f.write_str("EntryPointStatement"),
634
635            // wesl extension
636            Token::EntryPointImportStatement => f.write_str("EntryPointImportStatement"),
637
638            Token::LineComment => f.write_str("// line comment"),
639            Token::BlockComment => f.write_str("/* block comment */"),
640            Token::Ignored => unreachable!(),
641            Token::SymAnd => f.write_str("&"),
642            Token::SymAndAnd => f.write_str("&&"),
643            Token::SymArrow => f.write_str("->"),
644            Token::SymAttr => f.write_str("@"),
645            Token::SymForwardSlash => f.write_str("/"),
646            Token::SymBang => f.write_str("!"),
647            Token::SymBracketLeft => f.write_str("["),
648            Token::SymBracketRight => f.write_str("]"),
649            Token::SymBraceLeft => f.write_str("{"),
650            Token::SymBraceRight => f.write_str("}"),
651            Token::SymColon => f.write_str(":"),
652            Token::SymComma => f.write_str(","),
653            Token::SymEqual => f.write_str("="),
654            Token::SymEqualEqual => f.write_str("=="),
655            Token::SymNotEqual => f.write_str("!="),
656            Token::SymGreaterThan => f.write_str(">"),
657            Token::SymGreaterThanEqual => f.write_str(">="),
658            Token::SymShiftRight => f.write_str(">>"),
659            Token::SymLessThan => f.write_str("<"),
660            Token::SymLessThanEqual => f.write_str("<="),
661            Token::SymShiftLeft => f.write_str("<<"),
662            Token::SymModulo => f.write_str("%"),
663            Token::SymMinus => f.write_str("-"),
664            Token::SymMinusMinus => f.write_str("--"),
665            Token::SymPeriod => f.write_str("."),
666            Token::SymPlus => f.write_str("+"),
667            Token::SymPlusPlus => f.write_str("++"),
668            Token::SymOr => f.write_str("|"),
669            Token::SymOrOr => f.write_str("||"),
670            Token::SymParenLeft => f.write_str("("),
671            Token::SymParenRight => f.write_str(")"),
672            Token::SymSemicolon => f.write_str(";"),
673            Token::SymStar => f.write_str("*"),
674            Token::SymTilde => f.write_str("~"),
675            Token::SymUnderscore => f.write_str("_"),
676            Token::SymXor => f.write_str("^"),
677            Token::SymPlusEqual => f.write_str("+="),
678            Token::SymMinusEqual => f.write_str("-="),
679            Token::SymTimesEqual => f.write_str("*="),
680            Token::SymDivisionEqual => f.write_str("/="),
681            Token::SymModuloEqual => f.write_str("%="),
682            Token::SymAndEqual => f.write_str("&="),
683            Token::SymOrEqual => f.write_str("|="),
684            Token::SymXorEqual => f.write_str("^="),
685            Token::SymShiftRightAssign => f.write_str(">>="),
686            Token::SymShiftLeftAssign => f.write_str("<<="),
687            Token::KwAlias => f.write_str("alias"),
688            Token::KwBreak => f.write_str("break"),
689            Token::KwCase => f.write_str("case"),
690            Token::KwConst => f.write_str("const"),
691            Token::KwConstAssert => f.write_str("const_assert"),
692            Token::KwContinue => f.write_str("continue"),
693            Token::KwContinuing => f.write_str("continuing"),
694            Token::KwDefault => f.write_str("default"),
695            Token::KwDiagnostic => f.write_str("diagnostic"),
696            Token::KwDiscard => f.write_str("discard"),
697            Token::KwElse => f.write_str("else"),
698            Token::KwEnable => f.write_str("enable"),
699            Token::KwFalse => f.write_str("false"),
700            Token::KwFn => f.write_str("fn"),
701            Token::KwFor => f.write_str("for"),
702            Token::KwIf => f.write_str("if"),
703            Token::KwLet => f.write_str("let"),
704            Token::KwLoop => f.write_str("loop"),
705            Token::KwOverride => f.write_str("override"),
706            Token::KwRequires => f.write_str("requires"),
707            Token::KwReturn => f.write_str("return"),
708            Token::KwStruct => f.write_str("struct"),
709            Token::KwSwitch => f.write_str("switch"),
710            Token::KwTrue => f.write_str("true"),
711            Token::KwVar => f.write_str("var"),
712            Token::KwWhile => f.write_str("while"),
713            Token::Ident(s) => write!(f, "identifier `{s}`"),
714            Token::ReservedWord(s) => write!(f, "reserved word `{s}`"),
715            Token::AbstractInt(n) => write!(f, "{n}"),
716            Token::AbstractFloat(n) => write!(f, "{n}"),
717            Token::I32(n) => write!(f, "{n}i"),
718            Token::U32(n) => write!(f, "{n}u"),
719            Token::F32(n) => write!(f, "{n}f"),
720            Token::F16(n) => write!(f, "{n}h"),
721            Token::TemplateArgsStart => f.write_str("start of template"),
722            Token::TemplateArgsEnd => f.write_str("end of template"),
723
724            // naga extension
725            Token::I64(n) => write!(f, "{n}li"),
726            Token::U64(n) => write!(f, "{n}lu"),
727            Token::F64(n) => write!(f, "{n}lf"),
728
729            Token::KwPublic => f.write_str("public"),
730            Token::KwPrivate => f.write_str("private"),
731
732            // wesl import extension
733            Token::SymColonColon => f.write_str("::"),
734            Token::KwSelf => f.write_str("self"),
735            Token::KwSuper => f.write_str("super"),
736            Token::KwPackage => f.write_str("package"),
737            Token::KwAs => f.write_str("as"),
738            Token::KwImport => f.write_str("import"),
739        }
740    }
741}
742
743type Spanned<Tok, Loc, ParseError> = Result<(Loc, Tok, Loc), (Loc, ParseError, Loc)>;
744type NextToken = Option<(Result<Token, ParseError>, Span)>;
745
746#[derive(Clone)]
747pub struct Lexer<'s> {
748    source: &'s str,
749    token_stream: SpannedIter<'s, Token>,
750    next_token: NextToken,
751    recognizing_template: bool,
752    opened_templates: u32,
753}
754
755impl<'s> Lexer<'s> {
756    pub fn new(source: &'s str) -> Self {
757        let mut token_stream = Token::lexer_with_extras(source, LexerState::default()).spanned();
758        let next_token =
759            token_stream.find(|(tok, _)| tok.as_ref().is_ok_and(|tok| !tok.is_trivia()));
760
761        Self {
762            source,
763            token_stream,
764            next_token,
765            recognizing_template: false,
766            opened_templates: 0,
767        }
768    }
769
770    fn take_two_tokens(&mut self) -> (NextToken, NextToken) {
771        let mut tok1 = self.next_token.take();
772
773        let lookahead = self.token_stream.extras.lookahead.take();
774        let tok2 = match lookahead {
775            Some(tok) => {
776                let (_, span1) = tok1.as_mut().unwrap(); // safety: lookahead implies lexer looked at a `<` token
777                let span2 = span1.start + 1..span1.end;
778                Some((Ok(tok), span2))
779            }
780            None => self
781                .token_stream
782                .find(|(tok, _)| tok.as_ref().is_ok_and(|tok| !tok.is_trivia())),
783        };
784
785        (tok1, tok2)
786    }
787
788    fn next_token(&mut self) -> NextToken {
789        let (cur, mut next) = self.take_two_tokens();
790
791        let (cur_tok, cur_span) = match cur {
792            Some((Ok(tok), span)) => (tok, span),
793            Some((Err(e), span)) => return Some((Err(e), span)),
794            None => return None,
795        };
796
797        if let Some((Ok(next_tok), next_span)) = &mut next
798            && (matches!(cur_tok, Token::Ident(_)) || cur_tok.is_keyword())
799            && *next_tok == Token::SymLessThan
800        {
801            let source = &self.source[next_span.start..];
802            if recognize_template_list(source) {
803                *next_tok = Token::TemplateArgsStart;
804                let cur_depth = self.token_stream.extras.depth;
805                self.token_stream.extras.template_depths.push(cur_depth);
806                self.opened_templates += 1;
807            }
808        }
809
810        // if we finished recognition of a template
811        if self.recognizing_template && cur_tok == Token::TemplateArgsEnd {
812            self.opened_templates -= 1;
813            if self.opened_templates == 0 {
814                next = None; // push eof after end of template
815            }
816        }
817
818        self.next_token = next;
819        Some((Ok(cur_tok), cur_span))
820    }
821}
822
823/// Returns `true` if the source starts with a valid template list.
824///
825/// ## Specification
826///
827/// [3.9. Template Lists](https://www.w3.org/TR/WGSL/#template-lists-sec)
828///
829/// Contrary to the specification [template list discovery algorithm], this function also
830/// checks that the template is syntactically valid (syntax: [*template_list*]).
831///
832/// [template list discovery algorigthm]: https://www.w3.org/TR/WGSL/#template-list-discovery
833/// [*template_list*]: https://www.w3.org/TR/WGSL/#syntax-template_list
834pub fn recognize_template_list(source: &str) -> bool {
835    let mut lexer = Lexer::new(source);
836    match lexer.next_token {
837        Some((Ok(ref mut t), _)) if *t == Token::SymLessThan => *t = Token::TemplateArgsStart,
838        _ => return false,
839    };
840    lexer.recognizing_template = true;
841    lexer.opened_templates = 1;
842    lexer.token_stream.extras.template_depths.push(0);
843    crate::parser::recognize_template_list(lexer).is_ok()
844}
845
846#[test]
847fn test_recognize_template() {
848    // cases from the WGSL spec
849    assert!(recognize_template_list("<i32,select(2,3,a>b)>"));
850    assert!(!recognize_template_list("<d]>"));
851    assert!(recognize_template_list("<B<<C>"));
852    assert!(recognize_template_list("<B<=C>"));
853    assert!(recognize_template_list("<(B>=C)>"));
854    assert!(recognize_template_list("<(B!=C)>"));
855    assert!(recognize_template_list("<(B==C)>"));
856    // more cases
857    assert!(recognize_template_list("<X>"));
858    assert!(recognize_template_list("<X<Y>>"));
859    assert!(recognize_template_list("<X<Y<Z>>>"));
860    assert!(!recognize_template_list(""));
861    assert!(!recognize_template_list(""));
862    assert!(!recognize_template_list("<>"));
863    assert!(!recognize_template_list("<b || c>d"));
864}
865
866pub trait TokenIterator: IntoIterator<Item = Spanned<Token, usize, ParseError>> {}
867
868impl Iterator for Lexer<'_> {
869    type Item = Spanned<Token, usize, ParseError>;
870
871    fn next(&mut self) -> Option<Self::Item> {
872        let tok = self.next_token();
873        tok.map(|(tok, span)| match tok {
874            Ok(tok) => Ok((span.start, tok, span.end)),
875            Err(err) => Err((span.start, err, span.end)),
876        })
877    }
878}
879
880impl TokenIterator for Lexer<'_> {}