1 files changed, 29 insertions, 6 deletions

diff --git a/impl/EquationSystem.hpp b/impl/EquationSystem.hpp
index d09bef2..c7e5f4c 100644
--- a/impl/EquationSystem.hpp
+++ b/impl/EquationSystem.hpp
@@ -7,12 +7,16 @@
 #include "Operator.hpp"
 #include "Expression.hpp"
 #include "VariableAssignment.hpp"
+#include "IdMap.hpp"
 
 template<typename Domain>
 struct MaxStrategy;
 
 template<typename Domain>
 struct EquationSystem {
+  EquationSystem()
+    : _expr_to_var(NULL) { }
+
   virtual ~EquationSystem() {
     for (typename std::set<Expression<Domain>*>::iterator it = _expressions.begin();
          it != _expressions.end();
@@ -24,6 +28,8 @@ struct EquationSystem {
          ++it) {
       delete *it;
     }
+    if (_expr_to_var)
+      delete _expr_to_var;
   }
 
   MaxExpression<Domain>& maxExpression(const std::vector<Expression<Domain>*>& arguments) {
@@ -86,7 +92,7 @@ struct EquationSystem {
   StableVariableAssignment<Domain>* assignment(const Domain& value) const {
     return new StableVariableAssignment<Domain>(_variables.size(), value);
   }
-  VariableAssignment<Domain>* eval(const VariableAssignment<Domain>& rho) const {
+  StableVariableAssignment<Domain>* eval(const VariableAssignment<Domain>& rho) const {
     StableVariableAssignment<Domain>* result = this->assignment(-infinity<Domain>());
     for (unsigned int i = 0, length = _variables.size();
          i < length;
@@ -97,7 +103,7 @@ struct EquationSystem {
     }
     return result;
   }
-  VariableAssignment<Domain>* eval(const VariableAssignment<Domain>& rho, const MaxStrategy<Domain>& strat) const {
+  StableVariableAssignment<Domain>* eval(const VariableAssignment<Domain>& rho, const MaxStrategy<Domain>& strat) const {
     StableVariableAssignment<Domain>* result = this->assignment(-infinity<Domain>());
     for (unsigned int i = 0, length = _variables.size();
          i < length;
@@ -109,12 +115,28 @@ struct EquationSystem {
     return result;
   }
 
-  Variable<Domain>* varFromExpr(const Expression<Domain>& expr) const {
+  void indexMaxExpressions() {
+    _expr_to_var = new IdMap<MaxExpression<Domain>,Variable<Domain>*>(maxExpressionCount(), NULL); 
     for (unsigned int i = 0, length = _right_sides.size(); i < length; ++i) {
-      if (_right_sides[i] == &expr)
-        return _variables[i];
+      (*_expr_to_var)[*_right_sides[i]] = _variables[i];
+    }
+  }
+
+  Variable<Domain>* varFromExpr(const Expression<Domain>& expr) const {
+    if (_expr_to_var) { // we've indexed:
+      auto* maxExpr = dynamic_cast<const MaxExpression<Domain>*>(&expr);
+      if (maxExpr) {
+        return (*_expr_to_var)[*maxExpr];
+      } else {
+        return NULL;
+      }
+    } else {
+      for (unsigned int i = 0, length = _right_sides.size(); i < length; ++i) {
+        if (_right_sides[i] == &expr)
+          return _variables[i];
+      }
+      return NULL;
     }
-    return NULL;
   }
 
   virtual bool equalAssignments(const VariableAssignment<Domain>& l, const VariableAssignment<Domain>& r) const {
@@ -141,6 +163,7 @@ struct EquationSystem {
   std::set<Expression<Domain>*> _expressions;
   std::vector<Variable<Domain>*> _variables;
   std::map<std::string, Variable<Domain>*> _variable_names;
+  IdMap<MaxExpression<Domain>, Variable<Domain>*>* _expr_to_var;
   std::vector<MaxExpression<Domain>*> _max_expressions;
   std::vector<MaxExpression<Domain>*> _right_sides;
 };