1
use tree_sitter::{Node, Parser, Tree};
2
use tree_sitter_essence::LANGUAGE;
3

            
4
use super::traversal::WalkDFS;
5

            
6
/// Parse the given source code into a syntax tree using tree-sitter.
7
///
8
/// If successful, returns a tuple containing the syntax tree and the raw source code.
9
/// If the source code is not valid Essence, returns None.
10
///
11
/// NOTE: The new source code may be different from the original source code.
12
///       See implementation for details.
13
1209
pub fn get_tree(src: &str) -> Option<(Tree, String)> {
14
1209
    let mut parser = Parser::new();
15
1209
    parser.set_language(&LANGUAGE.into()).unwrap();
16

            
17
1209
    parser.parse(src, None).and_then(|tree| {
18
1209
        let root = tree.root_node();
19
1209
        if root.is_error() {
20
            return None;
21
1209
        }
22

            
23
1209
        let children: Vec<_> = named_children(&root).collect();
24
1209
        let first_child = children.first()?;
25

            
26
        // HACK: Tree-sitter can only parse a complete program from top to bottom, not an individual bit of syntax.
27
        // See: https://github.com/tree-sitter/tree-sitter/issues/711 and linked issues.
28
        // However, we can use a dummy _FRAGMENT_EXPRESSION prefix (which we insert as necessary)
29
        // to trick the parser into accepting an isolated expression.
30
        // This way we can parse an isolated expression and it is only slightly cursed :)
31
1209
        if first_child.is_error() {
32
310
            if src.starts_with("_FRAGMENT_EXPRESSION") {
33
1
                None
34
            } else {
35
309
                get_tree(&format!("_FRAGMENT_EXPRESSION {src}"))
36
            }
37
        } else {
38
899
            Some((tree, src.to_string()))
39
        }
40
1209
    })
41
1209
}
42

            
43
/// Get the named children of a node
44
2899
pub fn named_children<'a>(node: &'a Node<'a>) -> impl Iterator<Item = Node<'a>> + 'a {
45
2899
    (0..node.named_child_count())
46
5538
        .filter_map(|i| u32::try_from(i).ok().and_then(|i| node.named_child(i)))
47
2899
}
48

            
49
1196
pub fn node_is_expression(node: &Node) -> bool {
50
480
    matches!(
51
1196
        node.kind(),
52
1196
        "bool_expr" | "arithmetic_expr" | "comparison_expr" | "atom"
53
    )
54
1196
}
55

            
56
/// Get all top-level nodes that match the given predicate
57
239
pub fn query_toplevel<'a>(
58
239
    node: &'a Node<'a>,
59
239
    predicate: &'a dyn Fn(&Node<'a>) -> bool,
60
239
) -> impl Iterator<Item = Node<'a>> + 'a {
61
726
    WalkDFS::with_retract(node, predicate).filter(|n| n.is_named() && predicate(n))
62
239
}
63

            
64
/// Get all meta-variable names in a node
65
1
pub fn get_metavars<'a>(node: &'a Node<'a>, src: &'a str) -> impl Iterator<Item = String> + 'a {
66
14
    query_toplevel(node, &|n| n.kind() == "metavar").filter_map(|child| {
67
1
        child
68
1
            .named_child(0)
69
1
            .map(|name| src[name.start_byte()..name.end_byte()].to_string())
70
1
    })
71
1
}
72

            
73
mod test {
74
    #[allow(unused)]
75
    use super::*;
76

            
77
    #[test]
78
1
    fn test_get_metavars() {
79
1
        let src = "such that &x = y";
80
1
        let (tree, _) = get_tree(src).unwrap();
81
1
        let root = tree.root_node();
82
1
        let metavars = get_metavars(&root, src).collect::<Vec<_>>();
83
1
        assert_eq!(metavars, vec!["x"]);
84
1
    }
85
}