From 54d65e2dd06211867929dd1d25c1ffbd31e5bd1c Mon Sep 17 00:00:00 2001
From: Tobias Stolzmann <tobias.stolzmann@gmail.com>
Date: Tue, 14 Nov 2023 18:22:31 +0100
Subject: [PATCH] Add variable-methods to the FOL nodes

---
 src/syntax_tree/fol.rs | 105 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 95 insertions(+), 10 deletions(-)

diff --git a/src/syntax_tree/fol.rs b/src/syntax_tree/fol.rs
index 3cf7e7aa..dc805d41 100644
--- a/src/syntax_tree/fol.rs
+++ b/src/syntax_tree/fol.rs
@@ -1,12 +1,15 @@
-use crate::{
-    formatting::fol::default::Format,
-    parsing::fol::pest::{
-        AtomParser, AtomicFormulaParser, BasicIntegerTermParser, BinaryConnectiveParser,
-        BinaryOperatorParser, ComparisonParser, FormulaParser, GeneralTermParser, GuardParser,
-        IntegerTermParser, QuantificationParser, QuantifierParser, RelationParser, TheoryParser,
-        UnaryConnectiveParser, UnaryOperatorParser, VariableParser,
+use {
+    crate::{
+        formatting::fol::default::Format,
+        parsing::fol::pest::{
+            AtomParser, AtomicFormulaParser, BasicIntegerTermParser, BinaryConnectiveParser,
+            BinaryOperatorParser, ComparisonParser, FormulaParser, GeneralTermParser, GuardParser,
+            IntegerTermParser, QuantificationParser, QuantifierParser, RelationParser,
+            TheoryParser, UnaryConnectiveParser, UnaryOperatorParser, VariableParser,
+        },
+        syntax_tree::{impl_node, Node},
     },
-    syntax_tree::{impl_node, Node},
+    std::{collections::HashSet, hash::Hash},
 };
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
@@ -19,6 +22,18 @@ pub enum BasicIntegerTerm {
 
 impl_node!(BasicIntegerTerm, Format, BasicIntegerTermParser);
 
+impl BasicIntegerTerm {
+    pub fn variables(&self) -> HashSet<Variable> {
+        match &self {
+            BasicIntegerTerm::IntegerVariable(v) => HashSet::from([Variable {
+                name: v.to_string(),
+                sort: Sort::Integer,
+            }]),
+            _ => HashSet::new(),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum UnaryOperator {
     Negative,
@@ -51,6 +66,20 @@ pub enum IntegerTerm {
 
 impl_node!(IntegerTerm, Format, IntegerTermParser);
 
+impl IntegerTerm {
+    pub fn variables(&self) -> HashSet<Variable> {
+        match &self {
+            IntegerTerm::BasicIntegerTerm(t) => t.variables(),
+            IntegerTerm::UnaryOperation { arg: t, .. } => t.variables(),
+            IntegerTerm::BinaryOperation { lhs, rhs, .. } => {
+                let mut vars = lhs.variables();
+                vars.extend(rhs.variables());
+                vars
+            }
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum GeneralTerm {
     Symbol(String),
@@ -60,6 +89,19 @@ pub enum GeneralTerm {
 
 impl_node!(GeneralTerm, Format, GeneralTermParser);
 
+impl GeneralTerm {
+    pub fn variables(&self) -> HashSet<Variable> {
+        match &self {
+            GeneralTerm::Symbol(_) => HashSet::new(),
+            GeneralTerm::GeneralVariable(v) => HashSet::from([Variable {
+                name: v.to_string(),
+                sort: Sort::General,
+            }]),
+            GeneralTerm::IntegerTerm(t) => t.variables(),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Atom {
     pub predicate: String,
@@ -88,6 +130,12 @@ pub struct Guard {
 
 impl_node!(Guard, Format, GuardParser);
 
+impl Guard {
+    pub fn variables(&self) -> HashSet<Variable> {
+        self.term.variables()
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Comparison {
     pub term: GeneralTerm,
@@ -106,6 +154,28 @@ pub enum AtomicFormula {
 
 impl_node!(AtomicFormula, Format, AtomicFormulaParser);
 
+impl AtomicFormula {
+    pub fn variables(&self) -> HashSet<Variable> {
+        match &self {
+            AtomicFormula::Falsity | AtomicFormula::Truth => HashSet::new(),
+            AtomicFormula::Atom(a) => {
+                let mut vars = HashSet::new();
+                for t in a.terms.iter() {
+                    vars.extend(t.variables());
+                }
+                vars
+            }
+            AtomicFormula::Comparison(c) => {
+                let mut vars = c.term.variables();
+                for guard in c.guards.iter() {
+                    vars.extend(guard.variables())
+                }
+                vars
+            }
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum UnaryConnective {
     Negation,
@@ -129,7 +199,7 @@ pub struct Quantification {
 
 impl_node!(Quantification, Format, QuantificationParser);
 
-#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
 pub enum Sort {
     Integer,
     General,
@@ -137,7 +207,7 @@ pub enum Sort {
 
 // TODO: Should Sort be a Node?
 
-#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
 pub struct Variable {
     pub name: String,
     pub sort: Sort,
@@ -176,6 +246,21 @@ pub enum Formula {
 
 impl_node!(Formula, Format, FormulaParser);
 
+impl Formula {
+    pub fn variables(&self) -> HashSet<Variable> {
+        match &self {
+            Formula::AtomicFormula(f) => f.variables(),
+            Formula::UnaryFormula { formula, .. } => formula.variables(),
+            Formula::BinaryFormula { lhs, rhs, .. } => {
+                let mut vars = lhs.variables();
+                vars.extend(rhs.variables());
+                vars
+            }
+            Formula::QuantifiedFormula { formula, .. } => formula.variables(),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Theory {
     pub formulas: Vec<Formula>,