vortex_expr/traversal/
vars.rs

1use 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}