Skip to main content

vyre_libs/parsing/python/parse/
calls.rs

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/// Extract Python call sites plus top-level keyword arguments.
15#[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    // The dotted-name carriers are hoisted to the outer body so they outlive
35    // the if_then block that assigns them and remain in scope for the
36    // post-if_then `search_next_token` call that reads `name_end`.
37    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                        // Drop the explicit `Node::let_bind` siblings  -
144                        // `search_next_token` / `search_prev_token` each
145                        // emit their own outer let_bind, so the manual
146                        // ones here were duplicate-sibling V032 errors.
147                        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}