vortex_expr/traversal/
vars.rs1use vortex_error::VortexResult;
2use vortex_utils::aliases::hash_set::HashSet;
3
4use crate::traversal::{NodeVisitor, TraversalOrder};
5use crate::{ExprRef, Identifier, Var};
6
7#[derive(Default)]
8pub struct VarsCollector {
9 ids: HashSet<Identifier>,
10}
11
12impl VarsCollector {
13 pub fn new() -> Self {
14 Self {
15 ids: HashSet::new(),
16 }
17 }
18
19 pub fn into_vars(self) -> HashSet<Identifier> {
20 self.ids
21 }
22}
23
24impl NodeVisitor<'_> for VarsCollector {
25 type NodeTy = ExprRef;
26
27 fn visit_up(&mut self, node: &ExprRef) -> VortexResult<TraversalOrder> {
28 if let Some(var) = node.as_any().downcast_ref::<Var>() {
29 self.ids.insert(var.var().clone());
30 }
31 Ok(TraversalOrder::Continue)
32 }
33}