Add support for set comprehension (#3312)

8 files changed, 168 insertions, 41 deletions

diff --git a/src/theory/sets/kinds b/src/theory/sets/kinds
index 34fef7d64..058a20aeb 100644
--- a/src/theory/sets/kinds
+++ b/src/theory/sets/kinds
@@ -46,6 +46,21 @@ operator CARD          1  "set cardinality operator"
 operator COMPLEMENT    1  "set COMPLEMENT (with respect to finite universe)"
 nullaryoperator UNIVERSE_SET "(finite) universe set, all set variables must be interpreted as subsets of it."
 
+# A set comprehension is specified by:
+# (1) a bound variable list x1 ... xn,
+# (2) a predicate P[x1...xn], and
+# (3) a term t[x1...xn].
+# A comprehension C with the above form has members given by the following
+# semantics:
+# forall y. ( exists x1...xn. P[x1...xn] ^ t[x1...xn] = y ) <=> (member y C)
+# where y ranges over the element type of the (set) type of the comprehension.
+# Notice that since all sets must be interpreted as finite, this means that
+# CVC4 will not be able to construct a model for any set comprehension such
+# that there are infinitely many y that satisfy the left hand side of the
+# equivalence above. The same limitation occurs more generally when combining
+# finite sets with quantified formulas.
+operator COMPREHENSION 3 "set comprehension specified by a bound variable list, a predicate, and a term."
+
 operator JOIN 		   2  "set join"
 operator PRODUCT 	   2  "set cartesian product"
 operator TRANSPOSE 	   1  "set transpose"
@@ -64,6 +79,7 @@ typerule INSERT         ::CVC4::theory::sets::InsertTypeRule
 typerule CARD           ::CVC4::theory::sets::CardTypeRule
 typerule COMPLEMENT     ::CVC4::theory::sets::ComplementTypeRule
 typerule UNIVERSE_SET   ::CVC4::theory::sets::UniverseSetTypeRule
+typerule COMPREHENSION  ::CVC4::theory::sets::ComprehensionTypeRule
 
 typerule JOIN 			::CVC4::theory::sets::RelBinaryOperatorTypeRule
 typerule PRODUCT 		::CVC4::theory::sets::RelBinaryOperatorTypeRule
diff --git a/src/theory/sets/solver_state.cpp b/src/theory/sets/solver_state.cpp
index 7756fe081..ca7c0bf2b 100644
--- a/src/theory/sets/solver_state.cpp
+++ b/src/theory/sets/solver_state.cpp
@@ -44,12 +44,14 @@ void SolverState::reset()
   d_congruent.clear();
   d_nvar_sets.clear();
   d_var_set.clear();
+  d_compSets.clear();
   d_pol_mems[0].clear();
   d_pol_mems[1].clear();
   d_members_index.clear();
   d_singleton_index.clear();
   d_bop_index.clear();
   d_op_list.clear();
+  d_allCompSets.clear();
 }
 
 void SolverState::registerEqc(TypeNode tn, Node r)
@@ -137,6 +139,12 @@ void SolverState::registerTerm(Node r, TypeNode tnn, Node n)
     d_nvar_sets[r].push_back(n);
     Trace("sets-debug2") << "Non-var-set[" << r << "] : " << n << std::endl;
   }
+  else if (nk == COMPREHENSION)
+  {
+    d_compSets[r].push_back(n);
+    d_allCompSets.push_back(n);
+    Trace("sets-debug2") << "Comp-set[" << r << "] : " << n << std::endl;
+  }
   else if (n.isVar() && !d_skCache.isSkolem(n))
   {
     // it is important that we check it is a variable, but not an internally
@@ -146,9 +154,14 @@ void SolverState::registerTerm(Node r, TypeNode tnn, Node n)
       if (d_var_set.find(r) == d_var_set.end())
       {
         d_var_set[r] = n;
+        Trace("sets-debug2") << "var-set[" << r << "] : " << n << std::endl;
       }
     }
   }
+  else
+  {
+    Trace("sets-debug2") << "Unknown-set[" << r << "] : " << n << std::endl;
+  }
 }
 
 bool SolverState::areEqual(Node a, Node b) const
@@ -510,6 +523,17 @@ Node SolverState::getVariableSet(Node r) const
   }
   return Node::null();
 }
+
+const std::vector<Node>& SolverState::getComprehensionSets(Node r) const
+{
+  std::map<Node, std::vector<Node> >::const_iterator it = d_compSets.find(r);
+  if (it == d_compSets.end())
+  {
+    return d_emptyVec;
+  }
+  return it->second;
+}
+
 const std::map<Node, Node>& SolverState::getMembers(Node r) const
 {
   return getMembersInternal(r, 0);
@@ -550,6 +574,11 @@ const std::map<Kind, std::vector<Node> >& SolverState::getOperatorList() const
   return d_op_list;
 }
 
+const std::vector<Node>& SolverState::getComprehensionSets() const
+{
+  return d_allCompSets;
+}
+
 void SolverState::debugPrintSet(Node s, const char* c) const
 {
   if (s.getNumChildren() == 0)
diff --git a/src/theory/sets/solver_state.h b/src/theory/sets/solver_state.h
index 2b4d93ed3..7426a701b 100644
--- a/src/theory/sets/solver_state.h
+++ b/src/theory/sets/solver_state.h
@@ -133,6 +133,8 @@ class SolverState
    * if none exist.
    */
   Node getVariableSet(Node r) const;
+  /** Get comprehension sets in equivalence class with representative r */
+  const std::vector<Node>& getComprehensionSets(Node r) const;
   /** Get (positive) members of the set equivalence class r
    *
    * The members are return as a map, which maps members to their explanation.
@@ -161,6 +163,8 @@ class SolverState
    * map is a representative of its congruence class.
    */
   const std::map<Kind, std::vector<Node> >& getOperatorList() const;
+  /** Get the list of all comprehension sets in the current context */
+  const std::vector<Node>& getComprehensionSets() const;
 
   // --------------------------------------- commonly used terms
   /** Get type constraint skolem
@@ -225,6 +229,8 @@ class SolverState
   std::map<Node, Node> d_congruent;
   /** Map from equivalence classes to the list of non-variable sets in it */
   std::map<Node, std::vector<Node> > d_nvar_sets;
+  /** Map from equivalence classes to the list of comprehension sets in it */
+  std::map<Node, std::vector<Node> > d_compSets;
   /** Map from equivalence classes to a variable sets in it */
   std::map<Node, Node> d_var_set;
   /** polarity memberships
@@ -262,6 +268,8 @@ class SolverState
   std::map<Node, Node> d_singleton_index;
   /** Indices for the binary kinds INTERSECT, SETMINUS and UNION. */
   std::map<Kind, std::map<Node, std::map<Node, Node> > > d_bop_index;
+  /** A list of comprehension sets */
+  std::vector<Node> d_allCompSets;
   // -------------------------------- end term indices
   std::map<Kind, std::vector<Node> > d_op_list;
   /** the skolem cache */
diff --git a/src/theory/sets/theory_sets.cpp b/src/theory/sets/theory_sets.cpp
index 869cb8926..3b0427ec5 100644
--- a/src/theory/sets/theory_sets.cpp
+++ b/src/theory/sets/theory_sets.cpp
@@ -15,7 +15,11 @@
  **/
 
 #include "theory/sets/theory_sets.h"
+#include "options/sets_options.h"
 #include "theory/sets/theory_sets_private.h"
+#include "theory/theory_model.h"
+
+using namespace CVC4::kind;
 
 namespace CVC4 {
 namespace theory {
@@ -40,6 +44,13 @@ TheorySets::~TheorySets()
   // Do not move me to the header. See explanation in the constructor.
 }
 
+void TheorySets::finishInit()
+{
+  TheoryModel* tm = d_valuation.getModel();
+  Assert(tm != nullptr);
+  tm->setUnevaluatedKind(COMPREHENSION);
+}
+
 void TheorySets::addSharedTerm(TNode n) {
   d_internal->addSharedTerm(n);
 }
@@ -78,6 +89,28 @@ void TheorySets::preRegisterTerm(TNode node) {
 }
 
 Node TheorySets::expandDefinition(LogicRequest &logicRequest, Node n) {
+  Kind nk = n.getKind();
+  if (nk == UNIVERSE_SET || nk == COMPLEMENT || nk == JOIN_IMAGE
+      || nk == COMPREHENSION)
+  {
+    if (!options::setsExt())
+    {
+      std::stringstream ss;
+      ss << "Extended set operators are not supported in default mode, try "
+            "--sets-ext.";
+      throw LogicException(ss.str());
+    }
+  }
+  if (nk == COMPREHENSION)
+  {
+    // set comprehension is an implicit quantifier, require it in the logic
+    if (!getLogicInfo().isQuantified())
+    {
+      std::stringstream ss;
+      ss << "Set comprehensions require quantifiers in the background logic.";
+      throw LogicException(ss.str());
+    }
+  }
   return d_internal->expandDefinition(logicRequest, n);
 }
 
diff --git a/src/theory/sets/theory_sets.h b/src/theory/sets/theory_sets.h
index 4c145aedb..ed2459b32 100644
--- a/src/theory/sets/theory_sets.h
+++ b/src/theory/sets/theory_sets.h
@@ -42,6 +42,8 @@ class TheorySets : public Theory
              const LogicInfo& logicInfo);
   ~TheorySets() override;
 
+  /** finish initialization */
+  void finishInit() override;
   void addSharedTerm(TNode) override;
   void check(Effort) override;
   bool collectModelInfo(TheoryModel* m) override;
diff --git a/src/theory/sets/theory_sets_private.cpp b/src/theory/sets/theory_sets_private.cpp
index 94fef85f5..6a381d471 100644
--- a/src/theory/sets/theory_sets_private.cpp
+++ b/src/theory/sets/theory_sets_private.cpp
@@ -38,7 +38,7 @@ TheorySetsPrivate::TheorySetsPrivate(TheorySets& external,
                                      context::UserContext* u)
     : d_members(c),
       d_deq(c),
-      d_deq_processed(u),
+      d_termProcessed(u),
       d_keep(c),
       d_full_check_incomplete(false),
       d_external(external),
@@ -432,10 +432,16 @@ void TheorySetsPrivate::fullEffortCheck(){
           {
             checkDisequalities();
             d_im.flushPendingLemmas();
-            if (!d_im.hasProcessed() && d_card_enabled)
+            if (!d_im.hasProcessed())
             {
-              // call the check method of the cardinality solver
-              d_cardSolver->check();
+              checkReduceComprehensions();
+              d_im.flushPendingLemmas();
+
+              if (!d_im.hasProcessed() && d_card_enabled)
+              {
+                // call the check method of the cardinality solver
+                d_cardSolver->check();
+              }
             }
           }
         }
@@ -763,13 +769,13 @@ void TheorySetsPrivate::checkDisequalities()
       // already satisfied
       continue;
     }
-    if (d_deq_processed.find(deq) != d_deq_processed.end())
+    if (d_termProcessed.find(deq) != d_termProcessed.end())
     {
       // already added lemma
       continue;
     }
-    d_deq_processed.insert(deq);
-    d_deq_processed.insert(deq[1].eqNode(deq[0]));
+    d_termProcessed.insert(deq);
+    d_termProcessed.insert(deq[1].eqNode(deq[0]));
     Trace("sets") << "Process Disequality : " << deq.negate() << std::endl;
     TypeNode elementType = deq[0].getType().getSetElementType();
     Node x = d_state.getSkolemCache().mkTypedSkolemCached(
@@ -787,6 +793,41 @@ void TheorySetsPrivate::checkDisequalities()
   }
 }
 
+void TheorySetsPrivate::checkReduceComprehensions()
+{
+  NodeManager* nm = NodeManager::currentNM();
+  const std::vector<Node>& comps = d_state.getComprehensionSets();
+  for (const Node& n : comps)
+  {
+    if (d_termProcessed.find(n) != d_termProcessed.end())
+    {
+      // already reduced it
+      continue;
+    }
+    d_termProcessed.insert(n);
+    Node v = nm->mkBoundVar(n[2].getType());
+    Node body = nm->mkNode(AND, n[1], v.eqNode(n[2]));
+    // must do substitution
+    std::vector<Node> vars;
+    std::vector<Node> subs;
+    for (const Node& cv : n[0])
+    {
+      vars.push_back(cv);
+      Node cvs = nm->mkBoundVar(cv.getType());
+      subs.push_back(cvs);
+    }
+    body = body.substitute(vars.begin(), vars.end(), subs.begin(), subs.end());
+    Node bvl = nm->mkNode(BOUND_VAR_LIST, subs);
+    body = nm->mkNode(EXISTS, bvl, body);
+    Node mem = nm->mkNode(MEMBER, v, n);
+    Node lem =
+        nm->mkNode(FORALL, nm->mkNode(BOUND_VAR_LIST, v), body.eqNode(mem));
+    Trace("sets-comprehension")
+        << "Comprehension reduction: " << lem << std::endl;
+    d_im.flushLemma(lem);
+  }
+}
+
 /**************************** TheorySetsPrivate *****************************/
 /**************************** TheorySetsPrivate *****************************/
 /**************************** TheorySetsPrivate *****************************/
@@ -1190,35 +1231,22 @@ void TheorySetsPrivate::preRegisterTerm(TNode node)
                 << std::endl;
   switch(node.getKind()) {
   case kind::EQUAL:
-    // TODO: what's the point of this
     d_equalityEngine.addTriggerEquality(node);
     break;
   case kind::MEMBER:
-    // TODO: what's the point of this
     d_equalityEngine.addTriggerPredicate(node);
     break;
   case kind::CARD:
     d_equalityEngine.addTriggerTerm(node, THEORY_SETS);
     break;
   default:
-    //if( node.getType().isSet() ){
-    //  d_equalityEngine.addTriggerTerm(node, THEORY_SETS);
-    //}else{
     d_equalityEngine.addTerm(node);
-    //}
     break;
   }
 }
 
 Node TheorySetsPrivate::expandDefinition(LogicRequest &logicRequest, Node n) {
   Debug("sets-proc") << "expandDefinition : " << n << std::endl;
-  if( n.getKind()==kind::UNIVERSE_SET || n.getKind()==kind::COMPLEMENT || n.getKind()==kind::JOIN_IMAGE ){
-    if( !options::setsExt() ){
-      std::stringstream ss;
-      ss << "Extended set operators are not supported in default mode, try --sets-ext.";
-      throw LogicException(ss.str());
-    }
-  }
   return n;
 }
 
diff --git a/src/theory/sets/theory_sets_private.h b/src/theory/sets/theory_sets_private.h
index 0468e5fb7..5ef8c4825 100644
--- a/src/theory/sets/theory_sets_private.h
+++ b/src/theory/sets/theory_sets_private.h
@@ -101,6 +101,11 @@ class TheorySetsPrivate {
    * roughly corresponds the SET DISEQUALITY rule from Bansal et al IJCAR 2016.
    */
   void checkDisequalities();
+  /**
+   * Check comprehensions. This adds reduction lemmas for all set comprehensions
+   * in the current context.
+   */
+  void checkReduceComprehensions();
 
   void addCarePairs(TNodeTrie* t1,
                     TNodeTrie* t2,
@@ -112,7 +117,11 @@ class TheorySetsPrivate {
   Node d_false;
   Node d_zero;
   NodeBoolMap d_deq;
-  NodeSet d_deq_processed;
+  /**
+   * The set of terms that we have reduced via a lemma in the current user
+   * context.
+   */
+  NodeSet d_termProcessed;
   NodeSet d_keep;
   std::vector< Node > d_emp_exp;
   
diff --git a/src/theory/sets/theory_sets_type_rules.h b/src/theory/sets/theory_sets_type_rules.h
index 27aa58452..430163121 100644
--- a/src/theory/sets/theory_sets_type_rules.h
+++ b/src/theory/sets/theory_sets_type_rules.h
@@ -25,26 +25,6 @@ namespace CVC4 {
 namespace theory {
 namespace sets {
 
-class SetsTypeRule {
-public:
-
-  /**
-   * Compute the type for (and optionally typecheck) a term belonging
-   * to the theory of sets.
-   *
-   * @param check if true, the node's type should be checked as well
-   * as computed.
-   */
-  inline static TypeNode computeType(NodeManager* nodeManager, TNode n,
-                                     bool check)
-  {
-    // TODO: implement me!
-    Unimplemented();
-
-  }
-
-};/* class SetsTypeRule */
-
 struct SetsBinaryOperatorTypeRule {
   inline static TypeNode computeType(NodeManager* nodeManager, TNode n, bool check)
   {
@@ -216,6 +196,28 @@ struct UniverseSetTypeRule {
   }
 };/* struct ComplementTypeRule */
 
+struct ComprehensionTypeRule
+{
+  static TypeNode computeType(NodeManager* nodeManager, TNode n, bool check)
+  {
+    Assert(n.getKind() == kind::COMPREHENSION);
+    if (check)
+    {
+      if (n[0].getType(check) != nodeManager->boundVarListType())
+      {
+        throw TypeCheckingExceptionPrivate(
+            n, "first argument of set comprehension is not bound var list");
+      }
+      if (n[1].getType(check) != nodeManager->booleanType())
+      {
+        throw TypeCheckingExceptionPrivate(
+            n, "body of set comprehension is not boolean");
+      }
+    }
+    return nodeManager->mkSetType(n[2].getType(check));
+  }
+}; /* struct ComprehensionTypeRule */
+
 struct InsertTypeRule {
   inline static TypeNode computeType(NodeManager* nodeManager, TNode n, bool check)
   {