1 files changed, 46 insertions, 3 deletions

diff --git a/src/theory/uf/equality_engine.cpp b/src/theory/uf/equality_engine.cpp
index 90de8d0d2..636427ed1 100644
--- a/src/theory/uf/equality_engine.cpp
+++ b/src/theory/uf/equality_engine.cpp
@@ -79,8 +79,9 @@ EqualityEngine::~EqualityEngine() throw(AssertionException) {
 }
 
 
-EqualityEngine::EqualityEngine(context::Context* context, std::string name) 
+EqualityEngine::EqualityEngine(context::Context* context, std::string name)
 : ContextNotifyObj(context)
+, d_masterEqualityEngine(0)
 , d_context(context)
 , d_done(context, false)
 , d_performNotify(true)
@@ -102,6 +103,7 @@ EqualityEngine::EqualityEngine(context::Context* context, std::string name)
 
 EqualityEngine::EqualityEngine(EqualityEngineNotify& notify, context::Context* context, std::string name)
 : ContextNotifyObj(context)
+, d_masterEqualityEngine(0)
 , d_context(context)
 , d_done(context, false)
 , d_performNotify(true)
@@ -121,6 +123,11 @@ EqualityEngine::EqualityEngine(EqualityEngineNotify& notify, context::Context* c
   init();
 }
 
+void EqualityEngine::setMasterEqualityEngine(EqualityEngine* master) {
+  Assert(d_masterEqualityEngine == 0);
+  d_masterEqualityEngine = master;
+}
+
 void EqualityEngine::enqueue(const MergeCandidate& candidate, bool back) {
   Debug("equality") << d_name << "::eq::enqueue(" << d_nodes[candidate.t1Id] << ", " << d_nodes[candidate.t2Id] << ", " << candidate.type << ")" << std::endl;
   if (back) {
@@ -294,6 +301,11 @@ void EqualityEngine::addTerm(TNode t) {
     d_nodeIndividualTrigger[tId] = newTriggerTermSet();
   }
 
+  // If this is not an internal node, add it to the master
+  if (d_masterEqualityEngine && !d_isInternal[result]) {
+    d_masterEqualityEngine->addTerm(t);
+  }
+
   propagate();
 
   Debug("equality") << d_name << "::eq::addTerm(" << t << ") => " << result << std::endl;
@@ -1213,9 +1225,34 @@ void EqualityEngine::propagate() {
       continue;
     }
 
-    // Depending on the merge preference (such as size, or being a constant), merge them
+    // Vector to collect the triggered events
     std::vector<TriggerId> triggers;
-    if ((node2.getSize() > node1.getSize() && !d_isConstant[t1classId]) || d_isConstant[t2classId]) {
+
+    // Figure out the merge preference
+    EqualityNodeId mergeInto = t1classId;
+    if (d_isInternal[t2classId] != d_isInternal[t1classId]) {
+      // We always keep non-internal nodes as representatives: if any node in
+      // the class is non-internal, then the representative will be non-internal
+      if (d_isInternal[t1classId]) {
+        mergeInto = t2classId;
+      } else {
+        mergeInto = t1classId;
+      }
+    } else if (d_isConstant[t2classId] != d_isConstant[t1classId]) {
+      // We always keep constants as representatives: if any (at most one) node
+      // in the class in a constant, then the representative will be a constant
+      if (d_isConstant[t2classId]) {
+        mergeInto = t2classId;
+      } else {
+        mergeInto = t1classId;
+      }
+    } else if (node2.getSize() > node1.getSize()) {
+      // We always merge into the bigger class to reduce the amount of traversing
+      // we need to do
+      mergeInto = t2classId;
+    }
+
+    if (mergeInto == t2classId) {
       Debug("equality") << d_name << "::eq::propagate(): merging " << d_nodes[current.t1Id]<< " into " << d_nodes[current.t2Id] << std::endl;
       d_assertedEqualities.push_back(Equality(t2classId, t1classId));
       d_assertedEqualitiesCount = d_assertedEqualitiesCount + 1;
@@ -1231,6 +1268,12 @@ void EqualityEngine::propagate() {
       }
     }
 
+    // If not merging internal nodes, notify the master
+    if (d_masterEqualityEngine && !d_isInternal[mergeInto]) {
+      d_masterEqualityEngine->assertEqualityInternal(d_nodes[t1classId], d_nodes[t2classId], TNode::null());
+      d_masterEqualityEngine->propagate();
+    }
+
     // Notify the triggers
     if (d_performNotify && !d_done) {
       for (size_t trigger_i = 0, trigger_end = triggers.size(); trigger_i < trigger_end && !d_done; ++ trigger_i) {