Shared term manager tested and working

It is currently tracking all asserted equalities for simplicity.
Might want to check if this is a performance hit

1 files changed, 102 insertions, 29 deletions

diff --git a/src/theory/shared_term_manager.cpp b/src/theory/shared_term_manager.cpp
index 55087ecbd..564fb776f 100644
--- a/src/theory/shared_term_manager.cpp
+++ b/src/theory/shared_term_manager.cpp
@@ -31,11 +31,16 @@ SharedTermManager::SharedTermManager(TheoryEngine* engine,
 }
 
 
-SharedData* SharedTermManager::find(SharedData* pData) {
+SharedData* SharedTermManager::find(SharedData* pData) const {
+  // Check if pData is the representative
   SharedData* next = pData->getFind();
   if (next == pData) return pData;
+
+  // Check if its successor is the representative
   SharedData* nextNext = next->getFind();
   if (nextNext == next) return next;
+
+  // Otherwise, recurse and do find path-compression
   next = find(nextNext);
   pData->setFind(next);
   return next;
@@ -59,23 +64,49 @@ void SharedTermManager::addTerm(TNode n, Theory* parent, Theory* child) {
   // theory id
   uint64_t parentTag = (1 << parent->getId());
   uint64_t childTag = (1 << child->getId());
-  uint64_t newTags = parentTag | childTag;
+  uint64_t bothTags = parentTag | childTag;
+
+  // Create or update the SharedData for n
   SharedData* pData;
   if(n.getAttribute(SharedAttr(), pData)) {
     // The attribute already exists, just update it if necessary
     uint64_t tags = pData->getTheories();
-    newTags |= tags;
+    uint64_t newTags = tags | bothTags;
     if (tags == newTags) return;
+
+    // Get the equivalence class representative -- if this is different from
+    // the current node, then we will need to notify the theories of that and
+    // update the theory tags on the representative node
+    SharedData* pDataFind = find(pData);
+
+    // Notify parent theory if necessary
     if (!(tags & parentTag)) {
       parent->addSharedTerm(n);
+      if (pData != pDataFind) {
+        parent->notifyEq(n, pDataFind->getRep());
+      }
     }
+
+    // Notify child theory if necessary
     if (!(tags & childTag)) {
       child->addSharedTerm(n);
+      if (pData != pDataFind) {
+        child->notifyEq(n, pDataFind->getRep());
+      }
     }
+
+    // Update theory tags
     pData->setTheories(newTags);
+    if (pData != pDataFind) {
+      tags = pDataFind->getTheories();
+      newTags = tags | bothTags;
+      if (tags != newTags) {
+        pDataFind->setTheories(newTags);
+      }
+    }
   } else {
     // The attribute does not exist, so it is created and set
-    pData = new (true) SharedData(d_context, n, newTags);
+    pData = new (true) SharedData(d_context, n, bothTags);
     n.setAttribute(SharedAttr(), pData);
     parent->addSharedTerm(n);
     child->addSharedTerm(n);
@@ -83,22 +114,35 @@ void SharedTermManager::addTerm(TNode n, Theory* parent, Theory* child) {
 }
 
 
-void SharedTermManager::addEq(Theory* thReason, TNode eq) {
-  Assert(eq.getKind() == kind::EQUAL &&
-         eq[0].hasAttribute(SharedAttr()) && eq[1].hasAttribute(SharedAttr()),
+void SharedTermManager::addEq(TNode eq, Theory* thReason) {
+  Assert(eq.getKind() == kind::EQUAL,
          "SharedTermManager::addEq precondition violated");
 
   TNode x = eq[0];
   TNode y = eq[1];
 
-  SharedData* pDataX = x.getAttribute(SharedAttr());
-  SharedData* pDataY = y.getAttribute(SharedAttr());
+  SharedData* pDataX;
+  SharedData* pDataY;
+
+  // Grab the SharedData for each side of the equality, create if necessary
+  if(!x.getAttribute(SharedAttr(), pDataX)) {
+    pDataX = new (true) SharedData(d_context, x, 0);
+    x.setAttribute(SharedAttr(), pDataX);
+  }
+  if(!y.getAttribute(SharedAttr(), pDataY)) {
+    pDataY = new (true) SharedData(d_context, y, 0);
+    y.setAttribute(SharedAttr(), pDataY);
+  }
 
+  // Get the representatives
   SharedData* pDataXX = find(pDataX);
   SharedData* pDataYY = find(pDataY);
 
+  // If already in the same equivalence class, nothing to do
   if(pDataXX == pDataYY) return;
 
+  // Make sure we reverse the edges in the smaller tree
+  bool switched = false;
   if(pDataXX->getSize() > pDataYY->getSize()) {
     SharedData* tmp = pDataXX;
     pDataXX = pDataYY;
@@ -107,17 +151,27 @@ void SharedTermManager::addEq(Theory* thReason, TNode eq) {
     tmp = pDataX;
     pDataX = pDataY;
     pDataY = tmp;
+
+    switched = true;
   }
 
-  pDataX->reverseEdges();
-  pDataX->setEquality(eq);
+  // Reverse the edges in the left proof tree and link the two trees
+  if(!pDataX->isProofRoot()) {
+    pDataX->reverseEdges();
+  }
+  pDataX->setEdgeReversed(switched);
   pDataX->setExplainingTheory(thReason);
   pDataX->setProofNext(pDataY);
 
+  // Set the equivalence class representative for the left node to be the right node
   pDataXX->setFind(pDataYY);
   pDataYY->setSize(pDataYY->getSize() + pDataXX->getSize());
+
+  // Update the theory tags for the new representative
   uint64_t tags = pDataXX->getTheories();
   pDataYY->setTheories(pDataYY->getTheories() | tags);
+
+  // Notify the theories that care about the left node
   int id = 0;
   while (tags != 0) {
     if (tags & 1) {
@@ -129,8 +183,29 @@ void SharedTermManager::addEq(Theory* thReason, TNode eq) {
 }
 
 
+void SharedTermManager::collectExplanations(SharedData* pData, set<Node>& s) const {
+  Theory* expTh = pData->getExplainingTheory();
+  if(expTh == NULL) {
+    // This equality is directly from SAT, no need to ask a theory for an explanation
+    s.insert(pData->getEquality());
+  }
+  else {
+    // This equality was sent by a theory - ask the theory for the explanation
+    Node n = d_engine->getExplanation(pData->getEquality(), expTh);
+    if(n.getKind() == kind::AND) {
+      for (Node::iterator it = n.begin(), iend = n.end(); it != iend; ++it) {
+        s.insert(*it);
+      }
+    }
+    else {
+      s.insert(n);
+    }
+  }
+}
+
+
 // Explain this equality
-Node SharedTermManager::explain(TNode eq) {
+Node SharedTermManager::explain(TNode eq) const {
   Assert(eq.getKind() == kind::EQUAL &&
          eq[0].hasAttribute(SharedAttr()) && eq[1].hasAttribute(SharedAttr()),
          "SharedTermManager::explain precondition violated");
@@ -138,14 +213,15 @@ Node SharedTermManager::explain(TNode eq) {
   TNode x = eq[0];
   TNode y = eq[1];
 
+  // Get the SharedData for both sides of the equality
   SharedData* pDataX = x.getAttribute(SharedAttr());
   SharedData* pDataY = y.getAttribute(SharedAttr());
 
   Assert(find(pDataX) == find(pDataY),
          "invalid equality input to SharedTermManager::explain");
 
+  // Find the path between the two nodes and build up the explanation
   std::set<Node> s;
-  Node n;
   SharedData* tmp = pDataX;
 
   // Check to see if Y is on path from X to root
@@ -155,39 +231,27 @@ Node SharedTermManager::explain(TNode eq) {
   if (tmp == pDataY) {
     // Y is on path from X to root, just traverse that path
     while (pDataX != pDataY) {
-      n = d_engine->getExplanation(pDataX->getEquality(),
-                                   pDataX->getExplainingTheory());
-      for (Node::iterator it = n.begin(), iend = n.end(); it != iend; ++it) {
-        s.insert(*it);
-      }
+      collectExplanations(pDataX, s);
       pDataX = pDataX->getProofNext();
     }
   }
   else {
     // Y is not on path from X to root, so traverse from Y to root instead
     while (!pDataY->isProofRoot() && pDataX != pDataY) {
-      n = d_engine->getExplanation(pDataY->getEquality(),
-                                   pDataY->getExplainingTheory());
-      for (Node::iterator it = n.begin(), iend = n.end(); it != iend; ++it) {
-        s.insert(*it);
-      }
+      collectExplanations(pDataY, s);
       pDataY = pDataY->getProofNext();
     }
   }
   if (pDataX != pDataY) {
     // X and Y are on different branches - have to traverse both
     while (!pDataX->isProofRoot()) {
-      n = d_engine->getExplanation(pDataX->getEquality(),
-                                   pDataX->getExplainingTheory());
-      for (Node::iterator it = n.begin(), iend = n.end(); it != iend; ++it) {
-        s.insert(*it);
-      }
+      collectExplanations(pDataX, s);
       pDataX = pDataX->getProofNext();
     }
   }
 
   // Build explanation from s
-  NodeBuilder<kind::AND> nb;
+  NodeBuilder<> nb(kind::AND);
   set<Node>::iterator it = s.begin(), iend = s.end();
   for (; it != iend; ++it) {
     nb << *it;
@@ -196,4 +260,13 @@ Node SharedTermManager::explain(TNode eq) {
 }
 
 
+Node SharedTermManager::getRep(TNode n) const {
+  Assert(n.hasAttribute(SharedAttr()),
+         "SharedTermManager::getRep precondition violated");
+
+  SharedData* pData = n.getAttribute(SharedAttr());
+  return find(pData)->getRep();
+}
+
+
 }/* CVC4 namespace */