bel/parser/
references.rs

1use std::collections::HashSet;
2
3use crate::common::ast::{Expr, IdedExpr};
4
5/// A collection of all the references that an expression makes to variables and functions.
6pub struct ExpressionReferences<'expr> {
7    variables: HashSet<&'expr str>,
8    functions: HashSet<&'expr str>,
9}
10
11impl ExpressionReferences<'_> {
12    /// Returns true if the expression references the provided variable name.
13    ///
14    /// # Example
15    /// ```rust
16    /// # use bel::parser::Parser;
17    /// let expression = Parser::new().parse("foo.bar == true").unwrap();
18    /// let references = expression.references();
19    /// assert!(references.has_variable("foo"));
20    /// ```
21    pub fn has_variable(&self, name: impl AsRef<str>) -> bool {
22        self.variables.contains(name.as_ref())
23    }
24
25    /// Returns true if the expression references the provided function name.
26    ///
27    /// # Example
28    /// ```rust
29    /// # use bel::parser::Parser;
30    /// let expression = Parser::new().parse("length(foo) > 0").unwrap();
31    /// let references = expression.references();
32    /// assert!(references.has_function("length"));
33    /// ```
34    pub fn has_function(&self, name: impl AsRef<str>) -> bool {
35        self.functions.contains(name.as_ref())
36    }
37
38    /// Returns a list of all variables referenced in the expression.
39    ///
40    /// # Example
41    /// ```rust
42    /// # use bel::parser::Parser;
43    /// let expression = Parser::new().parse("foo.bar == true").unwrap();
44    /// let references = expression.references();
45    /// assert_eq!(vec!["foo"], references.variables());
46    /// ```
47    pub fn variables(&self) -> Vec<&str> {
48        self.variables.iter().copied().collect()
49    }
50
51    /// Returns a list of all functions referenced in the expression.
52    ///
53    /// # Example
54    /// ```rust
55    /// # use bel::parser::Parser;
56    /// let expression = Parser::new().parse("length(foo) > 0").unwrap();
57    /// let references = expression.references();
58    /// assert!(references.functions().contains(&"_>_"));
59    /// assert!(references.functions().contains(&"length"));
60    /// ```
61    pub fn functions(&self) -> Vec<&str> {
62        self.functions.iter().copied().collect()
63    }
64}
65
66impl IdedExpr {
67    /// Returns a set of all variables and functions referenced in the expression.
68    ///
69    /// # Example
70    /// ```rust
71    /// # use bel::parser::Parser;
72    /// let expression = Parser::new().parse("foo && length(foo) > 0").unwrap();
73    /// let references = expression.references();
74    ///
75    /// assert!(references.has_variable("foo"));
76    /// assert!(references.has_function("length"));
77    /// ```
78    pub fn references(&self) -> ExpressionReferences<'_> {
79        let mut variables = HashSet::new();
80        let mut functions = HashSet::new();
81        self._references(&mut variables, &mut functions);
82        ExpressionReferences {
83            variables,
84            functions,
85        }
86    }
87
88    /// Internal recursive function to collect all variable and function references in the expression.
89    fn _references<'expr>(&'expr self, variables: &mut HashSet<&'expr str>, functions: &mut HashSet<&'expr str>) {
90        match &self.expr {
91            Expr::Unspecified => {}
92            Expr::Call(call) => {
93                functions.insert(&call.func_name);
94                if let Some(target) = &call.target {
95                    target._references(variables, functions);
96                }
97                for arg in &call.args {
98                    arg._references(variables, functions);
99                }
100            }
101            Expr::Comprehension(comp) => {
102                comp.iter_range._references(variables, functions);
103                comp.accu_init._references(variables, functions);
104                comp.loop_cond._references(variables, functions);
105                comp.loop_step._references(variables, functions);
106                comp.result._references(variables, functions);
107            }
108            Expr::Ident(name) => {
109                // todo! Might want to make this "smarter" (are we in a comprehension?) and better encode these in const
110                if !name.starts_with('@') {
111                    variables.insert(name);
112                }
113            }
114            Expr::List(list) => {
115                for elem in &list.elements {
116                    elem._references(variables, functions);
117                }
118            }
119            Expr::Literal(_) => {}
120            Expr::Map(map) => {
121                for entry in &map.entries {
122                    match &entry.expr {
123                        crate::common::ast::EntryExpr::StructField(field) => {
124                            field.value._references(variables, functions);
125                        }
126                        crate::common::ast::EntryExpr::MapEntry(map_entry) => {
127                            map_entry.key._references(variables, functions);
128                            map_entry.value._references(variables, functions);
129                        }
130                    }
131                }
132            }
133            Expr::Select(select) => {
134                select.operand._references(variables, functions);
135            }
136            Expr::Struct(struct_expr) => {
137                for entry in &struct_expr.entries {
138                    match &entry.expr {
139                        crate::common::ast::EntryExpr::StructField(field) => {
140                            field.value._references(variables, functions);
141                        }
142                        crate::common::ast::EntryExpr::MapEntry(map_entry) => {
143                            map_entry.key._references(variables, functions);
144                            map_entry.value._references(variables, functions);
145                        }
146                    }
147                }
148            }
149        }
150    }
151}
bel/parser/references.rs

bel/parser/
references.rs
