Add various helper methods to the syntax tree

src/syntax_tree/asp.rs

@@ -1,14 +1,17 @@
-use crate::{
-    formatting::asp::default::Format,
-    parsing::asp::pest::{
-        AtomParser, AtomicFormulaParser, BinaryOperatorParser, BodyParser, ComparisonParser,
-        HeadParser, LiteralParser, PrecomputedTermParser, ProgramParser, RelationParser,
-        RuleParser, SignParser, TermParser, UnaryOperatorParser, VariableParser,
+use {
+    crate::{
+        formatting::asp::default::Format,
+        parsing::asp::pest::{
+            AtomParser, AtomicFormulaParser, BinaryOperatorParser, BodyParser, ComparisonParser,
+            HeadParser, LiteralParser, PrecomputedTermParser, ProgramParser, RelationParser,
+            RuleParser, SignParser, TermParser, UnaryOperatorParser, VariableParser,
+        },
+        syntax_tree::{impl_node, Node},
     },
-    syntax_tree::{impl_node, Node},
+    std::collections::HashSet,
 };
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum PrecomputedTerm {
     Infimum,
     Numeral(isize),
@@ -18,19 +21,19 @@ pub enum PrecomputedTerm {
 
 impl_node!(PrecomputedTerm, Format, PrecomputedTermParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Variable(pub String);
 
 impl_node!(Variable, Format, VariableParser);
 
-#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
 pub enum UnaryOperator {
     Negative,
 }
 
 impl_node!(UnaryOperator, Format, UnaryOperatorParser);
 
-#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
 pub enum BinaryOperator {
     Add,
     Subtract,
@@ -42,7 +45,7 @@ pub enum BinaryOperator {
 
 impl_node!(BinaryOperator, Format, BinaryOperatorParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum Term {
     PrecomputedTerm(PrecomputedTerm),
     Variable(Variable),
@@ -59,15 +62,40 @@ pub enum Term {
 
 impl_node!(Term, Format, TermParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+impl Term {
+    pub fn variables(&self) -> HashSet<Variable> {
+        match &self {
+            Term::PrecomputedTerm(_) => HashSet::new(),
+            Term::Variable(v) => HashSet::from([v.clone()]),
+            Term::UnaryOperation { arg, .. } => arg.variables(),
+            Term::BinaryOperation { lhs, rhs, .. } => {
+                let mut vars = lhs.variables();
+                vars.extend(rhs.variables());
+                vars
+            }
+        }
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Atom {
     pub predicate: String,
     pub terms: Vec<Term>,
 }
 
 impl_node!(Atom, Format, AtomParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+impl Atom {
+    pub fn variables(&self) -> HashSet<Variable> {
+        let mut vars = HashSet::new();
+        for term in self.terms.iter() {
+            vars.extend(term.variables())
+        }
+        vars
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum Sign {
     NoSign,
     Negation,
@@ -76,15 +104,21 @@ pub enum Sign {
 
 impl_node!(Sign, Format, SignParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Literal {
     pub sign: Sign,
     pub atom: Atom,
 }
 
 impl_node!(Literal, Format, LiteralParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+impl Literal {
+    pub fn variables(&self) -> HashSet<Variable> {
+        self.atom.variables()
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum Relation {
     Equal,
     NotEqual,
@@ -96,7 +130,7 @@ pub enum Relation {
 
 impl_node!(Relation, Format, RelationParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Comparison {
     pub relation: Relation,
     pub lhs: Term,
@@ -105,15 +139,32 @@ pub struct Comparison {
 
 impl_node!(Comparison, Format, ComparisonParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+impl Comparison {
+    pub fn variables(&self) -> HashSet<Variable> {
+        let mut vars = self.lhs.variables();
+        vars.extend(self.rhs.variables());
+        vars
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum AtomicFormula {
     Literal(Literal),
     Comparison(Comparison),
 }
 
 impl_node!(AtomicFormula, Format, AtomicFormulaParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+impl AtomicFormula {
+    pub fn variables(&self) -> HashSet<Variable> {
+        match &self {
+            AtomicFormula::Literal(l) => l.variables(),
+            AtomicFormula::Comparison(c) => c.variables(),
+        }
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum Head {
     Basic(Atom),
     Choice(Atom),
@@ -122,24 +173,85 @@ pub enum Head {
 
 impl_node!(Head, Format, HeadParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+impl Head {
+    pub fn predicate(&self) -> Option<&str> {
+        match self {
+            Head::Basic(a) => Some(&a.predicate),
+            Head::Choice(a) => Some(&a.predicate),
+            Head::Falsity => None,
+        }
+    }
+
+    pub fn terms(&self) -> Option<&[Term]> {
+        match self {
+            Head::Basic(a) => Some(&a.terms),
+            Head::Choice(a) => Some(&a.terms),
+            Head::Falsity => None,
+        }
+    }
+
+    pub fn arity(&self) -> usize {
+        match self {
+            Head::Basic(a) => a.terms.len(),
+            Head::Choice(a) => a.terms.len(),
+            Head::Falsity => 0,
+        }
+    }
+
+    pub fn variables(&self) -> HashSet<Variable> {
+        match &self {
+            Head::Basic(a) | Head::Choice(a) => a.variables(),
+            Head::Falsity => HashSet::new(),
+        }
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Body {
     pub formulas: Vec<AtomicFormula>,
 }
 
 impl_node!(Body, Format, BodyParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+impl Body {
+    pub fn variables(&self) -> HashSet<Variable> {
+        let mut vars = HashSet::new();
+        for formula in self.formulas.iter() {
+            vars.extend(formula.variables())
+        }
+        vars
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Rule {
     pub head: Head,
     pub body: Body,
 }
 
 impl_node!(Rule, Format, RuleParser);
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+impl Rule {
+    pub fn variables(&self) -> HashSet<Variable> {
+        let mut vars = self.head.variables();
+        vars.extend(self.body.variables());
+        vars
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Program {
     pub rules: Vec<Rule>,
 }
 
 impl_node!(Program, Format, ProgramParser);
+
+impl Program {
+    pub fn variables(&self) -> HashSet<Variable> {
+        let mut vars = HashSet::new();
+        for rule in self.rules.iter() {
+            vars.extend(rule.variables())
+        }
+        vars
+    }
+}