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