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#[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 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}