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
386
pub fn get_tree(src: &str) -> Option<(Tree, String)> {
14
386
    let mut parser = Parser::new();
15
386
    parser.set_language(&LANGUAGE.into()).unwrap();
16

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

            
23
386
        let children: Vec<_> = named_children(&root).collect();
24
386
        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
386
        if first_child.is_error() {
32
193
            if src.starts_with("_FRAGMENT_EXPRESSION") {
33
                None
34
            } else {
35
193
                get_tree(&format!("_FRAGMENT_EXPRESSION {src}"))
36
            }
37
        } else {
38
193
            Some((tree, src.to_string()))
39
        }
40
386
    })
41
386
}
42

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

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

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

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

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

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