Skip to main content

vyre_libs/parsing/python/parse/
mod.rs

1//! Python structural extractors.
2
3/// Python call-site extractor.
4pub mod calls;
5/// Python decorator extractor.
6pub mod decorators;
7/// Python declaration/span extractor.
8pub mod structure;
9
10use crate::parsing::python::INVALID_POS;
11use vyre::ir::{Expr, Node};
12
13pub(crate) fn store_words(buffer: &str, base_var: &str, words: &[Expr]) -> Vec<Node> {
14    words
15        .iter()
16        .enumerate()
17        .map(|(idx, value)| {
18            Node::store(
19                buffer,
20                Expr::add(Expr::var(base_var), Expr::u32(idx as u32)),
21                value.clone(),
22            )
23        })
24        .collect()
25}
26
27pub(crate) fn write_words(dst: &mut [u8], words: &[u32]) {
28    for (idx, word) in words.iter().enumerate() {
29        let base = idx * 4;
30        dst[base..base + 4].copy_from_slice(&word.to_le_bytes());
31    }
32}
33
34pub(crate) fn load_u32(buffer: &str, index: Expr) -> Expr {
35    Expr::load(buffer, index)
36}
37
38pub(crate) fn search_next_token(
39    out_var: &str,
40    start_expr: Expr,
41    tok_types: &str,
42    haystack_len: u32,
43) -> Vec<Node> {
44    let scan = format!("{out_var}_scan");
45    vec![
46        Node::let_bind(out_var, Expr::u32(INVALID_POS)),
47        Node::loop_for(
48            scan.clone(),
49            start_expr,
50            Expr::u32(haystack_len),
51            vec![Node::if_then(
52                Expr::and(
53                    Expr::eq(Expr::var(out_var), Expr::u32(INVALID_POS)),
54                    Expr::ne(load_u32(tok_types, Expr::var(scan.clone())), Expr::u32(0)),
55                ),
56                vec![Node::assign(out_var, Expr::var(scan))],
57            )],
58        ),
59    ]
60}
61
62pub(crate) fn search_prev_token(out_var: &str, start_expr: Expr, tok_types: &str) -> Vec<Node> {
63    let rev = format!("{out_var}_rev");
64    let cand = format!("{out_var}_cand");
65    vec![
66        Node::let_bind(out_var, Expr::u32(INVALID_POS)),
67        Node::loop_for(
68            rev.clone(),
69            Expr::u32(0),
70            start_expr.clone(),
71            vec![
72                Node::let_bind(
73                    cand.clone(),
74                    Expr::sub(Expr::sub(start_expr.clone(), Expr::u32(1)), Expr::var(rev)),
75                ),
76                Node::if_then(
77                    Expr::and(
78                        Expr::eq(Expr::var(out_var), Expr::u32(INVALID_POS)),
79                        Expr::ne(load_u32(tok_types, Expr::var(cand.clone())), Expr::u32(0)),
80                    ),
81                    vec![Node::assign(out_var, Expr::var(cand))],
82                ),
83            ],
84        ),
85    ]
86}
87
88/// Same as [`search_next_token`] but skips the leading `Node::let_bind`
89/// for `out_var`. The caller must declare `out_var` (typically with
90/// `Expr::u32(INVALID_POS)`) in an enclosing scope so the binding
91/// outlives the if/loop block this output is consumed inside.
92pub(crate) fn search_next_token_into(
93    out_var: &str,
94    start_expr: Expr,
95    tok_types: &str,
96    haystack_len: u32,
97) -> Vec<Node> {
98    let scan = format!("{out_var}_scan");
99    vec![Node::loop_for(
100        scan.clone(),
101        start_expr,
102        Expr::u32(haystack_len),
103        vec![Node::if_then(
104            Expr::and(
105                Expr::eq(Expr::var(out_var), Expr::u32(INVALID_POS)),
106                Expr::ne(load_u32(tok_types, Expr::var(scan.clone())), Expr::u32(0)),
107            ),
108            vec![Node::assign(out_var, Expr::var(scan))],
109        )],
110    )]
111}
112
113pub(crate) fn find_matching_delimiter(
114    out_var: &str,
115    open_pos: Expr,
116    tok_types: &str,
117    haystack_len: u32,
118    open_tok: u32,
119    close_tok: u32,
120) -> Vec<Node> {
121    find_matching_delimiter_nodes(
122        out_var,
123        open_pos,
124        tok_types,
125        haystack_len,
126        open_tok,
127        close_tok,
128        true,
129    )
130}
131
132/// Same as [`find_matching_delimiter`] but skips the leading
133/// `Node::let_bind` for `out_var`; caller pre-declares it in the
134/// enclosing scope so the binding outlives the if/loop block this
135/// output is consumed inside.
136pub(crate) fn find_matching_delimiter_into(
137    out_var: &str,
138    open_pos: Expr,
139    tok_types: &str,
140    haystack_len: u32,
141    open_tok: u32,
142    close_tok: u32,
143) -> Vec<Node> {
144    find_matching_delimiter_nodes(
145        out_var,
146        open_pos,
147        tok_types,
148        haystack_len,
149        open_tok,
150        close_tok,
151        false,
152    )
153}
154
155fn find_matching_delimiter_nodes(
156    out_var: &str,
157    open_pos: Expr,
158    tok_types: &str,
159    haystack_len: u32,
160    open_tok: u32,
161    close_tok: u32,
162    declare_out: bool,
163) -> Vec<Node> {
164    let depth = format!("{out_var}_depth");
165    let scan = format!("{out_var}_scan");
166    let tok = format!("{out_var}_tok");
167    let mut nodes = Vec::with_capacity(3);
168    if declare_out {
169        nodes.push(Node::let_bind(out_var, Expr::u32(INVALID_POS)));
170    }
171    nodes.push(Node::let_bind(depth.clone(), Expr::u32(0)));
172    nodes.push(Node::loop_for(
173        scan.clone(),
174        Expr::add(open_pos, Expr::u32(1)),
175        Expr::u32(haystack_len),
176        vec![
177            Node::let_bind(tok.clone(), load_u32(tok_types, Expr::var(scan.clone()))),
178            Node::if_then(
179                Expr::eq(Expr::var(out_var), Expr::u32(INVALID_POS)),
180                vec![
181                    Node::if_then(
182                        Expr::eq(Expr::var(tok.clone()), Expr::u32(open_tok)),
183                        vec![Node::assign(
184                            depth.clone(),
185                            Expr::add(Expr::var(depth.clone()), Expr::u32(1)),
186                        )],
187                    ),
188                    Node::if_then(
189                        Expr::eq(Expr::var(tok), Expr::u32(close_tok)),
190                        vec![Node::if_then_else(
191                            Expr::eq(Expr::var(depth.clone()), Expr::u32(0)),
192                            vec![Node::assign(out_var, Expr::var(scan))],
193                            vec![Node::assign(
194                                depth.clone(),
195                                Expr::sub(Expr::var(depth), Expr::u32(1)),
196                            )],
197                        )],
198                    ),
199                ],
200            ),
201        ],
202    ));
203    nodes
204}