1use crate::parsing::composition::child_phase;
2use crate::region::wrap_anonymous;
3use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
4
5pub const TOK_NONE: u32 = 0;
7pub const TOK_IDENTIFIER: u32 = 1;
9pub const TOK_NUMBER: u32 = 2;
11pub const TOK_STRING: u32 = 3;
13pub const TOK_NEWLINE: u32 = 4;
15pub const TOK_COMMENT: u32 = 5;
17
18pub const TOK_LPAREN: u32 = 10;
20pub const TOK_RPAREN: u32 = 11;
22pub const TOK_LBRACKET: u32 = 12;
24pub const TOK_RBRACKET: u32 = 13;
26pub const TOK_LBRACE: u32 = 14;
28pub const TOK_RBRACE: u32 = 15;
30pub const TOK_COLON: u32 = 16;
32pub const TOK_COMMA: u32 = 17;
34pub const TOK_DOT: u32 = 18;
36pub const TOK_EQ: u32 = 19;
38pub const TOK_AT: u32 = 20;
40pub const TOK_STAR: u32 = 21;
42
43pub const TOK_DEF: u32 = 100;
45pub const TOK_ASYNC: u32 = 101;
47pub const TOK_CLASS: u32 = 102;
49pub const TOK_IMPORT: u32 = 103;
51pub const TOK_FROM: u32 = 104;
53pub const TOK_AS: u32 = 105;
55pub const TOK_WITH: u32 = 106;
57pub const TOK_AWAIT: u32 = 107;
59pub const TOK_MATCH: u32 = 108;
61pub const TOK_CASE: u32 = 109;
63pub const TOK_EXCEPT: u32 = 110;
65
66fn load_byte(buffer: &str, index: Expr) -> Expr {
67 Expr::bitand(Expr::load(buffer, index), Expr::u32(0xFF))
68}
69
70fn ascii(ch: u8) -> Expr {
71 Expr::u32(ch as u32)
72}
73
74fn is_between(value: Expr, start: u8, end: u8) -> Expr {
75 Expr::and(
76 Expr::ge(value.clone(), ascii(start)),
77 Expr::le(value, ascii(end)),
78 )
79}
80
81fn is_alpha(value: Expr) -> Expr {
82 Expr::or(
83 is_between(value.clone(), b'a', b'z'),
84 is_between(value, b'A', b'Z'),
85 )
86}
87
88fn is_ident_continue(value: Expr) -> Expr {
89 Expr::or(
90 Expr::or(
91 is_alpha(value.clone()),
92 is_between(value.clone(), b'0', b'9'),
93 ),
94 Expr::eq(value, ascii(b'_')),
95 )
96}
97
98fn is_ident_start(value: Expr) -> Expr {
99 Expr::or(is_alpha(value.clone()), Expr::eq(value, ascii(b'_')))
100}
101
102fn keyword_match(haystack: &str, base: Expr, len_var: &str, word: &[u8]) -> Expr {
103 let mut expr = Expr::eq(Expr::var(len_var), Expr::u32(word.len() as u32));
104 for (offset, byte) in word.iter().enumerate() {
105 expr = Expr::and(
106 expr,
107 Expr::eq(
108 load_byte(haystack, Expr::add(base.clone(), Expr::u32(offset as u32))),
109 ascii(*byte),
110 ),
111 );
112 }
113 expr
114}
115
116fn classify_keyword(haystack: &str, base: Expr) -> Vec<Node> {
117 vec![
118 Node::if_then(
119 keyword_match(haystack, base.clone(), "token_len", b"def"),
120 vec![Node::assign("token_type", Expr::u32(TOK_DEF))],
121 ),
122 Node::if_then(
123 keyword_match(haystack, base.clone(), "token_len", b"async"),
124 vec![Node::assign("token_type", Expr::u32(TOK_ASYNC))],
125 ),
126 Node::if_then(
127 keyword_match(haystack, base.clone(), "token_len", b"class"),
128 vec![Node::assign("token_type", Expr::u32(TOK_CLASS))],
129 ),
130 Node::if_then(
131 keyword_match(haystack, base.clone(), "token_len", b"import"),
132 vec![Node::assign("token_type", Expr::u32(TOK_IMPORT))],
133 ),
134 Node::if_then(
135 keyword_match(haystack, base.clone(), "token_len", b"from"),
136 vec![Node::assign("token_type", Expr::u32(TOK_FROM))],
137 ),
138 Node::if_then(
139 keyword_match(haystack, base.clone(), "token_len", b"as"),
140 vec![Node::assign("token_type", Expr::u32(TOK_AS))],
141 ),
142 Node::if_then(
143 keyword_match(haystack, base.clone(), "token_len", b"with"),
144 vec![Node::assign("token_type", Expr::u32(TOK_WITH))],
145 ),
146 Node::if_then(
147 keyword_match(haystack, base.clone(), "token_len", b"await"),
148 vec![Node::assign("token_type", Expr::u32(TOK_AWAIT))],
149 ),
150 Node::if_then(
151 keyword_match(haystack, base.clone(), "token_len", b"match"),
152 vec![Node::assign("token_type", Expr::u32(TOK_MATCH))],
153 ),
154 Node::if_then(
155 keyword_match(haystack, base.clone(), "token_len", b"case"),
156 vec![Node::assign("token_type", Expr::u32(TOK_CASE))],
157 ),
158 Node::if_then(
159 keyword_match(haystack, base, "token_len", b"except"),
160 vec![Node::assign("token_type", Expr::u32(TOK_EXCEPT))],
161 ),
162 ]
163}
164
165#[must_use]
170#[allow(clippy::too_many_arguments)]
171pub fn python312_lexer(
172 haystack: &str,
173 out_tok_types: &str,
174 out_tok_starts: &str,
175 out_tok_lens: &str,
176 out_counts: &str,
177 haystack_len: u32,
178) -> Program {
179 let t = Expr::InvocationId { axis: 0 };
180 let body = vec![
181 Node::let_bind("ch", load_byte(haystack, t.clone())),
182 Node::let_bind(
183 "prev",
184 Expr::select(
185 Expr::gt(t.clone(), Expr::u32(0)),
186 load_byte(haystack, Expr::sub(t.clone(), Expr::u32(1))),
187 Expr::u32(0),
188 ),
189 ),
190 Node::let_bind("comment_scan_active", Expr::u32(1)),
191 Node::let_bind("in_comment_tail", Expr::u32(0)),
192 Node::loop_for(
193 "comment_rev",
194 Expr::u32(0),
195 t.clone(),
196 vec![Node::if_then(
197 Expr::eq(Expr::var("comment_scan_active"), Expr::u32(1)),
198 vec![
199 Node::let_bind(
200 "comment_pos",
201 Expr::sub(Expr::sub(t.clone(), Expr::u32(1)), Expr::var("comment_rev")),
202 ),
203 Node::let_bind("comment_ch", load_byte(haystack, Expr::var("comment_pos"))),
204 Node::if_then(
205 Expr::eq(Expr::var("comment_ch"), ascii(b'\n')),
206 vec![Node::assign("comment_scan_active", Expr::u32(0))],
207 ),
208 Node::if_then(
209 Expr::eq(Expr::var("comment_ch"), ascii(b'#')),
210 vec![
211 Node::assign("in_comment_tail", Expr::u32(1)),
212 Node::assign("comment_scan_active", Expr::u32(0)),
213 ],
214 ),
215 ],
216 )],
217 ),
218 Node::let_bind("emit", Expr::u32(0)),
219 Node::let_bind("token_type", Expr::u32(TOK_NONE)),
220 Node::let_bind("token_len", Expr::u32(0)),
221 Node::let_bind(
227 "is_ident_start",
228 Expr::select(is_ident_start(Expr::var("ch")), Expr::u32(1), Expr::u32(0)),
229 ),
230 Node::let_bind(
231 "prev_identish",
232 Expr::select(
233 is_ident_continue(Expr::var("prev")),
234 Expr::u32(1),
235 Expr::u32(0),
236 ),
237 ),
238 Node::if_then(
239 Expr::eq(Expr::var("ch"), ascii(b'\n')),
240 vec![
241 Node::assign("emit", Expr::u32(1)),
242 Node::assign("token_type", Expr::u32(TOK_NEWLINE)),
243 Node::assign("token_len", Expr::u32(1)),
244 ],
245 ),
246 Node::if_then(
247 Expr::and(
248 Expr::eq(Expr::var("emit"), Expr::u32(0)),
249 Expr::eq(Expr::var("ch"), ascii(b'#')),
250 ),
251 vec![
252 Node::let_bind("active", Expr::u32(1)),
253 Node::let_bind("scan_len", Expr::u32(1)),
254 Node::loop_for(
255 "j",
256 Expr::add(t.clone(), Expr::u32(1)),
257 Expr::u32(haystack_len),
258 vec![Node::if_then(
259 Expr::eq(Expr::var("active"), Expr::u32(1)),
260 vec![
261 Node::let_bind("cur", load_byte(haystack, Expr::var("j"))),
262 Node::if_then_else(
263 Expr::eq(Expr::var("cur"), ascii(b'\n')),
264 vec![Node::assign("active", Expr::u32(0))],
265 vec![Node::assign(
266 "scan_len",
267 Expr::add(Expr::var("scan_len"), Expr::u32(1)),
268 )],
269 ),
270 ],
271 )],
272 ),
273 Node::assign("emit", Expr::u32(1)),
274 Node::assign("token_type", Expr::u32(TOK_COMMENT)),
275 Node::assign("token_len", Expr::var("scan_len")),
276 ],
277 ),
278 Node::if_then(
279 Expr::and(
280 Expr::eq(Expr::var("emit"), Expr::u32(0)),
281 Expr::or(
282 Expr::eq(Expr::var("ch"), ascii(b'\'')),
283 Expr::eq(Expr::var("ch"), ascii(b'"')),
284 ),
285 ),
286 vec![
287 Node::let_bind("quote", Expr::var("ch")),
288 Node::let_bind("active", Expr::u32(1)),
289 Node::let_bind("escaped", Expr::u32(0)),
290 Node::let_bind("scan_len", Expr::u32(1)),
291 Node::loop_for(
292 "j",
293 Expr::add(t.clone(), Expr::u32(1)),
294 Expr::u32(haystack_len),
295 vec![Node::if_then(
296 Expr::eq(Expr::var("active"), Expr::u32(1)),
297 vec![
298 Node::let_bind("cur", load_byte(haystack, Expr::var("j"))),
299 Node::assign(
300 "scan_len",
301 Expr::add(Expr::var("scan_len"), Expr::u32(1)),
302 ),
303 Node::if_then_else(
304 Expr::eq(Expr::var("escaped"), Expr::u32(1)),
305 vec![Node::assign("escaped", Expr::u32(0))],
306 vec![
307 Node::if_then(
308 Expr::eq(Expr::var("cur"), ascii(b'\\')),
309 vec![Node::assign("escaped", Expr::u32(1))],
310 ),
311 Node::if_then(
312 Expr::eq(Expr::var("cur"), Expr::var("quote")),
313 vec![Node::assign("active", Expr::u32(0))],
314 ),
315 ],
316 ),
317 ],
318 )],
319 ),
320 Node::assign("emit", Expr::u32(1)),
321 Node::assign("token_type", Expr::u32(TOK_STRING)),
322 Node::assign("token_len", Expr::var("scan_len")),
323 ],
324 ),
325 Node::if_then(
326 Expr::and(
327 Expr::eq(Expr::var("emit"), Expr::u32(0)),
328 Expr::and(
329 Expr::eq(Expr::var("is_ident_start"), Expr::u32(1)),
330 Expr::eq(Expr::var("prev_identish"), Expr::u32(0)),
331 ),
332 ),
333 vec![
334 Node::let_bind("active", Expr::u32(1)),
335 Node::let_bind("scan_len", Expr::u32(0)),
336 Node::loop_for(
337 "j",
338 t.clone(),
339 Expr::u32(haystack_len),
340 vec![Node::if_then(
341 Expr::eq(Expr::var("active"), Expr::u32(1)),
342 vec![
343 Node::let_bind("cur", load_byte(haystack, Expr::var("j"))),
344 Node::if_then_else(
345 is_ident_continue(Expr::var("cur")),
346 vec![Node::assign(
347 "scan_len",
348 Expr::add(Expr::var("scan_len"), Expr::u32(1)),
349 )],
350 vec![Node::assign("active", Expr::u32(0))],
351 ),
352 ],
353 )],
354 ),
355 Node::assign("emit", Expr::u32(1)),
356 Node::assign("token_type", Expr::u32(TOK_IDENTIFIER)),
357 Node::assign("token_len", Expr::var("scan_len")),
358 ]
359 .into_iter()
360 .chain(classify_keyword(haystack, t.clone()))
361 .collect(),
362 ),
363 Node::if_then(
364 Expr::and(
365 Expr::eq(Expr::var("emit"), Expr::u32(0)),
366 Expr::and(
367 is_between(Expr::var("ch"), b'0', b'9'),
368 Expr::eq(Expr::var("prev_identish"), Expr::u32(0)),
369 ),
370 ),
371 vec![
372 Node::let_bind("active", Expr::u32(1)),
373 Node::let_bind("scan_len", Expr::u32(0)),
374 Node::loop_for(
375 "j",
376 t.clone(),
377 Expr::u32(haystack_len),
378 vec![Node::if_then(
379 Expr::eq(Expr::var("active"), Expr::u32(1)),
380 vec![
381 Node::let_bind("cur", load_byte(haystack, Expr::var("j"))),
382 Node::if_then_else(
383 Expr::or(
384 Expr::or(
385 is_between(Expr::var("cur"), b'0', b'9'),
386 Expr::eq(Expr::var("cur"), ascii(b'_')),
387 ),
388 Expr::eq(Expr::var("cur"), ascii(b'.')),
389 ),
390 vec![Node::assign(
391 "scan_len",
392 Expr::add(Expr::var("scan_len"), Expr::u32(1)),
393 )],
394 vec![Node::assign("active", Expr::u32(0))],
395 ),
396 ],
397 )],
398 ),
399 Node::assign("emit", Expr::u32(1)),
400 Node::assign("token_type", Expr::u32(TOK_NUMBER)),
401 Node::assign("token_len", Expr::var("scan_len")),
402 ],
403 ),
404 Node::if_then(
405 Expr::and(
406 Expr::eq(Expr::var("emit"), Expr::u32(0)),
407 Expr::eq(Expr::var("ch"), ascii(b'(')),
408 ),
409 vec![
410 Node::assign("emit", Expr::u32(1)),
411 Node::assign("token_type", Expr::u32(TOK_LPAREN)),
412 Node::assign("token_len", Expr::u32(1)),
413 ],
414 ),
415 Node::if_then(
416 Expr::and(
417 Expr::eq(Expr::var("emit"), Expr::u32(0)),
418 Expr::eq(Expr::var("ch"), ascii(b')')),
419 ),
420 vec![
421 Node::assign("emit", Expr::u32(1)),
422 Node::assign("token_type", Expr::u32(TOK_RPAREN)),
423 Node::assign("token_len", Expr::u32(1)),
424 ],
425 ),
426 Node::if_then(
427 Expr::and(
428 Expr::eq(Expr::var("emit"), Expr::u32(0)),
429 Expr::eq(Expr::var("ch"), ascii(b'[')),
430 ),
431 vec![
432 Node::assign("emit", Expr::u32(1)),
433 Node::assign("token_type", Expr::u32(TOK_LBRACKET)),
434 Node::assign("token_len", Expr::u32(1)),
435 ],
436 ),
437 Node::if_then(
438 Expr::and(
439 Expr::eq(Expr::var("emit"), Expr::u32(0)),
440 Expr::eq(Expr::var("ch"), ascii(b']')),
441 ),
442 vec![
443 Node::assign("emit", Expr::u32(1)),
444 Node::assign("token_type", Expr::u32(TOK_RBRACKET)),
445 Node::assign("token_len", Expr::u32(1)),
446 ],
447 ),
448 Node::if_then(
449 Expr::and(
450 Expr::eq(Expr::var("emit"), Expr::u32(0)),
451 Expr::eq(Expr::var("ch"), ascii(b'{')),
452 ),
453 vec![
454 Node::assign("emit", Expr::u32(1)),
455 Node::assign("token_type", Expr::u32(TOK_LBRACE)),
456 Node::assign("token_len", Expr::u32(1)),
457 ],
458 ),
459 Node::if_then(
460 Expr::and(
461 Expr::eq(Expr::var("emit"), Expr::u32(0)),
462 Expr::eq(Expr::var("ch"), ascii(b'}')),
463 ),
464 vec![
465 Node::assign("emit", Expr::u32(1)),
466 Node::assign("token_type", Expr::u32(TOK_RBRACE)),
467 Node::assign("token_len", Expr::u32(1)),
468 ],
469 ),
470 Node::if_then(
471 Expr::and(
472 Expr::eq(Expr::var("emit"), Expr::u32(0)),
473 Expr::eq(Expr::var("ch"), ascii(b':')),
474 ),
475 vec![
476 Node::assign("emit", Expr::u32(1)),
477 Node::assign("token_type", Expr::u32(TOK_COLON)),
478 Node::assign("token_len", Expr::u32(1)),
479 ],
480 ),
481 Node::if_then(
482 Expr::and(
483 Expr::eq(Expr::var("emit"), Expr::u32(0)),
484 Expr::eq(Expr::var("ch"), ascii(b',')),
485 ),
486 vec![
487 Node::assign("emit", Expr::u32(1)),
488 Node::assign("token_type", Expr::u32(TOK_COMMA)),
489 Node::assign("token_len", Expr::u32(1)),
490 ],
491 ),
492 Node::if_then(
493 Expr::and(
494 Expr::eq(Expr::var("emit"), Expr::u32(0)),
495 Expr::eq(Expr::var("ch"), ascii(b'.')),
496 ),
497 vec![
498 Node::assign("emit", Expr::u32(1)),
499 Node::assign("token_type", Expr::u32(TOK_DOT)),
500 Node::assign("token_len", Expr::u32(1)),
501 ],
502 ),
503 Node::if_then(
504 Expr::and(
505 Expr::eq(Expr::var("emit"), Expr::u32(0)),
506 Expr::eq(Expr::var("ch"), ascii(b'=')),
507 ),
508 vec![
509 Node::assign("emit", Expr::u32(1)),
510 Node::assign("token_type", Expr::u32(TOK_EQ)),
511 Node::assign("token_len", Expr::u32(1)),
512 ],
513 ),
514 Node::if_then(
515 Expr::and(
516 Expr::eq(Expr::var("emit"), Expr::u32(0)),
517 Expr::eq(Expr::var("ch"), ascii(b'@')),
518 ),
519 vec![
520 Node::assign("emit", Expr::u32(1)),
521 Node::assign("token_type", Expr::u32(TOK_AT)),
522 Node::assign("token_len", Expr::u32(1)),
523 ],
524 ),
525 Node::if_then(
526 Expr::and(
527 Expr::eq(Expr::var("emit"), Expr::u32(0)),
528 Expr::eq(Expr::var("ch"), ascii(b'*')),
529 ),
530 vec![
531 Node::assign("emit", Expr::u32(1)),
532 Node::assign("token_type", Expr::u32(TOK_STAR)),
533 Node::assign("token_len", Expr::u32(1)),
534 ],
535 ),
536 Node::if_then(
537 Expr::and(
538 Expr::eq(Expr::var("in_comment_tail"), Expr::u32(1)),
539 Expr::ne(Expr::var("ch"), ascii(b'\n')),
540 ),
541 vec![Node::assign("emit", Expr::u32(0))],
542 ),
543 Node::if_then(
544 Expr::eq(Expr::var("emit"), Expr::u32(1)),
545 vec![
546 Node::store(out_tok_types, t.clone(), Expr::var("token_type")),
547 Node::store(out_tok_starts, t.clone(), t.clone()),
548 Node::store(out_tok_lens, t.clone(), Expr::var("token_len")),
549 Node::let_bind(
550 "token_slot",
551 Expr::atomic_add(out_counts, Expr::u32(0), Expr::u32(1)),
552 ),
553 Node::assign("token_slot", Expr::var("token_slot")),
554 ],
555 ),
556 ];
557
558 Program::wrapped(
559 vec![
560 BufferDecl::storage(haystack, 0, BufferAccess::ReadOnly, DataType::U32)
561 .with_count(haystack_len),
562 BufferDecl::storage(out_tok_types, 1, BufferAccess::ReadWrite, DataType::U32)
563 .with_count(haystack_len),
564 BufferDecl::storage(out_tok_starts, 2, BufferAccess::ReadWrite, DataType::U32)
565 .with_count(haystack_len),
566 BufferDecl::storage(out_tok_lens, 3, BufferAccess::ReadWrite, DataType::U32)
567 .with_count(haystack_len),
568 BufferDecl::storage(out_counts, 4, BufferAccess::ReadWrite, DataType::U32)
569 .with_count(1),
570 ],
571 [256, 1, 1],
572 vec![wrap_anonymous(
573 "vyre-libs::parsing::python312_lexer",
574 vec![child_phase(
575 "vyre-libs::parsing::python312_lexer",
576 vyre_primitives::text::line_index::OP_ID,
577 vec![Node::if_then(
578 Expr::lt(t.clone(), Expr::u32(haystack_len)),
579 body,
580 )],
581 )],
582 )],
583 )
584 .with_entry_op_id("vyre-libs::parsing::python312_lexer")
585 .with_non_composable_with_self(true)
586}
587
588inventory::submit! {
589 crate::harness::OpEntry {
590 id: "vyre-libs::parsing::python312_lexer",
591 build: || python312_lexer("haystack", "tok_types", "tok_starts", "tok_lens", "counts", 16),
592 test_inputs: Some(lexer_fixture_inputs),
593 expected_output: Some(lexer_fixture_expected),
594 category: Some("parsing"),
595 }
596}
597
598fn lexer_fixture_inputs() -> Vec<Vec<Vec<u8>>> {
599 let source = b"def f(x):\n#z\n";
600 let mut haystack = vec![0u8; 16 * 4];
601 for (idx, byte) in source.iter().enumerate() {
602 haystack[idx * 4..idx * 4 + 4].copy_from_slice(&u32::from(*byte).to_le_bytes());
603 }
604 vec![vec![
605 haystack,
606 vec![0u8; 16 * 4],
607 vec![0u8; 16 * 4],
608 vec![0u8; 16 * 4],
609 vec![0u8; 4],
610 ]]
611}
612
613fn write_sparse_token(
614 tok_types: &mut [u8],
615 tok_starts: &mut [u8],
616 tok_lens: &mut [u8],
617 pos: usize,
618 tok: u32,
619 len: u32,
620) {
621 let base = pos * 4;
622 tok_types[base..base + 4].copy_from_slice(&tok.to_le_bytes());
623 tok_starts[base..base + 4].copy_from_slice(&(pos as u32).to_le_bytes());
624 tok_lens[base..base + 4].copy_from_slice(&len.to_le_bytes());
625}
626
627fn lexer_fixture_expected() -> Vec<Vec<Vec<u8>>> {
628 let mut tok_types = vec![0u8; 16 * 4];
629 let mut tok_starts = vec![0u8; 16 * 4];
630 let mut tok_lens = vec![0u8; 16 * 4];
631 for (pos, tok, len) in [
632 (0usize, TOK_DEF, 3u32),
633 (4, TOK_IDENTIFIER, 1),
634 (5, TOK_LPAREN, 1),
635 (6, TOK_IDENTIFIER, 1),
636 (7, TOK_RPAREN, 1),
637 (8, TOK_COLON, 1),
638 (9, TOK_NEWLINE, 1),
639 (10, TOK_COMMENT, 2),
640 (12, TOK_NEWLINE, 1),
641 ] {
642 write_sparse_token(
643 &mut tok_types,
644 &mut tok_starts,
645 &mut tok_lens,
646 pos,
647 tok,
648 len,
649 );
650 }
651
652 vec![vec![
653 tok_types,
654 tok_starts,
655 tok_lens,
656 9u32.to_le_bytes().to_vec(),
657 ]]
658}