1 files changed, 14 insertions, 4 deletions
diff --git a/impl/EquationSystem.hpp b/impl/EquationSystem.hpp
index 50a0e45..5549674 100644
--- a/impl/EquationSystem.hpp
+++ b/impl/EquationSystem.hpp
@@ -13,11 +13,13 @@ struct EquationSystem {
   virtual ~EquationSystem() { }
   virtual const Expression<Domain>* operator[](const Variable<Domain>&) const = 0;
   virtual VariableAssignment<Domain>* eval(const VariableAssignment<Domain>& rho) const = 0;
+
   virtual unsigned int variableCount() const = 0;
   virtual Variable<Domain>& variable(unsigned int) const = 0;
+
   virtual StableVariableAssignment<Domain>* assignment(const Domain& value) const = 0;
+  virtual Variable<Domain>* varFromExpr(const Expression<Domain>& expr) const = 0;
   virtual bool equalAssignments(const VariableAssignment<Domain>&, const VariableAssignment<Domain>&) const = 0;
-
   virtual void print(std::ostream& cout) const = 0;
 };
 
@@ -86,10 +88,10 @@ struct ConcreteEquationSystem : public EquationSystem<Domain> {
     return *constant;
   }
 
-  Expression<Domain>* operator[](const Variable<Domain>& var) const {
+  MaxExpression<Domain>* operator[](const Variable<Domain>& var) const {
     return _right_sides[var.id()];
   }
-  Expression<Domain>*& operator[](const Variable<Domain>& var) {
+  MaxExpression<Domain>*& operator[](const Variable<Domain>& var) {
     return _right_sides[var.id()];
   }
 
@@ -108,6 +110,14 @@ struct ConcreteEquationSystem : public EquationSystem<Domain> {
     return result;
   }
 
+  Variable<Domain>* varFromExpr(const Expression<Domain>& expr) const {
+    for (unsigned int i = 0, length = _right_sides.size(); i < length; ++i) {
+      if (_right_sides[i] == &expr)
+        return _variables[i];
+    }
+    return NULL;
+  }
+
   virtual bool equalAssignments(const VariableAssignment<Domain>& l, const VariableAssignment<Domain>& r) const {
     for (unsigned int i = 0, length = _variables.size();
          i < length;
@@ -133,7 +143,7 @@ struct ConcreteEquationSystem : public EquationSystem<Domain> {
   std::vector<Variable<Domain>*> _variables;
   std::map<std::string, Variable<Domain>*> _variable_names;
   std::vector<MaxExpression<Domain>*> _max_expressions;
-  std::vector<Expression<Domain>*> _right_sides;
+  std::vector<MaxExpression<Domain>*> _right_sides;
 };
 
 template<typename T>