conjure_cp_essence_parser/parser/
util.rs

1use tree_sitter::{Node, Parser, Tree};
2use tree_sitter_essence::LANGUAGE;
3
4use 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.
13pub fn get_tree(src: &str) -> Option<(Tree, String)> {
14    let mut parser = Parser::new();
15    parser.set_language(&LANGUAGE.into()).unwrap();
16
17    parser.parse(src, None).and_then(|tree| {
18        let root = tree.root_node();
19        if root.is_error() {
20            return None;
21        }
22
23        let children: Vec<_> = named_children(&root).collect();
24        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        if first_child.is_error() {
32            if src.starts_with("_FRAGMENT_EXPRESSION") {
33                None
34            } else {
35                get_tree(&format!("_FRAGMENT_EXPRESSION {src}"))
36            }
37        } else {
38            Some((tree, src.to_string()))
39        }
40    })
41}
42
43/// Get the named children of a node
44pub fn named_children<'a>(node: &'a Node<'a>) -> impl Iterator<Item = Node<'a>> + 'a {
45    (0..node.named_child_count())
46        .filter_map(|i| u32::try_from(i).ok().and_then(|i| node.named_child(i)))
47}
48
49pub fn node_is_expression(node: &Node) -> bool {
50    matches!(
51        node.kind(),
52        "bool_expr" | "arithmetic_expr" | "comparison_expr" | "atom"
53    )
54}
55
56/// Get all top-level nodes that match the given predicate
57pub fn query_toplevel<'a>(
58    node: &'a Node<'a>,
59    predicate: &'a dyn Fn(&Node<'a>) -> bool,
60) -> impl Iterator<Item = Node<'a>> + 'a {
61    WalkDFS::with_retract(node, predicate).filter(|n| n.is_named() && predicate(n))
62}
63
64/// Get all meta-variable names in a node
65pub fn get_metavars<'a>(node: &'a Node<'a>, src: &'a str) -> impl Iterator<Item = String> + 'a {
66    query_toplevel(node, &|n| n.kind() == "metavar").filter_map(|child| {
67        child
68            .named_child(0)
69            .map(|name| src[name.start_byte()..name.end_byte()].to_string())
70    })
71}
72
73mod test {
74    #[allow(unused)]
75    use super::*;
76
77    #[test]
78    fn test_get_metavars() {
79        let src = "such that &x = y";
80        let (tree, _) = get_tree(src).unwrap();
81        let root = tree.root_node();
82        let metavars = get_metavars(&root, src).collect::<Vec<_>>();
83        assert_eq!(metavars, vec!["x"]);
84    }
85}
conjure_cp_essence_parser/parser/util.rs

conjure_cp_essence_parser/parser/
util.rs
