1use super::walk::{pack_sparse_tokens, DottedName, TokenPass};
2use super::{
3 find_matching_delimiter, load_u32, search_next_token, search_prev_token, store_words,
4 write_words,
5};
6use crate::parsing::python::lex::{
7 TOK_AWAIT, TOK_DOT, TOK_EQ, TOK_IDENTIFIER, TOK_LPAREN, TOK_NUMBER, TOK_RPAREN,
8};
9use crate::parsing::python::{CALL_RECORD_WORDS, INVALID_POS, KWARG_RECORD_WORDS};
10use vyre_foundation::ir::{Expr, Node, Program};
11
12const OP_ID: &str = "vyre-libs::parsing::python312_extract_calls";
13
14#[must_use]
16#[allow(clippy::too_many_arguments)]
17pub fn python312_extract_calls(
18 tok_types: &str,
19 tok_starts: &str,
20 tok_lens: &str,
21 out_calls: &str,
22 out_call_counts: &str,
23 out_kwargs: &str,
24 out_kw_counts: &str,
25 haystack_len: u32,
26) -> Program {
27 let t = Expr::InvocationId { axis: 0 };
28 let name = DottedName {
29 tok_types,
30 haystack_len,
31 head: t.clone(),
32 accumulator: "name_end",
33 };
34 let mut body = vec![
38 Node::let_bind("tok", load_u32(tok_types, t.clone())),
39 Node::let_bind("emit", Expr::u32(0)),
40 Node::let_bind("is_call_head", Expr::u32(0)),
41 ];
42 body.extend(name.carriers());
43 body.extend(search_prev_token("prev_tok", t.clone(), tok_types));
44 body.push(Node::if_then(
45 Expr::and(
46 Expr::eq(Expr::var("tok"), Expr::u32(TOK_IDENTIFIER)),
47 Expr::ne(
48 load_u32(tok_types, Expr::var("prev_tok")),
49 Expr::u32(TOK_DOT),
50 ),
51 ),
52 vec![Node::assign("is_call_head", Expr::u32(1)), name.walk()],
53 ));
54 body.extend(search_next_token(
55 "after_name",
56 Expr::add(Expr::var("name_end"), Expr::u32(1)),
57 tok_types,
58 haystack_len,
59 ));
60 body.extend(find_matching_delimiter(
61 "rparen",
62 Expr::var("after_name"),
63 tok_types,
64 haystack_len,
65 TOK_LPAREN,
66 TOK_RPAREN,
67 ));
68 body.push(Node::if_then(
69 Expr::and(
70 Expr::eq(Expr::var("is_call_head"), Expr::u32(1)),
71 Expr::and(
72 Expr::eq(
73 load_u32(tok_types, Expr::var("after_name")),
74 Expr::u32(TOK_LPAREN),
75 ),
76 Expr::ne(Expr::var("rparen"), Expr::u32(INVALID_POS)),
77 ),
78 ),
79 vec![Node::assign("emit", Expr::u32(1))],
80 ));
81 let span = name.span(tok_starts, tok_lens);
82 body.push(Node::if_then(
83 Expr::eq(Expr::var("emit"), Expr::u32(1)),
84 vec![
85 Node::let_bind("kw_base", Expr::load(out_kw_counts, Expr::u32(0))),
86 Node::let_bind("kw_count", Expr::u32(0)),
87 Node::let_bind("paren_depth", Expr::u32(0)),
88 Node::let_bind("bracket_depth", Expr::u32(0)),
89 Node::loop_for(
90 "scan",
91 Expr::add(Expr::var("after_name"), Expr::u32(1)),
92 Expr::var("rparen"),
93 vec![
94 Node::let_bind("scan_tok", load_u32(tok_types, Expr::var("scan"))),
95 Node::if_then(
96 Expr::eq(Expr::var("scan_tok"), Expr::u32(TOK_LPAREN)),
97 vec![Node::assign(
98 "paren_depth",
99 Expr::add(Expr::var("paren_depth"), Expr::u32(1)),
100 )],
101 ),
102 Node::if_then(
103 Expr::eq(Expr::var("scan_tok"), Expr::u32(TOK_RPAREN)),
104 vec![Node::if_then(
105 Expr::gt(Expr::var("paren_depth"), Expr::u32(0)),
106 vec![Node::assign(
107 "paren_depth",
108 Expr::sub(Expr::var("paren_depth"), Expr::u32(1)),
109 )],
110 )],
111 ),
112 Node::if_then(
113 Expr::eq(
114 Expr::var("scan_tok"),
115 Expr::u32(crate::parsing::python::lex::TOK_LBRACKET),
116 ),
117 vec![Node::assign(
118 "bracket_depth",
119 Expr::add(Expr::var("bracket_depth"), Expr::u32(1)),
120 )],
121 ),
122 Node::if_then(
123 Expr::eq(
124 Expr::var("scan_tok"),
125 Expr::u32(crate::parsing::python::lex::TOK_RBRACKET),
126 ),
127 vec![Node::if_then(
128 Expr::gt(Expr::var("bracket_depth"), Expr::u32(0)),
129 vec![Node::assign(
130 "bracket_depth",
131 Expr::sub(Expr::var("bracket_depth"), Expr::u32(1)),
132 )],
133 )],
134 ),
135 Node::if_then(
136 Expr::and(
137 Expr::and(
138 Expr::eq(Expr::var("scan_tok"), Expr::u32(TOK_IDENTIFIER)),
139 Expr::eq(Expr::var("paren_depth"), Expr::u32(0)),
140 ),
141 Expr::eq(Expr::var("bracket_depth"), Expr::u32(0)),
142 ),
143 search_next_token(
148 "kw_eq_pos",
149 Expr::add(Expr::var("scan"), Expr::u32(1)),
150 tok_types,
151 haystack_len,
152 )
153 .into_iter()
154 .chain(search_prev_token("kw_prev", Expr::var("scan"), tok_types))
155 .chain(vec![Node::if_then(
156 Expr::and(
157 Expr::eq(
158 load_u32(tok_types, Expr::var("kw_eq_pos")),
159 Expr::u32(TOK_EQ),
160 ),
161 Expr::ne(
162 load_u32(tok_types, Expr::var("kw_prev")),
163 Expr::u32(TOK_DOT),
164 ),
165 ),
166 vec![
167 Node::let_bind(
168 "kw_slot",
169 Expr::atomic_add(
170 out_kw_counts,
171 Expr::u32(0),
172 Expr::u32(KWARG_RECORD_WORDS),
173 ),
174 ),
175 Node::store(
176 out_kwargs,
177 Expr::var("kw_slot"),
178 load_u32(tok_starts, Expr::var("scan")),
179 ),
180 Node::store(
181 out_kwargs,
182 Expr::add(Expr::var("kw_slot"), Expr::u32(1)),
183 load_u32(tok_lens, Expr::var("scan")),
184 ),
185 Node::assign(
186 "kw_count",
187 Expr::add(Expr::var("kw_count"), Expr::u32(1)),
188 ),
189 ],
190 )])
191 .collect(),
192 ),
193 ],
194 ),
195 Node::let_bind(
196 "call_slot",
197 Expr::atomic_add(out_call_counts, Expr::u32(0), Expr::u32(CALL_RECORD_WORDS)),
198 ),
199 ]
200 .into_iter()
201 .chain(store_words(
202 out_calls,
203 "call_slot",
204 &[
205 span[0].clone(),
206 span[1].clone(),
207 Expr::var("after_name"),
208 Expr::var("rparen"),
209 Expr::var("kw_base"),
210 Expr::var("kw_count"),
211 Expr::select(
212 Expr::eq(
213 load_u32(tok_types, Expr::var("prev_tok")),
214 Expr::u32(TOK_AWAIT),
215 ),
216 Expr::u32(1),
217 Expr::u32(0),
218 ),
219 ],
220 ))
221 .collect(),
222 ));
223
224 let pass = TokenPass {
225 op_id: OP_ID,
226 child_op_id: vyre_primitives::parsing::core_delimiter_match::OP_ID,
227 tok_types,
228 tok_starts,
229 tok_lens,
230 haystack_len,
231 };
232 let mut buffers = pass.token_buffers();
233 buffers.extend(pass.record_buffers(out_calls, out_call_counts, 3, CALL_RECORD_WORDS));
234 buffers.extend(pass.record_buffers(out_kwargs, out_kw_counts, 5, KWARG_RECORD_WORDS));
235 pass.program(buffers, body)
236}
237
238inventory::submit! {
239 vyre_foundation::operation::OperationRegistration {
240 semantic_version: 1,
241 signature: None,
242 tier: vyre_foundation::operation::OperationTier::Library,
243 laws: &[],
244 tolerance: vyre_foundation::operation::TolerancePolicy::EXACT,
245 id: OP_ID,
246 build: Some(|| python312_extract_calls(
247 "tok_types", "tok_starts", "tok_lens", "out_calls", "out_call_counts", "out_kwargs", "out_kw_counts", 16
248 )),
249 test_inputs: Some(call_fixture_inputs),
250 expected_output: Some(call_fixture_expected),
251 category: Some("parsing"),
252 }
253}
254
255fn call_fixture_inputs() -> Vec<Vec<Vec<u8>>> {
256 let (tok_types, tok_starts, tok_lens) = pack_sparse_tokens(
257 &[
258 (0, TOK_AWAIT, 5),
259 (6, TOK_IDENTIFIER, 3),
260 (9, TOK_LPAREN, 1),
261 (10, TOK_IDENTIFIER, 1),
262 (11, TOK_EQ, 1),
263 (12, TOK_NUMBER, 1),
264 (13, TOK_RPAREN, 1),
265 ],
266 16,
267 );
268
269 vec![vec![
270 tok_types,
271 tok_starts,
272 tok_lens,
273 vec![0u8; 16 * CALL_RECORD_WORDS as usize * 4],
274 vec![0u8; 4],
275 vec![0u8; 16 * KWARG_RECORD_WORDS as usize * 4],
276 vec![0u8; 4],
277 ]]
278}
279
280fn call_fixture_expected() -> Vec<Vec<Vec<u8>>> {
281 let mut calls = vec![0u8; 16 * CALL_RECORD_WORDS as usize * 4];
282 write_words(&mut calls, &[6, 3, 9, 13, 0, 1, 1]);
283
284 let mut kwargs = vec![0u8; 16 * KWARG_RECORD_WORDS as usize * 4];
285 write_words(&mut kwargs, &[10, 1]);
286
287 vec![vec![
288 calls,
289 CALL_RECORD_WORDS.to_le_bytes().to_vec(),
290 kwargs,
291 KWARG_RECORD_WORDS.to_le_bytes().to_vec(),
292 ]]
293}