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()).filter_map(|i| node.named_child(i))
46}
47
48pub fn node_is_expression(node: &Node) -> bool {
49    matches!(
50        node.kind(),
51        "bool_expr" | "arithmetic_expr" | "comparison_expr" | "atom"
52    )
53}
54
55/// Get all top-level nodes that match the given predicate
56pub fn query_toplevel<'a>(
57    node: &'a Node<'a>,
58    predicate: &'a dyn Fn(&Node<'a>) -> bool,
59) -> impl Iterator<Item = Node<'a>> + 'a {
60    WalkDFS::with_retract(node, predicate).filter(|n| n.is_named() && predicate(n))
61}
62
63/// Get all meta-variable names in a node
64pub 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
72mod 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}
conjure_cp_essence_parser/parser/util.rs

conjure_cp_essence_parser/parser/
util.rs
