vyre_libs/parsing/python/parse/
decorators.rs

1use super::{
2    find_matching_delimiter, load_u32, search_next_token, search_next_token_into, store_words,
3};
4use crate::parsing::composition::child_phase;
5use crate::parsing::python::lex::{
6    TOK_ASYNC, TOK_AT, TOK_CLASS, TOK_DEF, TOK_DOT, TOK_IDENTIFIER, TOK_LPAREN, TOK_RPAREN,
7};
8use crate::parsing::python::{DECORATOR_RECORD_WORDS, INVALID_POS, MAX_DOTTED_SEGMENTS};
9use crate::region::wrap_anonymous;
10use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
11
12/// Extract decorator occurrences and their immediate target.
13#[must_use]
14pub fn python312_extract_decorators(
15    tok_types: &str,
16    tok_starts: &str,
17    tok_lens: &str,
18    out_records: &str,
19    out_counts: &str,
20    haystack_len: u32,
21) -> Program {
22    let t = Expr::InvocationId { axis: 0 };
23    let mut body = Vec::new();
24    body.extend(search_next_token(
25        "decorator_name",
26        Expr::add(t.clone(), Expr::u32(1)),
27        tok_types,
28        haystack_len,
29    ));
30    body.push(Node::let_bind("tok", load_u32(tok_types, t.clone())));
31    // Hoist every cross-block name to the outer body so it survives
32    // the if_then / loop_for scopes that assign it. Each helper we
33    // call inside an if_then uses the `_into` (assign-only) variant
34    // so the outer let_bind isn't redeclared (V008/V032 noise).
35    body.push(Node::let_bind("decorator_end", Expr::var("decorator_name")));
36    body.push(Node::let_bind("cursor", Expr::var("decorator_name")));
37    body.push(Node::let_bind("dot_pos", Expr::u32(INVALID_POS)));
38    body.push(Node::let_bind("after_dot", Expr::u32(INVALID_POS)));
39    body.push(Node::let_bind("after_decorator", Expr::u32(INVALID_POS)));
40    body.push(Node::let_bind("target_tok", Expr::u32(INVALID_POS)));
41    body.push(Node::let_bind("target_name", Expr::u32(INVALID_POS)));
42    body.push(Node::let_bind("target_kind", Expr::u32(0)));
43    body.push(Node::let_bind("async_def", Expr::u32(INVALID_POS)));
44    body.extend(find_matching_delimiter(
45        "decorator_rparen",
46        Expr::var("decorator_name"),
47        tok_types,
48        haystack_len,
49        TOK_LPAREN,
50        TOK_RPAREN,
51    ));
52    body.push(Node::if_then(
53        Expr::and(
54            Expr::eq(Expr::var("tok"), Expr::u32(TOK_AT)),
55            Expr::eq(
56                load_u32(tok_types, Expr::var("decorator_name")),
57                Expr::u32(TOK_IDENTIFIER),
58            ),
59        ),
60        vec![Node::loop_for(
61            "seg",
62            Expr::u32(0),
63            Expr::u32(MAX_DOTTED_SEGMENTS),
64            vec![
65                Node::assign("dot_pos", Expr::u32(INVALID_POS)),
66                Node::assign("after_dot", Expr::u32(INVALID_POS)),
67            ]
68            .into_iter()
69            .chain(search_next_token_into(
70                "dot_pos",
71                Expr::add(Expr::var("cursor"), Expr::u32(1)),
72                tok_types,
73                haystack_len,
74            ))
75            .chain(vec![Node::if_then(
76                Expr::eq(
77                    load_u32(tok_types, Expr::var("dot_pos")),
78                    Expr::u32(TOK_DOT),
79                ),
80                search_next_token_into(
81                    "after_dot",
82                    Expr::add(Expr::var("dot_pos"), Expr::u32(1)),
83                    tok_types,
84                    haystack_len,
85                ),
86            )])
87            .chain(vec![
88                Node::if_then(
89                    Expr::eq(
90                        load_u32(tok_types, Expr::var("after_dot")),
91                        Expr::u32(TOK_IDENTIFIER),
92                    ),
93                    vec![
94                        Node::assign("decorator_end", Expr::var("after_dot")),
95                        Node::assign("cursor", Expr::var("after_dot")),
96                    ],
97                ),
98                Node::if_then(
99                    Expr::ne(
100                        load_u32(tok_types, Expr::var("after_dot")),
101                        Expr::u32(TOK_IDENTIFIER),
102                    ),
103                    vec![Node::assign("cursor", Expr::u32(INVALID_POS))],
104                ),
105            ])
106            .collect(),
107        )]
108        .into_iter()
109        .chain(search_next_token_into(
110            "after_decorator",
111            Expr::add(Expr::var("decorator_end"), Expr::u32(1)),
112            tok_types,
113            haystack_len,
114        ))
115        .chain(vec![Node::if_then_else(
116            Expr::eq(
117                load_u32(tok_types, Expr::var("after_decorator")),
118                Expr::u32(TOK_LPAREN),
119            ),
120            search_next_token_into(
121                "target_tok",
122                Expr::add(Expr::var("decorator_rparen"), Expr::u32(1)),
123                tok_types,
124                haystack_len,
125            ),
126            search_next_token_into(
127                "target_tok",
128                Expr::add(Expr::var("decorator_end"), Expr::u32(1)),
129                tok_types,
130                haystack_len,
131            ),
132        )])
133        .chain(vec![
134            Node::if_then(
135                Expr::eq(
136                    load_u32(tok_types, Expr::var("target_tok")),
137                    Expr::u32(TOK_DEF),
138                ),
139                vec![
140                    Node::assign("target_kind", Expr::u32(1)),
141                    Node::assign("target_name", Expr::u32(INVALID_POS)),
142                ]
143                .into_iter()
144                .chain(search_next_token_into(
145                    "target_name",
146                    Expr::add(Expr::var("target_tok"), Expr::u32(1)),
147                    tok_types,
148                    haystack_len,
149                ))
150                .collect(),
151            ),
152            Node::if_then(
153                Expr::eq(
154                    load_u32(tok_types, Expr::var("target_tok")),
155                    Expr::u32(TOK_CLASS),
156                ),
157                vec![
158                    Node::assign("target_kind", Expr::u32(3)),
159                    Node::assign("target_name", Expr::u32(INVALID_POS)),
160                ]
161                .into_iter()
162                .chain(search_next_token_into(
163                    "target_name",
164                    Expr::add(Expr::var("target_tok"), Expr::u32(1)),
165                    tok_types,
166                    haystack_len,
167                ))
168                .collect(),
169            ),
170            Node::if_then(
171                Expr::eq(
172                    load_u32(tok_types, Expr::var("target_tok")),
173                    Expr::u32(TOK_ASYNC),
174                ),
175                vec![
176                    Node::assign("target_kind", Expr::u32(2)),
177                    Node::assign("target_name", Expr::u32(INVALID_POS)),
178                ]
179                .into_iter()
180                .chain(search_next_token_into(
181                    "async_def",
182                    Expr::add(Expr::var("target_tok"), Expr::u32(1)),
183                    tok_types,
184                    haystack_len,
185                ))
186                .chain(search_next_token_into(
187                    "target_name",
188                    Expr::add(Expr::var("async_def"), Expr::u32(1)),
189                    tok_types,
190                    haystack_len,
191                ))
192                .collect(),
193            ),
194            Node::let_bind(
195                "slot",
196                Expr::atomic_add(out_counts, Expr::u32(0), Expr::u32(DECORATOR_RECORD_WORDS)),
197            ),
198        ])
199        .chain(store_words(
200            out_records,
201            "slot",
202            &[
203                load_u32(tok_starts, Expr::var("decorator_name")),
204                Expr::add(
205                    Expr::sub(
206                        load_u32(tok_starts, Expr::var("decorator_end")),
207                        load_u32(tok_starts, Expr::var("decorator_name")),
208                    ),
209                    load_u32(tok_lens, Expr::var("decorator_end")),
210                ),
211                Expr::var("target_kind"),
212                load_u32(tok_starts, Expr::var("target_name")),
213                load_u32(tok_lens, Expr::var("target_name")),
214                Expr::var("target_tok"),
215            ],
216        ))
217        .collect(),
218    ));
219
220    Program::wrapped(
221        vec![
222            BufferDecl::storage(tok_types, 0, BufferAccess::ReadOnly, DataType::U32)
223                .with_count(haystack_len),
224            BufferDecl::storage(tok_starts, 1, BufferAccess::ReadOnly, DataType::U32)
225                .with_count(haystack_len),
226            BufferDecl::storage(tok_lens, 2, BufferAccess::ReadOnly, DataType::U32)
227                .with_count(haystack_len),
228            BufferDecl::storage(out_records, 3, BufferAccess::ReadWrite, DataType::U32)
229                .with_count(haystack_len.saturating_mul(DECORATOR_RECORD_WORDS)),
230            BufferDecl::storage(out_counts, 4, BufferAccess::ReadWrite, DataType::U32)
231                .with_count(1),
232        ],
233        [256, 1, 1],
234        vec![wrap_anonymous(
235            "vyre-libs::parsing::python312_extract_decorators",
236            vec![child_phase(
237                "vyre-libs::parsing::python312_extract_decorators",
238                vyre_primitives::parsing::core_delimiter_match::OP_ID,
239                vec![Node::if_then(
240                    Expr::lt(t.clone(), Expr::u32(haystack_len)),
241                    body,
242                )],
243            )],
244        )],
245    )
246    .with_entry_op_id("vyre-libs::parsing::python312_extract_decorators")
247    .with_non_composable_with_self(true)
248}
249
250inventory::submit! {
251    crate::harness::OpEntry {
252        id: "vyre-libs::parsing::python312_extract_decorators",
253        build: || python312_extract_decorators("tok_types", "tok_starts", "tok_lens", "out_records", "out_counts", 16),
254        test_inputs: Some(decorator_fixture_inputs),
255        expected_output: Some(decorator_fixture_expected),
256        category: Some("parsing"),
257    }
258}
259
260fn decorator_fixture_inputs() -> Vec<Vec<Vec<u8>>> {
261    let mut tok_types = vec![0u8; 16 * 4];
262    let mut tok_starts = vec![0u8; 16 * 4];
263    let mut tok_lens = vec![0u8; 16 * 4];
264    for (pos, tok, len) in [
265        (0usize, TOK_AT, 1u32),
266        (1, TOK_IDENTIFIER, 1),
267        (3, TOK_ASYNC, 5),
268        (9, TOK_DEF, 3),
269        (13, TOK_IDENTIFIER, 1),
270    ] {
271        let base = pos * 4;
272        tok_types[base..base + 4].copy_from_slice(&tok.to_le_bytes());
273        tok_starts[base..base + 4].copy_from_slice(&(pos as u32).to_le_bytes());
274        tok_lens[base..base + 4].copy_from_slice(&len.to_le_bytes());
275    }
276
277    vec![vec![
278        tok_types,
279        tok_starts,
280        tok_lens,
281        vec![0u8; 16 * DECORATOR_RECORD_WORDS as usize * 4],
282        vec![0u8; 4],
283    ]]
284}
285
286fn decorator_fixture_expected() -> Vec<Vec<Vec<u8>>> {
287    let mut records = vec![0u8; 16 * DECORATOR_RECORD_WORDS as usize * 4];
288    for (idx, word) in [1u32, 1, 2, 13, 1, 3].into_iter().enumerate() {
289        let base = idx * 4;
290        records[base..base + 4].copy_from_slice(&word.to_le_bytes());
291    }
292
293    vec![vec![records, DECORATOR_RECORD_WORDS.to_le_bytes().to_vec()]]
294}
vyre_libs/parsing/python/parse/decorators.rs

vyre_libs/parsing/python/parse/
decorators.rs
