1 files changed, 121 insertions, 4 deletions
diff --git a/src/theory/sets/solver_state.cpp b/src/theory/sets/solver_state.cpp
index 941f59bc6..79c7bc1c8 100644
--- a/src/theory/sets/solver_state.cpp
+++ b/src/theory/sets/solver_state.cpp
@@ -29,14 +29,12 @@ SolverState::SolverState(context::Context* c,
                          context::UserContext* u,
                          Valuation val,
                          SkolemCache& skc)
-    : TheoryState(c, u, val), d_skCache(skc), d_parent(nullptr)
+    : TheoryState(c, u, val), d_skCache(skc), d_members(c)
 {
   d_true = NodeManager::currentNM()->mkConst(true);
   d_false = NodeManager::currentNM()->mkConst(false);
 }
 
-void SolverState::setParent(TheorySetsPrivate* p) { d_parent = p; }
-
 void SolverState::reset()
 {
   d_set_eqc.clear();
@@ -249,7 +247,7 @@ bool SolverState::isEntailed(Node n, bool polarity) const
     if (polarity && d_ee->hasTerm(n[1]))
     {
       Node r = d_ee->getRepresentative(n[1]);
-      if (d_parent->isMember(n[0], r))
+      if (isMember(n[0], r))
       {
         return true;
       }
@@ -469,6 +467,125 @@ const vector<Node> SolverState::getSetsEqClasses(const TypeNode& t) const
   return representatives;
 }
 
+bool SolverState::isMember(TNode x, TNode s) const
+{
+  Assert(hasTerm(s) && getRepresentative(s) == s);
+  NodeIntMap::const_iterator mem_i = d_members.find(s);
+  if (mem_i != d_members.end())
+  {
+    std::map<Node, std::vector<Node> >::const_iterator itd =
+        d_members_data.find(s);
+    Assert(itd != d_members_data.end());
+    const std::vector<Node>& members = itd->second;
+    Assert((*mem_i).second <= members.size());
+    for (size_t i = 0, nmem = (*mem_i).second; i < nmem; i++)
+    {
+      if (areEqual(members[i][0], x))
+      {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+void SolverState::addMember(TNode r, TNode atom)
+{
+  NodeIntMap::iterator mem_i = d_members.find(r);
+  size_t n_members = 0;
+  if (mem_i != d_members.end())
+  {
+    n_members = (*mem_i).second;
+  }
+  d_members[r] = n_members + 1;
+  if (n_members < d_members_data[r].size())
+  {
+    d_members_data[r][n_members] = atom;
+  }
+  else
+  {
+    d_members_data[r].push_back(atom);
+  }
+}
+
+bool SolverState::merge(TNode t1,
+                        TNode t2,
+                        std::vector<Node>& facts,
+                        TNode cset)
+{
+  NodeIntMap::iterator mem_i2 = d_members.find(t2);
+  if (mem_i2 == d_members.end())
+  {
+    // no members in t2, we are done
+    return true;
+  }
+  NodeIntMap::iterator mem_i1 = d_members.find(t1);
+  size_t n_members = 0;
+  if (mem_i1 != d_members.end())
+  {
+    n_members = (*mem_i1).second;
+  }
+  for (size_t i = 0, nmem2 = (*mem_i2).second; i < nmem2; i++)
+  {
+    Assert(i < d_members_data[t2].size()
+           && d_members_data[t2][i].getKind() == MEMBER);
+    Node m2 = d_members_data[t2][i];
+    // check if redundant
+    bool add = true;
+    for (size_t j = 0; j < n_members; j++)
+    {
+      Assert(j < d_members_data[t1].size()
+             && d_members_data[t1][j].getKind() == MEMBER);
+      if (areEqual(m2[0], d_members_data[t1][j][0]))
+      {
+        add = false;
+        break;
+      }
+    }
+    if (add)
+    {
+      // if there is a concrete set in t1, propagate new facts or conflicts
+      if (!cset.isNull())
+      {
+        NodeManager* nm = NodeManager::currentNM();
+        Assert(areEqual(m2[1], cset));
+        Node exp = nm->mkNode(AND, m2[1].eqNode(cset), m2);
+        if (cset.getKind() == SINGLETON)
+        {
+          if (cset[0] != m2[0])
+          {
+            Node eq = cset[0].eqNode(m2[0]);
+            Trace("sets-prop") << "Propagate eq-mem eq inference : " << exp
+                               << " => " << eq << std::endl;
+            Node fact = nm->mkNode(IMPLIES, exp, eq);
+            facts.push_back(fact);
+          }
+        }
+        else
+        {
+          // conflict
+          Assert(facts.empty());
+          Trace("sets-prop")
+              << "Propagate eq-mem conflict : " << exp << std::endl;
+          facts.push_back(exp);
+          return false;
+        }
+      }
+      if (n_members < d_members_data[t1].size())
+      {
+        d_members_data[t1][n_members] = m2;
+      }
+      else
+      {
+        d_members_data[t1].push_back(m2);
+      }
+      n_members++;
+    }
+  }
+  d_members[t1] = n_members;
+  return true;
+}
+
 }  // namespace sets
 }  // namespace theory
 }  // namespace CVC4