Compute free variables of a Formula

potassco · Dec 19, 2023 · 2a032a1 · 2a032a1
1 parent ee0af51
commit 2a032a1
Showing 1 changed file with 34 additions and 1 deletion.
diff --git a/src/syntax_tree/fol.rs b/src/syntax_tree/fol.rs
@@ -337,6 +337,22 @@ impl Formula {
         }
     }
 
+    pub fn free_variables(&self) -> HashSet<Variable> {
+        match &self {
+            Formula::QuantifiedFormula {
+                quantification,
+                formula,
+            } => {
+                let mut vars = formula.variables();
+                for var in &quantification.variables {
+                    vars.remove(var);
+                }
+                vars
+            }
+            x => x.variables(),
+        }
+    }
+
     pub fn predicates(&self) -> HashSet<Predicate> {
         match &self {
             Formula::AtomicFormula(f) => f.predicates(),
@@ -360,7 +376,10 @@ impl_node!(Theory, Format, TheoryParser);
 
 #[cfg(test)]
 mod tests {
-    use super::Formula;
+    use {
+        super::{Formula, Sort, Variable},
+        std::collections::HashSet,
+    };
 
     #[test]
     fn test_formula_conjoin() {
@@ -389,4 +408,18 @@ mod tests {
             )
         }
     }
+
+    #[test]
+    fn test_formula_free_variables() {
+        assert_eq!(
+            "forall X (X = Y)"
+                .parse::<Formula>()
+                .unwrap()
+                .free_variables(),
+            HashSet::from([Variable {
+                name: "Y".into(),
+                sort: Sort::General
+            }])
+        )
+    }
 }