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>) -> Option<()> {
215    let mut depth = 1;
216    while depth > 0 {
217        let rem = lex.remainder();
218        if rem.is_empty() {
219            return None;
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    Some(())
235}
236
237fn parse_ident(lex: &mut logos::Lexer<Token>) -> Token {
238    let ident = lex.slice().to_string();
239    if RESERVED_WORDS.iter().contains(&ident.as_str()) {
240        Token::ReservedWord(ident)
241    } else {
242        Token::Ident(ident)
243    }
244}
245
246#[derive(Default, Clone, Debug, PartialEq)]
247pub struct LexerState {
248    depth: i32,
249    template_depths: Vec<i32>,
250    lookahead: Option<Token>,
251}
252
253// following the spec at this date: https://www.w3.org/TR/2024/WD-WGSL-20240731/
254#[derive(Logos, Clone, Debug, PartialEq)]
255#[logos(
256    // see blankspace and line breaks: https://www.w3.org/TR/WGSL/#blankspace-and-line-breaks
257    skip r"[\s\u0085\u200e\u200f\u2028\u2029]+", // blankspace
258    extras = LexerState,
259    error = ParseError)]
260pub enum Token {
261    // HACK: parsing entrypoints.
262    // See: https://github.com/lalrpop/lalrpop/issues/65#issuecomment-516769995
263    // The first token provided by the lexer to the parser must be one of these tokens.
264    // They tell the parser which syntax node to parse. The parser returns a union type
265    // containing the corresponding syntax node type.
266    EntryPointTryTemplateList,
267    EntryPointTranslationUnit,
268    EntryPointGlobalDecl,
269    EntryPointLiteral,
270    EntryPointGlobalDirective,
271    EntryPointExpression,
272    EntryPointStatement,
273
274    // wesl import extension
275    EntryPointImportStatement,
276
277    #[token("//", parse_line_comment)]
278    LineComment,
279    #[token("/*", parse_block_comment, priority = 2)]
280    BlockComment,
281    // the parse_ident function can return either Token::Ident or Token::ReservedWord.
282    // Token::Ignored variant is never produced.
283    // It serves as a placeholder for running parse_ident.
284    #[regex(
285        r#"([_\p{XID_Start}][\p{XID_Continue}]+)|([\p{XID_Start}])"#,
286        parse_ident,
287        priority = 1
288    )]
289    Ignored,
290    // syntactic tokens
291    // https://www.w3.org/TR/WGSL/#syntactic-tokens
292    #[token("&")]
293    SymAnd,
294    #[token("&&", maybe_fail_template)]
295    SymAndAnd,
296    #[token("->")]
297    SymArrow,
298    #[token("@")]
299    SymAttr,
300    #[token("/")]
301    SymForwardSlash,
302    #[token("!")]
303    SymBang,
304    #[token("[", incr_depth)]
305    SymBracketLeft,
306    #[token("]", decr_depth)]
307    SymBracketRight,
308    #[token("{")]
309    SymBraceLeft,
310    #[token("}")]
311    SymBraceRight,
312    #[token(":")]
313    SymColon,
314    #[token(",")]
315    SymComma,
316    #[token("=")]
317    SymEqual,
318    #[token("==")]
319    SymEqualEqual,
320    #[token("!=")]
321    SymNotEqual,
322    #[token(">", |lex| maybe_template_end(lex, Token::SymGreaterThan, None))]
323    SymGreaterThan,
324    #[token(">=", |lex| maybe_template_end(lex, Token::SymGreaterThanEqual, Some(Token::SymEqual)))]
325    SymGreaterThanEqual,
326    #[token(">>", |lex| maybe_template_end(lex, Token::SymShiftRight, Some(Token::SymGreaterThan)))]
327    SymShiftRight,
328    #[token("<")]
329    SymLessThan,
330    #[token("<=")]
331    SymLessThanEqual,
332    #[token("<<")]
333    SymShiftLeft,
334    #[token("%")]
335    SymModulo,
336    #[token("-")]
337    SymMinus,
338    #[token("--")]
339    SymMinusMinus,
340    #[token(".")]
341    SymPeriod,
342    #[token("+")]
343    SymPlus,
344    #[token("++")]
345    SymPlusPlus,
346    #[token("|")]
347    SymOr,
348    #[token("||", maybe_fail_template)]
349    SymOrOr,
350    #[token("(", incr_depth)]
351    SymParenLeft,
352    #[token(")", decr_depth)]
353    SymParenRight,
354    #[token(";")]
355    SymSemicolon,
356    #[token("*")]
357    SymStar,
358    #[token("~")]
359    SymTilde,
360    #[token("_")]
361    SymUnderscore,
362    #[token("^")]
363    SymXor,
364    #[token("+=")]
365    SymPlusEqual,
366    #[token("-=")]
367    SymMinusEqual,
368    #[token("*=")]
369    SymTimesEqual,
370    #[token("/=")]
371    SymDivisionEqual,
372    #[token("%=")]
373    SymModuloEqual,
374    #[token("&=")]
375    SymAndEqual,
376    #[token("|=")]
377    SymOrEqual,
378    #[token("^=")]
379    SymXorEqual,
380    #[token(">>=", |lex| maybe_template_end(lex, Token::SymShiftRightAssign, Some(Token::SymGreaterThanEqual)))]
381    SymShiftRightAssign,
382    #[token("<<=")]
383    SymShiftLeftAssign,
384
385    // keywords
386    // https://www.w3.org/TR/WGSL/#keyword-summary
387    #[token("alias")]
388    KwAlias,
389    #[token("break")]
390    KwBreak,
391    #[token("case")]
392    KwCase,
393    #[token("const", priority = 2)]
394    KwConst,
395    #[token("const_assert")]
396    KwConstAssert,
397    #[token("continue")]
398    KwContinue,
399    #[token("continuing")]
400    KwContinuing,
401    #[token("default")]
402    KwDefault,
403    #[token("diagnostic")]
404    KwDiagnostic,
405    #[token("discard")]
406    KwDiscard,
407    #[token("else")]
408    KwElse,
409    #[token("enable")]
410    KwEnable,
411    #[token("false")]
412    KwFalse,
413    #[token("fn")]
414    KwFn,
415    #[token("for")]
416    KwFor,
417    #[token("if")]
418    KwIf,
419    #[token("let")]
420    KwLet,
421    #[token("loop")]
422    KwLoop,
423    #[token("override")]
424    KwOverride,
425    #[token("requires")]
426    KwRequires,
427    #[token("return")]
428    KwReturn,
429    #[token("struct")]
430    KwStruct,
431    #[token("switch")]
432    KwSwitch,
433    #[token("true")]
434    KwTrue,
435    #[token("var")]
436    KwVar,
437    #[token("while")]
438    KwWhile,
439
440    // Idents and ReservedWord tokens are parsed on the Ignored variant, because of a current
441    // limitation of logos. See logos#295.
442    Ident(String),
443    // variant produced by parse_ident for reserved words.
444    // Reserved words can be used in context-dependent words, e.g. attribute names and module names.
445    ReservedWord(String),
446
447    #[regex(r#"0|[1-9]\d*"#, parse_dec_abstract_int)]
448    #[regex(r#"0[xX][\da-fA-F]+"#, parse_hex_abstract_int)]
449    AbstractInt(i64),
450    #[regex(r#"(\d+\.\d*|\.\d+)([eE][+-]?\d+)?"#, parse_dec_abs_float)]
451    #[regex(r#"\d+[eE][+-]?\d+"#, parse_dec_abs_float)]
452    #[regex(r#"0[xX][\da-fA-F]+\.[\da-fA-F]*([pP][+-]?\d+)?"#, parse_hex_abs_float)]
453    #[regex(r#"0[xX]\.[\da-fA-F]+([pP][+-]?\d+)?"#, parse_hex_abs_float)]
454    #[regex(r#"0[xX][\da-fA-F]+[pP][+-]?\d+"#, parse_hex_abs_float)]
455    AbstractFloat(f64),
456    #[regex(r#"(0|[1-9]\d*)i"#, parse_dec_i32)]
457    #[regex(r#"0[xX][\da-fA-F]+i"#, parse_hex_i32)]
458    I32(i32),
459    #[regex(r#"(0|[1-9]\d*)u"#, parse_dec_u32)]
460    #[regex(r#"0[xX][\da-fA-F]+u"#, parse_hex_u32)]
461    U32(u32),
462    #[regex(r#"(\d+\.\d*|\.\d+)([eE][+-]?\d+)?f"#, parse_dec_f32)]
463    #[regex(r#"\d+([eE][+-]?\d+)?f"#, parse_dec_f32)]
464    #[regex(r#"0[xX][\da-fA-F]+\.[\da-fA-F]*[pP][+-]?\d+f"#, parse_hex_f32)]
465    #[regex(r#"0[xX]\.[\da-fA-F]+[pP][+-]?\d+f"#, parse_hex_f32)]
466    #[regex(r#"0[xX][\da-fA-F]+[pP][+-]?\d+f"#, parse_hex_f32)]
467    F32(f32),
468    #[regex(r#"(\d+\.\d*|\.\d+)([eE][+-]?\d+)?h"#, parse_dec_f16)]
469    #[regex(r#"\d+([eE][+-]?\d+)?h"#, parse_dec_f16)]
470    #[regex(r#"0[xX][\da-fA-F]+\.[\da-fA-F]*[pP][+-]?\d+h"#, parse_hex_f16)]
471    #[regex(r#"0[xX]\.[\da-fA-F]+[pP][+-]?\d+h"#, parse_hex_f16)]
472    #[regex(r#"0[xX][\da-fA-F]+[pP][+-]?\d+h"#, parse_hex_f16)]
473    F16(f32),
474    TemplateArgsStart,
475    TemplateArgsEnd,
476
477    // naga extension
478    #[regex(r#"(0|[1-9]\d*)li"#, parse_dec_i64)]
479    #[regex(r#"0[xX][\da-fA-F]+li"#, parse_hex_i64)]
480    I64(i64),
481    #[regex(r#"(0|[1-9]\d*)lu"#, parse_dec_u64)]
482    #[regex(r#"0[xX][\da-fA-F]+lu"#, parse_hex_u64)]
483    U64(u64),
484    #[regex(r#"(\d+\.\d*|\.\d+)([eE][+-]?\d+)?lf"#, parse_dec_f64)]
485    #[regex(r#"\d+([eE][+-]?\d+)?lf"#, parse_dec_f64)]
486    #[regex(r#"0[xX][\da-fA-F]+\.[\da-fA-F]*[pP][+-]?\d+lf"#, parse_hex_f64)]
487    #[regex(r#"0[xX]\.[\da-fA-F]+[pP][+-]?\d+lf"#, parse_hex_f64)]
488    #[regex(r#"0[xX][\da-fA-F]+[pP][+-]?\d+lf"#, parse_hex_f64)]
489    F64(f64),
490
491    // extension: wesl-visibility
492    // https://github.com/webgpu-tools/wesl-spec/blob/main/visibility.md
493    // date: 2026-09-06, hash: af8236a875332f672a0fc5bcc4f9a034aec4f1ee
494    #[token("public")]
495    KwPublic,
496    #[token("private")]
497    KwPrivate,
498
499    // extension: wesl-imports
500    // https://github.com/webgpu-tools/wesl-spec/blob/imports-update/Imports.md
501    // date: 2025-01-18, hash: 2db8e7f681087db6bdcd4a254963deb5c0159775
502    #[token("::")]
503    SymColonColon,
504    #[token("self")]
505    KwSelf,
506    #[token("super")]
507    KwSuper,
508    #[token("package")]
509    KwPackage,
510    #[token("as")]
511    KwAs,
512    #[token("import")]
513    KwImport,
514}
515
516impl Token {
517    pub fn is_trivia(&self) -> bool {
518        matches!(
519            self,
520            Token::LineComment | Token::BlockComment | Token::Ignored
521        )
522    }
523
524    #[allow(unused)]
525    pub fn is_symbol(&self) -> bool {
526        matches!(
527            self,
528            Token::SymAnd
529                | Token::SymAndAnd
530                | Token::SymArrow
531                | Token::SymAttr
532                | Token::SymForwardSlash
533                | Token::SymBang
534                | Token::SymBracketLeft
535                | Token::SymBracketRight
536                | Token::SymBraceLeft
537                | Token::SymBraceRight
538                | Token::SymColon
539                | Token::SymComma
540                | Token::SymEqual
541                | Token::SymEqualEqual
542                | Token::SymNotEqual
543                | Token::SymGreaterThan
544                | Token::SymGreaterThanEqual
545                | Token::SymShiftRight
546                | Token::SymLessThan
547                | Token::SymLessThanEqual
548                | Token::SymShiftLeft
549                | Token::SymModulo
550                | Token::SymMinus
551                | Token::SymMinusMinus
552                | Token::SymPeriod
553                | Token::SymPlus
554                | Token::SymPlusPlus
555                | Token::SymOr
556                | Token::SymOrOr
557                | Token::SymParenLeft
558                | Token::SymParenRight
559                | Token::SymSemicolon
560                | Token::SymStar
561                | Token::SymTilde
562                | Token::SymUnderscore
563                | Token::SymXor
564                | Token::SymPlusEqual
565                | Token::SymMinusEqual
566                | Token::SymTimesEqual
567                | Token::SymDivisionEqual
568                | Token::SymModuloEqual
569                | Token::SymAndEqual
570                | Token::SymOrEqual
571                | Token::SymXorEqual
572                | Token::SymShiftRightAssign
573                | Token::SymShiftLeftAssign
574        )
575    }
576
577    #[allow(unused)]
578    pub fn is_keyword(&self) -> bool {
579        matches!(
580            self,
581            Token::KwAlias
582                | Token::KwBreak
583                | Token::KwCase
584                | Token::KwConst
585                | Token::KwConstAssert
586                | Token::KwContinue
587                | Token::KwContinuing
588                | Token::KwDefault
589                | Token::KwDiagnostic
590                | Token::KwDiscard
591                | Token::KwElse
592                | Token::KwEnable
593                | Token::KwFalse
594                | Token::KwFn
595                | Token::KwFor
596                | Token::KwIf
597                | Token::KwLet
598                | Token::KwLoop
599                | Token::KwOverride
600                | Token::KwRequires
601                | Token::KwReturn
602                | Token::KwStruct
603                | Token::KwSwitch
604                | Token::KwTrue
605                | Token::KwVar
606                | Token::KwWhile
607        )
608    }
609
610    #[allow(unused)]
611    pub fn is_numeric_literal(&self) -> bool {
612        matches!(
613            self,
614            Token::AbstractInt(_)
615                | Token::AbstractFloat(_)
616                | Token::I32(_)
617                | Token::U32(_)
618                | Token::F32(_)
619                | Token::F16(_)
620        )
621    }
622}
623
624impl Display for Token {
625    /// This display implementation is used for error messages.
626    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
627        match self {
628            Token::EntryPointTryTemplateList => f.write_str("EntryPointTryTemplateList"),
629            Token::EntryPointTranslationUnit => f.write_str("EntryPointTranslationUnit"),
630            Token::EntryPointGlobalDecl => f.write_str("EntryPointGlobalDecl"),
631            Token::EntryPointLiteral => f.write_str("EntryPointLiteral"),
632            Token::EntryPointGlobalDirective => f.write_str("EntryPointGlobalDirective"),
633            Token::EntryPointExpression => f.write_str("EntryPointExpression"),
634            Token::EntryPointStatement => f.write_str("EntryPointStatement"),
635
636            // wesl extension
637            Token::EntryPointImportStatement => f.write_str("EntryPointImportStatement"),
638
639            Token::LineComment => f.write_str("// line comment"),
640            Token::BlockComment => f.write_str("/* block comment */"),
641            Token::Ignored => unreachable!(),
642            Token::SymAnd => f.write_str("&"),
643            Token::SymAndAnd => f.write_str("&&"),
644            Token::SymArrow => f.write_str("->"),
645            Token::SymAttr => f.write_str("@"),
646            Token::SymForwardSlash => f.write_str("/"),
647            Token::SymBang => f.write_str("!"),
648            Token::SymBracketLeft => f.write_str("["),
649            Token::SymBracketRight => f.write_str("]"),
650            Token::SymBraceLeft => f.write_str("{"),
651            Token::SymBraceRight => f.write_str("}"),
652            Token::SymColon => f.write_str(":"),
653            Token::SymComma => f.write_str(","),
654            Token::SymEqual => f.write_str("="),
655            Token::SymEqualEqual => f.write_str("=="),
656            Token::SymNotEqual => f.write_str("!="),
657            Token::SymGreaterThan => f.write_str(">"),
658            Token::SymGreaterThanEqual => f.write_str(">="),
659            Token::SymShiftRight => f.write_str(">>"),
660            Token::SymLessThan => f.write_str("<"),
661            Token::SymLessThanEqual => f.write_str("<="),
662            Token::SymShiftLeft => f.write_str("<<"),
663            Token::SymModulo => f.write_str("%"),
664            Token::SymMinus => f.write_str("-"),
665            Token::SymMinusMinus => f.write_str("--"),
666            Token::SymPeriod => f.write_str("."),
667            Token::SymPlus => f.write_str("+"),
668            Token::SymPlusPlus => f.write_str("++"),
669            Token::SymOr => f.write_str("|"),
670            Token::SymOrOr => f.write_str("||"),
671            Token::SymParenLeft => f.write_str("("),
672            Token::SymParenRight => f.write_str(")"),
673            Token::SymSemicolon => f.write_str(";"),
674            Token::SymStar => f.write_str("*"),
675            Token::SymTilde => f.write_str("~"),
676            Token::SymUnderscore => f.write_str("_"),
677            Token::SymXor => f.write_str("^"),
678            Token::SymPlusEqual => f.write_str("+="),
679            Token::SymMinusEqual => f.write_str("-="),
680            Token::SymTimesEqual => f.write_str("*="),
681            Token::SymDivisionEqual => f.write_str("/="),
682            Token::SymModuloEqual => f.write_str("%="),
683            Token::SymAndEqual => f.write_str("&="),
684            Token::SymOrEqual => f.write_str("|="),
685            Token::SymXorEqual => f.write_str("^="),
686            Token::SymShiftRightAssign => f.write_str(">>="),
687            Token::SymShiftLeftAssign => f.write_str("<<="),
688            Token::KwAlias => f.write_str("alias"),
689            Token::KwBreak => f.write_str("break"),
690            Token::KwCase => f.write_str("case"),
691            Token::KwConst => f.write_str("const"),
692            Token::KwConstAssert => f.write_str("const_assert"),
693            Token::KwContinue => f.write_str("continue"),
694            Token::KwContinuing => f.write_str("continuing"),
695            Token::KwDefault => f.write_str("default"),
696            Token::KwDiagnostic => f.write_str("diagnostic"),
697            Token::KwDiscard => f.write_str("discard"),
698            Token::KwElse => f.write_str("else"),
699            Token::KwEnable => f.write_str("enable"),
700            Token::KwFalse => f.write_str("false"),
701            Token::KwFn => f.write_str("fn"),
702            Token::KwFor => f.write_str("for"),
703            Token::KwIf => f.write_str("if"),
704            Token::KwLet => f.write_str("let"),
705            Token::KwLoop => f.write_str("loop"),
706            Token::KwOverride => f.write_str("override"),
707            Token::KwRequires => f.write_str("requires"),
708            Token::KwReturn => f.write_str("return"),
709            Token::KwStruct => f.write_str("struct"),
710            Token::KwSwitch => f.write_str("switch"),
711            Token::KwTrue => f.write_str("true"),
712            Token::KwVar => f.write_str("var"),
713            Token::KwWhile => f.write_str("while"),
714            Token::Ident(s) => write!(f, "identifier `{s}`"),
715            Token::ReservedWord(s) => write!(f, "reserved word `{s}`"),
716            Token::AbstractInt(n) => write!(f, "{n}"),
717            Token::AbstractFloat(n) => write!(f, "{n}"),
718            Token::I32(n) => write!(f, "{n}i"),
719            Token::U32(n) => write!(f, "{n}u"),
720            Token::F32(n) => write!(f, "{n}f"),
721            Token::F16(n) => write!(f, "{n}h"),
722            Token::TemplateArgsStart => f.write_str("start of template"),
723            Token::TemplateArgsEnd => f.write_str("end of template"),
724
725            // naga extension
726            Token::I64(n) => write!(f, "{n}li"),
727            Token::U64(n) => write!(f, "{n}lu"),
728            Token::F64(n) => write!(f, "{n}lf"),
729
730            Token::KwPublic => f.write_str("public"),
731            Token::KwPrivate => f.write_str("private"),
732
733            // wesl import extension
734            Token::SymColonColon => f.write_str("::"),
735            Token::KwSelf => f.write_str("self"),
736            Token::KwSuper => f.write_str("super"),
737            Token::KwPackage => f.write_str("package"),
738            Token::KwAs => f.write_str("as"),
739            Token::KwImport => f.write_str("import"),
740        }
741    }
742}
743
744type Spanned<Tok, Loc, ParseError> = Result<(Loc, Tok, Loc), (Loc, ParseError, Loc)>;
745type NextToken = Option<(Result<Token, ParseError>, Span)>;
746
747#[derive(Clone)]
748pub struct Lexer<'s> {
749    source: &'s str,
750    token_stream: SpannedIter<'s, Token>,
751    next_token: NextToken,
752    recognizing_template: bool,
753    opened_templates: u32,
754}
755
756impl<'s> Lexer<'s> {
757    pub fn new(source: &'s str) -> Self {
758        let mut token_stream = Token::lexer_with_extras(source, LexerState::default()).spanned();
759        let next_token =
760            token_stream.find(|(tok, _)| !tok.as_ref().is_ok_and(|tok| tok.is_trivia()));
761
762        Self {
763            source,
764            token_stream,
765            next_token,
766            recognizing_template: false,
767            opened_templates: 0,
768        }
769    }
770
771    fn take_two_tokens(&mut self) -> (NextToken, NextToken) {
772        let mut tok1 = self.next_token.take();
773
774        let lookahead = self.token_stream.extras.lookahead.take();
775        let tok2 = match lookahead {
776            Some(tok) => {
777                let (_, span1) = tok1.as_mut().unwrap(); // safety: lookahead implies lexer looked at a `<` token
778                let span2 = span1.start + 1..span1.end;
779                Some((Ok(tok), span2))
780            }
781            None => self
782                .token_stream
783                .find(|(tok, _)| !tok.as_ref().is_ok_and(|tok| tok.is_trivia())),
784        };
785
786        (tok1, tok2)
787    }
788
789    fn next_token(&mut self) -> NextToken {
790        let (cur, mut next) = self.take_two_tokens();
791
792        let (cur_tok, cur_span) = match cur {
793            Some((Ok(tok), span)) => (tok, span),
794            Some((Err(e), span)) => return Some((Err(e), span)),
795            None => return None,
796        };
797
798        if let Some((Ok(next_tok), next_span)) = &mut next
799            && (matches!(cur_tok, Token::Ident(_)) || cur_tok.is_keyword())
800            && *next_tok == Token::SymLessThan
801        {
802            let source = &self.source[next_span.start..];
803            if recognize_template_list(source) {
804                *next_tok = Token::TemplateArgsStart;
805                let cur_depth = self.token_stream.extras.depth;
806                self.token_stream.extras.template_depths.push(cur_depth);
807                self.opened_templates += 1;
808            }
809        }
810
811        // if we finished recognition of a template
812        if self.recognizing_template && cur_tok == Token::TemplateArgsEnd {
813            self.opened_templates -= 1;
814            if self.opened_templates == 0 {
815                next = None; // push eof after end of template
816            }
817        }
818
819        self.next_token = next;
820        Some((Ok(cur_tok), cur_span))
821    }
822}
823
824/// Returns `true` if the source starts with a valid template list.
825///
826/// ## Specification
827///
828/// [3.9. Template Lists](https://www.w3.org/TR/WGSL/#template-lists-sec)
829///
830/// Contrary to the specification [template list discovery algorithm], this function also
831/// checks that the template is syntactically valid (syntax: [*template_list*]).
832///
833/// [template list discovery algorigthm]: https://www.w3.org/TR/WGSL/#template-list-discovery
834/// [*template_list*]: https://www.w3.org/TR/WGSL/#syntax-template_list
835pub fn recognize_template_list(source: &str) -> bool {
836    let mut lexer = Lexer::new(source);
837    match lexer.next_token {
838        Some((Ok(ref mut t), _)) if *t == Token::SymLessThan => *t = Token::TemplateArgsStart,
839        _ => return false,
840    };
841    lexer.recognizing_template = true;
842    lexer.opened_templates = 1;
843    lexer.token_stream.extras.template_depths.push(0);
844    crate::parser::recognize_template_list(lexer).is_ok()
845}
846
847#[test]
848fn test_recognize_template() {
849    // cases from the WGSL spec
850    assert!(recognize_template_list("<i32,select(2,3,a>b)>"));
851    assert!(!recognize_template_list("<d]>"));
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    assert!(recognize_template_list("<(B==C)>"));
857    // more cases
858    assert!(recognize_template_list("<X>"));
859    assert!(recognize_template_list("<X<Y>>"));
860    assert!(recognize_template_list("<X<Y<Z>>>"));
861    assert!(!recognize_template_list(""));
862    assert!(!recognize_template_list(""));
863    assert!(!recognize_template_list("<>"));
864    assert!(!recognize_template_list("<b || c>d"));
865}
866
867pub trait TokenIterator: IntoIterator<Item = Spanned<Token, usize, ParseError>> {}
868
869impl Iterator for Lexer<'_> {
870    type Item = Spanned<Token, usize, ParseError>;
871
872    fn next(&mut self) -> Option<Self::Item> {
873        let tok = self.next_token();
874        tok.map(|(tok, span)| match tok {
875            Ok(tok) => Ok((span.start, tok, span.end)),
876            Err(err) => Err((span.start, err, span.end)),
877        })
878    }
879}
880
881impl TokenIterator for Lexer<'_> {}
882
883#[test]
884fn test_lexer_errors() {
885    assert!(crate::parse_str("const value = 1 $ + 2;").is_err());
886    assert!(crate::parse_str("/* outer /* inner */").is_err());
887    assert!(crate::parse_str("/* outer /* inner */ */").is_ok());
888}