2 files changed, 16 insertions, 8 deletions

diff --git a/src/theory/uf/equality_engine.cpp b/src/theory/uf/equality_engine.cpp
index 9376c858d..b9114cb51 100644
--- a/src/theory/uf/equality_engine.cpp
+++ b/src/theory/uf/equality_engine.cpp
@@ -214,6 +214,8 @@ EqualityNodeId EqualityEngine::newNode(TNode node) {
   d_isConstant.push_back(node.isConst());
   // Mark Boolean nodes
   d_isBoolean.push_back(false);
+  // Mark the node as internal by default
+  d_isInternal.push_back(true);
   // Add the equality node to the nodes
   d_equalityNodes.push_back(EqualityNode(newId));
 
@@ -249,7 +251,8 @@ void EqualityEngine::addTerm(TNode t) {
   if (t.getKind() == kind::EQUAL) {
     addTerm(t[0]);
     addTerm(t[1]);
-    result = newApplicationNode(t, getNodeId(t[0]), getNodeId(t[1]), true); 
+    result = newApplicationNode(t, getNodeId(t[0]), getNodeId(t[1]), true);
+    d_isInternal[result] = false;
   } else if (t.getNumChildren() > 0 && d_congruenceKinds[t.getKind()]) {
     // Add the operator
     TNode tOp = t.getOperator();
@@ -262,9 +265,11 @@ void EqualityEngine::addTerm(TNode t) {
       // Add the application
       result = newApplicationNode(t, result, getNodeId(t[i]), false);
     }
+    d_isInternal[result] = false;
   } else {
     // Otherwise we just create the new id
     result = newNode(t);
+    d_isInternal[result] = false;
   }
 
   if (t.getType().isBoolean()) {
@@ -798,6 +803,7 @@ void EqualityEngine::backtrack() {
     d_nodeIndividualTrigger.resize(d_nodesCount);
     d_isConstant.resize(d_nodesCount);
     d_isBoolean.resize(d_nodesCount);
+    d_isInternal.resize(d_nodesCount);
     d_equalityGraph.resize(d_nodesCount);
     d_equalityNodes.resize(d_nodesCount);
   }
diff --git a/src/theory/uf/equality_engine.h b/src/theory/uf/equality_engine.h
index 45d1b4acf..cae169081 100644
--- a/src/theory/uf/equality_engine.h
+++ b/src/theory/uf/equality_engine.h
@@ -402,6 +402,12 @@ private:
   std::vector<bool> d_isBoolean;
 
   /**
+   * Map from ids to whether the nods is internal. An internal node is a node
+   * that corresponds to a partially currified node, for example.
+   */
+  std::vector<bool> d_isInternal;
+
+  /**
    * Adds the trigger with triggerId to the beginning of the trigger list of the node with id nodeId.
    */
   void addTriggerToList(EqualityNodeId nodeId, TriggerId triggerId);
@@ -791,14 +797,13 @@ class EqClassesIterator {
 
   const eq::EqualityEngine* d_ee;
   size_t d_it;
-  std::vector< Node > d_visited;
 public:
 
   EqClassesIterator(): d_ee(NULL), d_it(0){ }
   EqClassesIterator(const eq::EqualityEngine* ee) : d_ee(ee) {
     d_it = 0;
-    if ( d_it < d_ee->d_nodesCount &&
-         d_ee->getRepresentative(d_ee->d_nodes[d_it]) != d_ee->d_nodes[d_it] ) {
+    // Go to the first non-internal node that is it's own representative
+    if (d_it < d_ee->d_nodesCount && (d_ee->d_isInternal[d_it] || d_ee->getRepresentative(d_ee->d_nodes[d_it]) != d_ee->d_nodes[d_it])) {
       ++*this;
     }
   }
@@ -812,11 +817,8 @@ public:
     return !(*this == i);
   }
   EqClassesIterator& operator++() {
-    d_visited.push_back( d_ee->d_nodes[d_it] );
     ++d_it;
-    while ( d_it<d_ee->d_nodesCount &&
-            ( d_ee->getRepresentative(d_ee->d_nodes[d_it]) != d_ee->d_nodes[d_it] ||
-              std::find( d_visited.begin(), d_visited.end(), d_ee->d_nodes[d_it] )!=d_visited.end() ) ) { // this line is necessary for ignoring duplicates
+    while (d_it<d_ee->d_nodesCount && (d_ee->d_isInternal[d_it] || d_ee->getRepresentative(d_ee->d_nodes[d_it]) != d_ee->d_nodes[d_it])) {
       ++d_it;
     }
     return *this;