vyre_libs/parsing/python/parse/
structure.rs

1use super::walk::{pack_sparse_tokens, DottedName, TokenPass};
2use super::{
3    find_matching_delimiter, find_matching_delimiter_into, load_u32, search_next_token,
4    search_next_token_into, search_prev_token, store_words, write_words,
5};
6use crate::parsing::python::lex::{
7    TOK_ASYNC, TOK_CLASS, TOK_COLON, TOK_COMMA, TOK_DEF, TOK_FROM, TOK_IDENTIFIER, TOK_IMPORT,
8    TOK_LBRACKET, TOK_LPAREN, TOK_RBRACKET, TOK_WITH,
9};
10use crate::parsing::python::{
11    DEF_RECORD_WORDS, IMPORT_RECORD_WORDS, INVALID_POS, WITH_RECORD_WORDS,
12};
13use vyre_foundation::ir::{Expr, Node, Program};
14
15const STRUCTURE_OP_ID: &str = "vyre-libs::parsing::python312_extract_structure";
16const IMPORTS_OP_ID: &str = "vyre-libs::parsing::python312_extract_imports";
17const WITH_BLOCKS_OP_ID: &str = "vyre-libs::parsing::python312_extract_with_blocks";
18
19fn line_index_pass<'a>(
20    op_id: &'a str,
21    tok_types: &'a str,
22    tok_starts: &'a str,
23    tok_lens: &'a str,
24    haystack_len: u32,
25) -> TokenPass<'a> {
26    TokenPass {
27        op_id,
28        child_op_id: vyre_primitives::text::line_index::OP_ID,
29        tok_types,
30        tok_starts,
31        tok_lens,
32        haystack_len,
33    }
34}
35
36/// Extract `def`, `async def`, and `class` declarations.
37#[must_use]
38pub fn python312_extract_structure(
39    tok_types: &str,
40    tok_starts: &str,
41    tok_lens: &str,
42    out_records: &str,
43    out_counts: &str,
44    haystack_len: u32,
45) -> Program {
46    let t = Expr::InvocationId { axis: 0 };
47    let mut body = vec![
48        Node::let_bind("tok", load_u32(tok_types, t.clone())),
49        Node::let_bind("emit_kind", Expr::u32(0)),
50        Node::let_bind("keyword_pos", Expr::u32(INVALID_POS)),
51        Node::if_then(
52            Expr::eq(Expr::var("tok"), Expr::u32(TOK_DEF)),
53            vec![
54                Node::assign("emit_kind", Expr::u32(1)),
55                Node::assign("keyword_pos", t.clone()),
56            ],
57        ),
58        Node::if_then(
59            Expr::eq(Expr::var("tok"), Expr::u32(TOK_CLASS)),
60            vec![
61                Node::assign("emit_kind", Expr::u32(3)),
62                Node::assign("keyword_pos", t.clone()),
63            ],
64        ),
65    ];
66    body.extend(search_next_token(
67        "async_next",
68        Expr::add(t.clone(), Expr::u32(1)),
69        tok_types,
70        haystack_len,
71    ));
72    body.push(Node::if_then(
73        Expr::and(
74            Expr::eq(Expr::var("tok"), Expr::u32(TOK_ASYNC)),
75            Expr::eq(
76                load_u32(tok_types, Expr::var("async_next")),
77                Expr::u32(TOK_DEF),
78            ),
79        ),
80        vec![
81            Node::assign("emit_kind", Expr::u32(2)),
82            Node::assign("keyword_pos", Expr::var("async_next")),
83        ],
84    ));
85    body.extend(search_next_token(
86        "name_pos",
87        Expr::add(Expr::var("keyword_pos"), Expr::u32(1)),
88        tok_types,
89        haystack_len,
90    ));
91    body.extend(search_next_token(
92        "post_name",
93        Expr::add(Expr::var("name_pos"), Expr::u32(1)),
94        tok_types,
95        haystack_len,
96    ));
97    body.extend(find_matching_delimiter(
98        "type_params_end",
99        Expr::var("post_name"),
100        tok_types,
101        haystack_len,
102        TOK_LBRACKET,
103        TOK_RBRACKET,
104    ));
105    body.push(Node::if_then(
106        Expr::and(
107            Expr::ne(Expr::var("emit_kind"), Expr::u32(0)),
108            Expr::eq(
109                load_u32(tok_types, Expr::var("name_pos")),
110                Expr::u32(TOK_IDENTIFIER),
111            ),
112        ),
113        vec![
114            Node::let_bind("params_start", Expr::u32(INVALID_POS)),
115            Node::let_bind("params_end", Expr::u32(INVALID_POS)),
116            Node::let_bind("colon_pos", Expr::u32(INVALID_POS)),
117            // Hoist `after_type_params` and `after_params` to the
118            // outer scope so the if-block bodies (which assign them)
119            // and the later if-blocks (which read them) share one
120            // binding. Pre-T-V2 the per-branch `Node::let_bind` lived
121            // inside each block, the validator scoped the binding to
122            // the block, and the read sites failed with "reference to
123            // undeclared variable `after_type_params`" / `after_params`.
124            Node::let_bind("after_type_params", Expr::u32(INVALID_POS)),
125            Node::let_bind("after_params", Expr::u32(INVALID_POS)),
126            Node::if_then_else(
127                Expr::eq(
128                    load_u32(tok_types, Expr::var("post_name")),
129                    Expr::u32(TOK_LBRACKET),
130                ),
131                search_next_token_into(
132                    "after_type_params",
133                    Expr::add(Expr::var("type_params_end"), Expr::u32(1)),
134                    tok_types,
135                    haystack_len,
136                ),
137                vec![Node::assign("after_type_params", Expr::var("post_name"))],
138            ),
139            Node::if_then(
140                Expr::eq(
141                    load_u32(tok_types, Expr::var("after_type_params")),
142                    Expr::u32(TOK_LPAREN),
143                ),
144                vec![
145                    Node::assign("params_start", Expr::var("after_type_params")),
146                    Node::assign("params_end", Expr::u32(INVALID_POS)),
147                ]
148                .into_iter()
149                .chain(find_matching_delimiter_into(
150                    "params_end",
151                    Expr::var("after_type_params"),
152                    tok_types,
153                    haystack_len,
154                    TOK_LPAREN,
155                    crate::parsing::python::lex::TOK_RPAREN,
156                ))
157                .collect(),
158            ),
159            Node::if_then_else(
160                Expr::ne(Expr::var("params_end"), Expr::u32(INVALID_POS)),
161                search_next_token_into(
162                    "after_params",
163                    Expr::add(Expr::var("params_end"), Expr::u32(1)),
164                    tok_types,
165                    haystack_len,
166                ),
167                vec![Node::assign("after_params", Expr::var("after_type_params"))],
168            ),
169            Node::if_then(
170                Expr::eq(
171                    load_u32(tok_types, Expr::var("after_params")),
172                    Expr::u32(TOK_COLON),
173                ),
174                vec![Node::assign("colon_pos", Expr::var("after_params"))],
175            ),
176            Node::let_bind(
177                "slot",
178                Expr::atomic_add(out_counts, Expr::u32(0), Expr::u32(DEF_RECORD_WORDS)),
179            ),
180        ]
181        .into_iter()
182        .chain(store_words(
183            out_records,
184            "slot",
185            &[
186                Expr::var("emit_kind"),
187                load_u32(tok_starts, Expr::var("name_pos")),
188                load_u32(tok_lens, Expr::var("name_pos")),
189                Expr::var("params_start"),
190                Expr::var("params_end"),
191                Expr::var("colon_pos"),
192            ],
193        ))
194        .collect(),
195    ));
196
197    let pass = line_index_pass(STRUCTURE_OP_ID, tok_types, tok_starts, tok_lens, haystack_len);
198    let mut buffers = pass.token_buffers();
199    buffers.extend(pass.record_buffers(out_records, out_counts, 3, DEF_RECORD_WORDS));
200    pass.program(buffers, body)
201}
202
203/// Extract `import` and `from ... import ...` statements.
204#[must_use]
205pub fn python312_extract_imports(
206    tok_types: &str,
207    tok_starts: &str,
208    tok_lens: &str,
209    out_records: &str,
210    out_counts: &str,
211    haystack_len: u32,
212) -> Program {
213    let t = Expr::InvocationId { axis: 0 };
214    let name = DottedName {
215        tok_types,
216        haystack_len,
217        head: t.clone(),
218        accumulator: "name_end",
219    };
220    let mut body = vec![
221        Node::let_bind("tok", load_u32(tok_types, t.clone())),
222        Node::let_bind("record_kind", Expr::u32(0)),
223    ];
224    body.extend(search_prev_token("prev_tok", t.clone(), tok_types));
225    body.extend(search_next_token(
226        "next_tok",
227        Expr::add(t.clone(), Expr::u32(1)),
228        tok_types,
229        haystack_len,
230    ));
231    body.push(Node::if_then(
232        Expr::and(
233            Expr::eq(Expr::var("tok"), Expr::u32(TOK_IDENTIFIER)),
234            Expr::or(
235                Expr::eq(
236                    load_u32(tok_types, Expr::var("prev_tok")),
237                    Expr::u32(TOK_IMPORT),
238                ),
239                Expr::eq(
240                    load_u32(tok_types, Expr::var("prev_tok")),
241                    Expr::u32(TOK_FROM),
242                ),
243            ),
244        ),
245        vec![Node::assign(
246            "record_kind",
247            Expr::select(
248                Expr::eq(
249                    load_u32(tok_types, Expr::var("prev_tok")),
250                    Expr::u32(TOK_IMPORT),
251                ),
252                Expr::u32(1),
253                Expr::u32(2),
254            ),
255        )],
256    ));
257    body.push(Node::if_then(
258        Expr::and(
259            Expr::eq(Expr::var("tok"), Expr::u32(TOK_IDENTIFIER)),
260            Expr::eq(
261                load_u32(tok_types, Expr::var("prev_tok")),
262                Expr::u32(TOK_COMMA),
263            ),
264        ),
265        vec![Node::assign("record_kind", Expr::u32(1))],
266    ));
267    let span = name.span(tok_starts, tok_lens);
268    body.push(Node::if_then(
269        Expr::ne(Expr::var("record_kind"), Expr::u32(0)),
270        name.carriers()
271            .into_iter()
272            .chain([
273                name.walk(),
274                Node::let_bind(
275                    "slot",
276                    Expr::atomic_add(out_counts, Expr::u32(0), Expr::u32(IMPORT_RECORD_WORDS)),
277                ),
278            ])
279            .chain(store_words(
280                out_records,
281                "slot",
282                &[
283                    Expr::var("record_kind"),
284                    span[0].clone(),
285                    span[1].clone(),
286                    Expr::var("prev_tok"),
287                    Expr::var("name_end"),
288                    Expr::var("next_tok"),
289                ],
290            ))
291            .collect(),
292    ));
293
294    let pass = line_index_pass(IMPORTS_OP_ID, tok_types, tok_starts, tok_lens, haystack_len);
295    let mut buffers = pass.token_buffers();
296    buffers.extend(pass.record_buffers(out_records, out_counts, 3, IMPORT_RECORD_WORDS));
297    pass.program(buffers, body)
298}
299
300/// Extract `with` / `async with` headers.
301#[must_use]
302pub fn python312_extract_with_blocks(
303    tok_types: &str,
304    tok_starts: &str,
305    tok_lens: &str,
306    out_records: &str,
307    out_counts: &str,
308    haystack_len: u32,
309) -> Program {
310    let t = Expr::InvocationId { axis: 0 };
311    let name = DottedName {
312        tok_types,
313        haystack_len,
314        head: Expr::var("manager_pos"),
315        accumulator: "manager_end",
316    };
317    let mut body = vec![
318        Node::let_bind("tok", load_u32(tok_types, t.clone())),
319        Node::let_bind("with_pos", Expr::u32(INVALID_POS)),
320        Node::let_bind("flags", Expr::u32(0)),
321    ];
322    body.extend(search_prev_token("prev_tok", t.clone(), tok_types));
323    body.push(Node::if_then(
324        Expr::and(
325            Expr::eq(Expr::var("tok"), Expr::u32(TOK_WITH)),
326            Expr::ne(
327                load_u32(tok_types, Expr::var("prev_tok")),
328                Expr::u32(TOK_ASYNC),
329            ),
330        ),
331        vec![Node::assign("with_pos", t.clone())],
332    ));
333    body.extend(search_next_token(
334        "async_next",
335        Expr::add(t.clone(), Expr::u32(1)),
336        tok_types,
337        haystack_len,
338    ));
339    body.push(Node::if_then(
340        Expr::and(
341            Expr::eq(Expr::var("tok"), Expr::u32(TOK_ASYNC)),
342            Expr::eq(
343                load_u32(tok_types, Expr::var("async_next")),
344                Expr::u32(TOK_WITH),
345            ),
346        ),
347        vec![
348            Node::assign("with_pos", Expr::var("async_next")),
349            Node::assign("flags", Expr::u32(1)),
350        ],
351    ));
352    body.extend(search_next_token(
353        "manager_pos",
354        Expr::add(Expr::var("with_pos"), Expr::u32(1)),
355        tok_types,
356        haystack_len,
357    ));
358    body.extend(search_next_token(
359        "after_manager",
360        Expr::add(Expr::var("manager_pos"), Expr::u32(1)),
361        tok_types,
362        haystack_len,
363    ));
364    let span = name.span(tok_starts, tok_lens);
365    body.push(Node::if_then(
366        Expr::and(
367            Expr::ne(Expr::var("with_pos"), Expr::u32(INVALID_POS)),
368            Expr::eq(
369                load_u32(tok_types, Expr::var("manager_pos")),
370                Expr::u32(TOK_IDENTIFIER),
371            ),
372        ),
373        name.carriers()
374            .into_iter()
375            .chain([
376                name.walk(),
377                Node::let_bind("colon_pos", Expr::u32(INVALID_POS)),
378                Node::loop_for(
379                    "scan",
380                    Expr::add(Expr::var("manager_end"), Expr::u32(1)),
381                    Expr::u32(haystack_len),
382                    vec![Node::if_then(
383                        Expr::and(
384                            Expr::eq(Expr::var("colon_pos"), Expr::u32(INVALID_POS)),
385                            Expr::eq(
386                                load_u32(tok_types, Expr::var("scan")),
387                                Expr::u32(TOK_COLON),
388                            ),
389                        ),
390                        vec![Node::assign("colon_pos", Expr::var("scan"))],
391                    )],
392                ),
393                Node::let_bind(
394                    "slot",
395                    Expr::atomic_add(out_counts, Expr::u32(0), Expr::u32(WITH_RECORD_WORDS)),
396                ),
397            ])
398            .chain(store_words(
399                out_records,
400                "slot",
401                &[
402                    span[0].clone(),
403                    span[1].clone(),
404                    Expr::var("with_pos"),
405                    Expr::var("colon_pos"),
406                    Expr::var("flags"),
407                    Expr::u32(0),
408                ],
409            ))
410            .collect(),
411    ));
412
413    let pass = line_index_pass(
414        WITH_BLOCKS_OP_ID,
415        tok_types,
416        tok_starts,
417        tok_lens,
418        haystack_len,
419    );
420    let mut buffers = pass.token_buffers();
421    buffers.extend(pass.record_buffers(out_records, out_counts, 3, WITH_RECORD_WORDS));
422    pass.program(buffers, body)
423}
424
425inventory::submit! {
426    vyre_foundation::operation::OperationRegistration {
427        semantic_version: 1,
428        signature: None,
429        tier: vyre_foundation::operation::OperationTier::Library,
430        laws: &[],
431        tolerance: vyre_foundation::operation::TolerancePolicy::EXACT,
432        id: STRUCTURE_OP_ID,
433        build: Some(|| python312_extract_structure("tok_types", "tok_starts", "tok_lens", "out_records", "out_counts", 16)),
434        test_inputs: Some(structure_fixture_inputs),
435        expected_output: Some(structure_fixture_expected),
436        category: Some("parsing"),
437    }
438}
439
440inventory::submit! {
441    vyre_foundation::operation::OperationRegistration {
442        semantic_version: 1,
443        signature: None,
444        tier: vyre_foundation::operation::OperationTier::Library,
445        laws: &[],
446        tolerance: vyre_foundation::operation::TolerancePolicy::EXACT,
447        id: IMPORTS_OP_ID,
448        build: Some(|| python312_extract_imports("tok_types", "tok_starts", "tok_lens", "out_records", "out_counts", 16)),
449        test_inputs: Some(import_fixture_inputs),
450        expected_output: Some(import_fixture_expected),
451        category: Some("parsing"),
452    }
453}
454
455inventory::submit! {
456    vyre_foundation::operation::OperationRegistration {
457        semantic_version: 1,
458        signature: None,
459        tier: vyre_foundation::operation::OperationTier::Library,
460        laws: &[],
461        tolerance: vyre_foundation::operation::TolerancePolicy::EXACT,
462        id: WITH_BLOCKS_OP_ID,
463        build: Some(|| python312_extract_with_blocks("tok_types", "tok_starts", "tok_lens", "out_records", "out_counts", 16)),
464        test_inputs: Some(with_fixture_inputs),
465        expected_output: Some(with_fixture_expected),
466        category: Some("parsing"),
467    }
468}
469
470fn structure_fixture_inputs() -> Vec<Vec<Vec<u8>>> {
471    let (tok_types, tok_starts, tok_lens) = pack_sparse_tokens(
472        &[
473            (0, TOK_DEF, 3),
474            (4, TOK_IDENTIFIER, 1),
475            (5, TOK_LPAREN, 1),
476            (6, crate::parsing::python::lex::TOK_RPAREN, 1),
477            (7, TOK_COLON, 1),
478        ],
479        16,
480    );
481    vec![vec![
482        tok_types,
483        tok_starts,
484        tok_lens,
485        vec![0u8; 16 * DEF_RECORD_WORDS as usize * 4],
486        vec![0u8; 4],
487    ]]
488}
489
490fn structure_fixture_expected() -> Vec<Vec<Vec<u8>>> {
491    let mut records = vec![0u8; 16 * DEF_RECORD_WORDS as usize * 4];
492    write_words(&mut records, &[1, 4, 1, 5, 6, 7]);
493    vec![vec![records, DEF_RECORD_WORDS.to_le_bytes().to_vec()]]
494}
495
496fn import_fixture_inputs() -> Vec<Vec<Vec<u8>>> {
497    let (tok_types, tok_starts, tok_lens) =
498        pack_sparse_tokens(&[(0, TOK_IMPORT, 6), (7, TOK_IDENTIFIER, 2)], 16);
499    vec![vec![
500        tok_types,
501        tok_starts,
502        tok_lens,
503        vec![0u8; 16 * IMPORT_RECORD_WORDS as usize * 4],
504        vec![0u8; 4],
505    ]]
506}
507
508fn import_fixture_expected() -> Vec<Vec<Vec<u8>>> {
509    let mut records = vec![0u8; 16 * IMPORT_RECORD_WORDS as usize * 4];
510    write_words(
511        &mut records,
512        &[1, 7, 2, 0, 7, crate::parsing::python::INVALID_POS],
513    );
514    vec![vec![records, IMPORT_RECORD_WORDS.to_le_bytes().to_vec()]]
515}
516
517fn with_fixture_inputs() -> Vec<Vec<Vec<u8>>> {
518    let (tok_types, tok_starts, tok_lens) = pack_sparse_tokens(
519        &[
520            (0, TOK_ASYNC, 5),
521            (6, TOK_WITH, 4),
522            (11, TOK_IDENTIFIER, 3),
523            (14, TOK_COLON, 1),
524        ],
525        16,
526    );
527    vec![vec![
528        tok_types,
529        tok_starts,
530        tok_lens,
531        vec![0u8; 16 * WITH_RECORD_WORDS as usize * 4],
532        vec![0u8; 4],
533    ]]
534}
535
536fn with_fixture_expected() -> Vec<Vec<Vec<u8>>> {
537    let mut records = vec![0u8; 16 * WITH_RECORD_WORDS as usize * 4];
538    write_words(&mut records, &[11, 3, 6, 14, 1, 0]);
539    vec![vec![records, WITH_RECORD_WORDS.to_le_bytes().to_vec()]]
540}
vyre_libs/parsing/python/parse/structure.rs

vyre_libs/parsing/python/parse/
structure.rs
