1use super::{
2 find_matching_delimiter, load_u32, search_next_token, search_next_token_into,
3 search_prev_token, store_words, write_words,
4};
5use crate::parsing::composition::child_phase;
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::{
10 CALL_RECORD_WORDS, INVALID_POS, KWARG_RECORD_WORDS, MAX_DOTTED_SEGMENTS,
11};
12use crate::region::wrap_anonymous;
13use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
14
15#[must_use]
17#[allow(clippy::too_many_arguments)]
18pub fn python312_extract_calls(
19 tok_types: &str,
20 tok_starts: &str,
21 tok_lens: &str,
22 out_calls: &str,
23 out_call_counts: &str,
24 out_kwargs: &str,
25 out_kw_counts: &str,
26 haystack_len: u32,
27) -> Program {
28 let t = Expr::InvocationId { axis: 0 };
29 let mut body = vec![
30 Node::let_bind("tok", load_u32(tok_types, t.clone())),
31 Node::let_bind("emit", Expr::u32(0)),
32 Node::let_bind("is_call_head", Expr::u32(0)),
37 Node::let_bind("name_end", t.clone()),
38 Node::let_bind("cursor", t.clone()),
39 Node::let_bind("dot_pos", Expr::u32(INVALID_POS)),
40 Node::let_bind("after_dot", Expr::u32(INVALID_POS)),
41 ];
42 body.extend(search_prev_token("prev_tok", t.clone(), tok_types));
43 body.push(Node::if_then(
44 Expr::and(
45 Expr::eq(Expr::var("tok"), Expr::u32(TOK_IDENTIFIER)),
46 Expr::ne(
47 load_u32(tok_types, Expr::var("prev_tok")),
48 Expr::u32(TOK_DOT),
49 ),
50 ),
51 vec![
52 Node::assign("is_call_head", Expr::u32(1)),
53 Node::loop_for(
54 "seg",
55 Expr::u32(0),
56 Expr::u32(MAX_DOTTED_SEGMENTS),
57 vec![
58 Node::assign("dot_pos", Expr::u32(INVALID_POS)),
59 Node::assign("after_dot", Expr::u32(INVALID_POS)),
60 Node::if_then(
61 Expr::ne(Expr::var("cursor"), Expr::u32(INVALID_POS)),
62 search_next_token_into(
63 "dot_pos",
64 Expr::add(Expr::var("cursor"), Expr::u32(1)),
65 tok_types,
66 haystack_len,
67 ),
68 ),
69 Node::if_then(
70 Expr::eq(
71 load_u32(tok_types, Expr::var("dot_pos")),
72 Expr::u32(TOK_DOT),
73 ),
74 search_next_token_into(
75 "after_dot",
76 Expr::add(Expr::var("dot_pos"), Expr::u32(1)),
77 tok_types,
78 haystack_len,
79 ),
80 ),
81 Node::if_then(
82 Expr::eq(
83 load_u32(tok_types, Expr::var("after_dot")),
84 Expr::u32(TOK_IDENTIFIER),
85 ),
86 vec![
87 Node::assign("name_end", Expr::var("after_dot")),
88 Node::assign("cursor", Expr::var("after_dot")),
89 ],
90 ),
91 Node::if_then(
92 Expr::ne(
93 load_u32(tok_types, Expr::var("after_dot")),
94 Expr::u32(TOK_IDENTIFIER),
95 ),
96 vec![Node::assign("cursor", Expr::u32(INVALID_POS))],
97 ),
98 ],
99 ),
100 ],
101 ));
102 body.extend(search_next_token(
103 "after_name",
104 Expr::add(Expr::var("name_end"), Expr::u32(1)),
105 tok_types,
106 haystack_len,
107 ));
108 body.extend(find_matching_delimiter(
109 "rparen",
110 Expr::var("after_name"),
111 tok_types,
112 haystack_len,
113 TOK_LPAREN,
114 TOK_RPAREN,
115 ));
116 body.push(Node::if_then(
117 Expr::and(
118 Expr::eq(Expr::var("is_call_head"), Expr::u32(1)),
119 Expr::and(
120 Expr::eq(
121 load_u32(tok_types, Expr::var("after_name")),
122 Expr::u32(TOK_LPAREN),
123 ),
124 Expr::ne(Expr::var("rparen"), Expr::u32(INVALID_POS)),
125 ),
126 ),
127 vec![Node::assign("emit", Expr::u32(1))],
128 ));
129 body.push(Node::if_then(
130 Expr::eq(Expr::var("emit"), Expr::u32(1)),
131 vec![
132 Node::let_bind("kw_base", Expr::load(out_kw_counts, Expr::u32(0))),
133 Node::let_bind("kw_count", Expr::u32(0)),
134 Node::let_bind("paren_depth", Expr::u32(0)),
135 Node::let_bind("bracket_depth", Expr::u32(0)),
136 Node::loop_for(
137 "scan",
138 Expr::add(Expr::var("after_name"), Expr::u32(1)),
139 Expr::var("rparen"),
140 vec![
141 Node::let_bind("scan_tok", load_u32(tok_types, Expr::var("scan"))),
142 Node::if_then(
143 Expr::eq(Expr::var("scan_tok"), Expr::u32(TOK_LPAREN)),
144 vec![Node::assign(
145 "paren_depth",
146 Expr::add(Expr::var("paren_depth"), Expr::u32(1)),
147 )],
148 ),
149 Node::if_then(
150 Expr::eq(Expr::var("scan_tok"), Expr::u32(TOK_RPAREN)),
151 vec![Node::if_then(
152 Expr::gt(Expr::var("paren_depth"), Expr::u32(0)),
153 vec![Node::assign(
154 "paren_depth",
155 Expr::sub(Expr::var("paren_depth"), Expr::u32(1)),
156 )],
157 )],
158 ),
159 Node::if_then(
160 Expr::eq(
161 Expr::var("scan_tok"),
162 Expr::u32(crate::parsing::python::lex::TOK_LBRACKET),
163 ),
164 vec![Node::assign(
165 "bracket_depth",
166 Expr::add(Expr::var("bracket_depth"), Expr::u32(1)),
167 )],
168 ),
169 Node::if_then(
170 Expr::eq(
171 Expr::var("scan_tok"),
172 Expr::u32(crate::parsing::python::lex::TOK_RBRACKET),
173 ),
174 vec![Node::if_then(
175 Expr::gt(Expr::var("bracket_depth"), Expr::u32(0)),
176 vec![Node::assign(
177 "bracket_depth",
178 Expr::sub(Expr::var("bracket_depth"), Expr::u32(1)),
179 )],
180 )],
181 ),
182 Node::if_then(
183 Expr::and(
184 Expr::and(
185 Expr::eq(Expr::var("scan_tok"), Expr::u32(TOK_IDENTIFIER)),
186 Expr::eq(Expr::var("paren_depth"), Expr::u32(0)),
187 ),
188 Expr::eq(Expr::var("bracket_depth"), Expr::u32(0)),
189 ),
190 search_next_token(
195 "kw_eq_pos",
196 Expr::add(Expr::var("scan"), Expr::u32(1)),
197 tok_types,
198 haystack_len,
199 )
200 .into_iter()
201 .chain(search_prev_token("kw_prev", Expr::var("scan"), tok_types))
202 .chain(vec![Node::if_then(
203 Expr::and(
204 Expr::eq(
205 load_u32(tok_types, Expr::var("kw_eq_pos")),
206 Expr::u32(TOK_EQ),
207 ),
208 Expr::ne(
209 load_u32(tok_types, Expr::var("kw_prev")),
210 Expr::u32(TOK_DOT),
211 ),
212 ),
213 vec![
214 Node::let_bind(
215 "kw_slot",
216 Expr::atomic_add(
217 out_kw_counts,
218 Expr::u32(0),
219 Expr::u32(KWARG_RECORD_WORDS),
220 ),
221 ),
222 Node::store(
223 out_kwargs,
224 Expr::var("kw_slot"),
225 load_u32(tok_starts, Expr::var("scan")),
226 ),
227 Node::store(
228 out_kwargs,
229 Expr::add(Expr::var("kw_slot"), Expr::u32(1)),
230 load_u32(tok_lens, Expr::var("scan")),
231 ),
232 Node::assign(
233 "kw_count",
234 Expr::add(Expr::var("kw_count"), Expr::u32(1)),
235 ),
236 ],
237 )])
238 .collect(),
239 ),
240 ],
241 ),
242 Node::let_bind(
243 "call_slot",
244 Expr::atomic_add(out_call_counts, Expr::u32(0), Expr::u32(CALL_RECORD_WORDS)),
245 ),
246 ]
247 .into_iter()
248 .chain(store_words(
249 out_calls,
250 "call_slot",
251 &[
252 load_u32(tok_starts, t.clone()),
253 Expr::add(
254 Expr::sub(
255 load_u32(tok_starts, Expr::var("name_end")),
256 load_u32(tok_starts, t.clone()),
257 ),
258 load_u32(tok_lens, Expr::var("name_end")),
259 ),
260 Expr::var("after_name"),
261 Expr::var("rparen"),
262 Expr::var("kw_base"),
263 Expr::var("kw_count"),
264 Expr::select(
265 Expr::eq(
266 load_u32(tok_types, Expr::var("prev_tok")),
267 Expr::u32(TOK_AWAIT),
268 ),
269 Expr::u32(1),
270 Expr::u32(0),
271 ),
272 ],
273 ))
274 .collect(),
275 ));
276
277 Program::wrapped(
278 vec![
279 BufferDecl::storage(tok_types, 0, BufferAccess::ReadOnly, DataType::U32)
280 .with_count(haystack_len),
281 BufferDecl::storage(tok_starts, 1, BufferAccess::ReadOnly, DataType::U32)
282 .with_count(haystack_len),
283 BufferDecl::storage(tok_lens, 2, BufferAccess::ReadOnly, DataType::U32)
284 .with_count(haystack_len),
285 BufferDecl::storage(out_calls, 3, BufferAccess::ReadWrite, DataType::U32)
286 .with_count(haystack_len.saturating_mul(CALL_RECORD_WORDS)),
287 BufferDecl::storage(out_call_counts, 4, BufferAccess::ReadWrite, DataType::U32)
288 .with_count(1),
289 BufferDecl::storage(out_kwargs, 5, BufferAccess::ReadWrite, DataType::U32)
290 .with_count(haystack_len.saturating_mul(KWARG_RECORD_WORDS)),
291 BufferDecl::storage(out_kw_counts, 6, BufferAccess::ReadWrite, DataType::U32)
292 .with_count(1),
293 ],
294 [256, 1, 1],
295 vec![wrap_anonymous(
296 "vyre-libs::parsing::python312_extract_calls",
297 vec![child_phase(
298 "vyre-libs::parsing::python312_extract_calls",
299 vyre_primitives::parsing::core_delimiter_match::OP_ID,
300 vec![Node::if_then(
301 Expr::lt(t.clone(), Expr::u32(haystack_len)),
302 body,
303 )],
304 )],
305 )],
306 )
307 .with_entry_op_id("vyre-libs::parsing::python312_extract_calls")
308 .with_non_composable_with_self(true)
309}
310
311inventory::submit! {
312 crate::harness::OpEntry {
313 id: "vyre-libs::parsing::python312_extract_calls",
314 build: || python312_extract_calls(
315 "tok_types", "tok_starts", "tok_lens", "out_calls", "out_call_counts", "out_kwargs", "out_kw_counts", 16
316 ),
317 test_inputs: Some(call_fixture_inputs),
318 expected_output: Some(call_fixture_expected),
319 category: Some("parsing"),
320 }
321}
322
323fn call_fixture_inputs() -> Vec<Vec<Vec<u8>>> {
324 let mut tok_types = vec![0u8; 16 * 4];
325 let mut tok_starts = vec![0u8; 16 * 4];
326 let mut tok_lens = vec![0u8; 16 * 4];
327 for (pos, tok, len) in [
328 (0usize, TOK_AWAIT, 5u32),
329 (6, TOK_IDENTIFIER, 3),
330 (9, TOK_LPAREN, 1),
331 (10, TOK_IDENTIFIER, 1),
332 (11, TOK_EQ, 1),
333 (12, TOK_NUMBER, 1),
334 (13, TOK_RPAREN, 1),
335 ] {
336 let base = pos * 4;
337 tok_types[base..base + 4].copy_from_slice(&tok.to_le_bytes());
338 tok_starts[base..base + 4].copy_from_slice(&(pos as u32).to_le_bytes());
339 tok_lens[base..base + 4].copy_from_slice(&len.to_le_bytes());
340 }
341
342 vec![vec![
343 tok_types,
344 tok_starts,
345 tok_lens,
346 vec![0u8; 16 * CALL_RECORD_WORDS as usize * 4],
347 vec![0u8; 4],
348 vec![0u8; 16 * KWARG_RECORD_WORDS as usize * 4],
349 vec![0u8; 4],
350 ]]
351}
352
353fn call_fixture_expected() -> Vec<Vec<Vec<u8>>> {
354 let mut calls = vec![0u8; 16 * CALL_RECORD_WORDS as usize * 4];
355 write_words(&mut calls, &[6, 3, 9, 13, 0, 1, 1]);
356
357 let mut kwargs = vec![0u8; 16 * KWARG_RECORD_WORDS as usize * 4];
358 write_words(&mut kwargs, &[10, 1]);
359
360 vec![vec![
361 calls,
362 CALL_RECORD_WORDS.to_le_bytes().to_vec(),
363 kwargs,
364 KWARG_RECORD_WORDS.to_le_bytes().to_vec(),
365 ]]
366}