From 5c28922e5b239d1bc6849235fd6c2ffbc7185c42 Mon Sep 17 00:00:00 2001
From: Tobias Stolzmann <tobias.stolzmann@gmail.com>
Date: Tue, 19 Dec 2023 21:44:14 +0100
Subject: [PATCH] Fix the computation of free variables

---
 src/syntax_tree/fol.rs | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/src/syntax_tree/fol.rs b/src/syntax_tree/fol.rs
index 8e9ae77a..de02cffc 100644
--- a/src/syntax_tree/fol.rs
+++ b/src/syntax_tree/fol.rs
@@ -339,6 +339,13 @@ impl Formula {
 
     pub fn free_variables(&self) -> HashSet<Variable> {
         match &self {
+            Formula::AtomicFormula(f) => f.variables(),
+            Formula::UnaryFormula { formula, .. } => formula.free_variables(),
+            Formula::BinaryFormula { lhs, rhs, .. } => {
+                let mut vars = lhs.free_variables();
+                vars.extend(rhs.free_variables());
+                vars
+            }
             Formula::QuantifiedFormula {
                 quantification,
                 formula,
@@ -349,7 +356,6 @@ impl Formula {
                 }
                 vars
             }
-            x => x.variables(),
         }
     }
 
@@ -376,10 +382,7 @@ impl_node!(Theory, Format, TheoryParser);
 
 #[cfg(test)]
 mod tests {
-    use {
-        super::{Formula, Sort, Variable},
-        std::collections::HashSet,
-    };
+    use super::Formula;
 
     #[test]
     fn test_formula_conjoin() {
@@ -411,15 +414,15 @@ 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
-            }])
-        )
+        for (src, target) in [
+            ("forall X (X = Y)", vec!["Y"]),
+            ("forall X (X = Y) and Y = Z", vec!["Y", "Z"]),
+            ("forall X exists Y (X = Y)", vec![]),
+        ] {
+            assert_eq!(
+                src.parse::<Formula>().unwrap().free_variables(),
+                target.iter().map(|x| x.parse().unwrap()).collect(),
+            )
+        }
     }
 }