Skip to main content

vyre_libs/parsing/python/parse/
decorators.rs

1use super::walk::{pack_sparse_tokens, DottedName, TokenPass};
2use super::{
3    find_matching_delimiter, load_u32, search_next_token, search_next_token_into, store_words,
4    write_words,
5};
6use crate::parsing::python::lex::{
7    TOK_ASYNC, TOK_AT, TOK_CLASS, TOK_DEF, TOK_IDENTIFIER, TOK_LPAREN, TOK_RPAREN,
8};
9use crate::parsing::python::{DECORATOR_RECORD_WORDS, INVALID_POS};
10use vyre_foundation::ir::{Expr, Node, Program};
11
12const OP_ID: &str = "vyre-libs::parsing::python312_extract_decorators";
13
14/// Extract decorator occurrences and their immediate target.
15#[must_use]
16pub fn python312_extract_decorators(
17    tok_types: &str,
18    tok_starts: &str,
19    tok_lens: &str,
20    out_records: &str,
21    out_counts: &str,
22    haystack_len: u32,
23) -> Program {
24    let t = Expr::InvocationId { axis: 0 };
25    let name = DottedName {
26        tok_types,
27        haystack_len,
28        head: Expr::var("decorator_name"),
29        accumulator: "decorator_end",
30    };
31    let mut body = Vec::new();
32    body.extend(search_next_token(
33        "decorator_name",
34        Expr::add(t.clone(), Expr::u32(1)),
35        tok_types,
36        haystack_len,
37    ));
38    body.push(Node::let_bind("tok", load_u32(tok_types, t.clone())));
39    // Hoist every cross-block name to the outer body so it survives
40    // the if_then / loop_for scopes that assign it. Each helper we
41    // call inside an if_then uses the `_into` (assign-only) variant
42    // so the outer let_bind isn't redeclared (V008/V032 noise).
43    body.extend(name.carriers());
44    body.push(Node::let_bind("after_decorator", Expr::u32(INVALID_POS)));
45    body.push(Node::let_bind("target_tok", Expr::u32(INVALID_POS)));
46    body.push(Node::let_bind("target_name", Expr::u32(INVALID_POS)));
47    body.push(Node::let_bind("target_kind", Expr::u32(0)));
48    body.push(Node::let_bind("async_def", Expr::u32(INVALID_POS)));
49    body.extend(find_matching_delimiter(
50        "decorator_rparen",
51        Expr::var("decorator_name"),
52        tok_types,
53        haystack_len,
54        TOK_LPAREN,
55        TOK_RPAREN,
56    ));
57    let span = name.span(tok_starts, tok_lens);
58    body.push(Node::if_then(
59        Expr::and(
60            Expr::eq(Expr::var("tok"), Expr::u32(TOK_AT)),
61            Expr::eq(
62                load_u32(tok_types, Expr::var("decorator_name")),
63                Expr::u32(TOK_IDENTIFIER),
64            ),
65        ),
66        vec![name.walk()]
67            .into_iter()
68            .chain(search_next_token_into(
69                "after_decorator",
70                Expr::add(Expr::var("decorator_end"), Expr::u32(1)),
71                tok_types,
72                haystack_len,
73            ))
74            .chain(vec![Node::if_then_else(
75                Expr::eq(
76                    load_u32(tok_types, Expr::var("after_decorator")),
77                    Expr::u32(TOK_LPAREN),
78                ),
79                search_next_token_into(
80                    "target_tok",
81                    Expr::add(Expr::var("decorator_rparen"), Expr::u32(1)),
82                    tok_types,
83                    haystack_len,
84                ),
85                search_next_token_into(
86                    "target_tok",
87                    Expr::add(Expr::var("decorator_end"), Expr::u32(1)),
88                    tok_types,
89                    haystack_len,
90                ),
91            )])
92            .chain(vec![
93                Node::if_then(
94                    Expr::eq(
95                        load_u32(tok_types, Expr::var("target_tok")),
96                        Expr::u32(TOK_DEF),
97                    ),
98                    vec![
99                        Node::assign("target_kind", Expr::u32(1)),
100                        Node::assign("target_name", Expr::u32(INVALID_POS)),
101                    ]
102                    .into_iter()
103                    .chain(search_next_token_into(
104                        "target_name",
105                        Expr::add(Expr::var("target_tok"), Expr::u32(1)),
106                        tok_types,
107                        haystack_len,
108                    ))
109                    .collect(),
110                ),
111                Node::if_then(
112                    Expr::eq(
113                        load_u32(tok_types, Expr::var("target_tok")),
114                        Expr::u32(TOK_CLASS),
115                    ),
116                    vec![
117                        Node::assign("target_kind", Expr::u32(3)),
118                        Node::assign("target_name", Expr::u32(INVALID_POS)),
119                    ]
120                    .into_iter()
121                    .chain(search_next_token_into(
122                        "target_name",
123                        Expr::add(Expr::var("target_tok"), Expr::u32(1)),
124                        tok_types,
125                        haystack_len,
126                    ))
127                    .collect(),
128                ),
129                Node::if_then(
130                    Expr::eq(
131                        load_u32(tok_types, Expr::var("target_tok")),
132                        Expr::u32(TOK_ASYNC),
133                    ),
134                    vec![
135                        Node::assign("target_kind", Expr::u32(2)),
136                        Node::assign("target_name", Expr::u32(INVALID_POS)),
137                    ]
138                    .into_iter()
139                    .chain(search_next_token_into(
140                        "async_def",
141                        Expr::add(Expr::var("target_tok"), Expr::u32(1)),
142                        tok_types,
143                        haystack_len,
144                    ))
145                    .chain(search_next_token_into(
146                        "target_name",
147                        Expr::add(Expr::var("async_def"), Expr::u32(1)),
148                        tok_types,
149                        haystack_len,
150                    ))
151                    .collect(),
152                ),
153                Node::let_bind(
154                    "slot",
155                    Expr::atomic_add(out_counts, Expr::u32(0), Expr::u32(DECORATOR_RECORD_WORDS)),
156                ),
157            ])
158            .chain(store_words(
159                out_records,
160                "slot",
161                &[
162                    span[0].clone(),
163                    span[1].clone(),
164                    Expr::var("target_kind"),
165                    load_u32(tok_starts, Expr::var("target_name")),
166                    load_u32(tok_lens, Expr::var("target_name")),
167                    Expr::var("target_tok"),
168                ],
169            ))
170            .collect(),
171    ));
172
173    let pass = TokenPass {
174        op_id: OP_ID,
175        child_op_id: vyre_primitives::parsing::core_delimiter_match::OP_ID,
176        tok_types,
177        tok_starts,
178        tok_lens,
179        haystack_len,
180    };
181    let mut buffers = pass.token_buffers();
182    buffers.extend(pass.record_buffers(out_records, out_counts, 3, DECORATOR_RECORD_WORDS));
183    pass.program(buffers, body)
184}
185
186inventory::submit! {
187    vyre_foundation::operation::OperationRegistration {
188        semantic_version: 1,
189        signature: None,
190        tier: vyre_foundation::operation::OperationTier::Library,
191        laws: &[],
192        tolerance: vyre_foundation::operation::TolerancePolicy::EXACT,
193        id: OP_ID,
194        build: Some(|| python312_extract_decorators("tok_types", "tok_starts", "tok_lens", "out_records", "out_counts", 16)),
195        test_inputs: Some(decorator_fixture_inputs),
196        expected_output: Some(decorator_fixture_expected),
197        category: Some("parsing"),
198    }
199}
200
201fn decorator_fixture_inputs() -> Vec<Vec<Vec<u8>>> {
202    let (tok_types, tok_starts, tok_lens) = pack_sparse_tokens(
203        &[
204            (0, TOK_AT, 1),
205            (1, TOK_IDENTIFIER, 1),
206            (3, TOK_ASYNC, 5),
207            (9, TOK_DEF, 3),
208            (13, TOK_IDENTIFIER, 1),
209        ],
210        16,
211    );
212
213    vec![vec![
214        tok_types,
215        tok_starts,
216        tok_lens,
217        vec![0u8; 16 * DECORATOR_RECORD_WORDS as usize * 4],
218        vec![0u8; 4],
219    ]]
220}
221
222fn decorator_fixture_expected() -> Vec<Vec<Vec<u8>>> {
223    let mut records = vec![0u8; 16 * DECORATOR_RECORD_WORDS as usize * 4];
224    write_words(&mut records, &[1, 1, 2, 13, 1, 3]);
225
226    vec![vec![records, DECORATOR_RECORD_WORDS.to_le_bytes().to_vec()]]
227}