Add operator set.map to theory of sets (#7641)

Towards supporting set.map operator in the sets solver.

11 files changed, 220 insertions, 6 deletions

diff --git a/src/api/cpp/cvc5.cpp b/src/api/cpp/cvc5.cpp
index e1a2a547c..0c7e86c86 100644
--- a/src/api/cpp/cvc5.cpp
+++ b/src/api/cpp/cvc5.cpp
@@ -288,6 +288,7 @@ const static std::unordered_map<Kind, cvc5::Kind> s_kinds{
     {SET_COMPREHENSION, cvc5::Kind::SET_COMPREHENSION},
     {SET_CHOOSE, cvc5::Kind::SET_CHOOSE},
     {SET_IS_SINGLETON, cvc5::Kind::SET_IS_SINGLETON},
+    {SET_MAP, cvc5::Kind::SET_MAP},
     /* Relations ----------------------------------------------------------- */
     {RELATION_JOIN, cvc5::Kind::RELATION_JOIN},
     {RELATION_PRODUCT, cvc5::Kind::RELATION_PRODUCT},
@@ -598,6 +599,7 @@ const static std::unordered_map<cvc5::Kind, Kind, cvc5::kind::KindHashFunction>
         {cvc5::Kind::SET_COMPREHENSION, SET_COMPREHENSION},
         {cvc5::Kind::SET_CHOOSE, SET_CHOOSE},
         {cvc5::Kind::SET_IS_SINGLETON, SET_IS_SINGLETON},
+        {cvc5::Kind::SET_MAP, SET_MAP},
         /* Relations ------------------------------------------------------- */
         {cvc5::Kind::RELATION_JOIN, RELATION_JOIN},
         {cvc5::Kind::RELATION_PRODUCT, RELATION_PRODUCT},
diff --git a/src/api/cpp/cvc5_kind.h b/src/api/cpp/cvc5_kind.h
index 8da58ebd6..29cb130d4 100644
--- a/src/api/cpp/cvc5_kind.h
+++ b/src/api/cpp/cvc5_kind.h
@@ -2252,9 +2252,9 @@ enum Kind : int32_t
   SET_COMPREHENSION,
   /**
    * Returns an element from a given set.
-   * If a set A = {x}, then the term (choose A) is equivalent to the term x.
-   * If the set is empty, then (choose A) is an arbitrary value.
-   * If the set has cardinality > 1, then (choose A) will deterministically
+   * If a set A = {x}, then the term (set.choose A) is equivalent to the term x.
+   * If the set is empty, then (set.choose A) is an arbitrary value.
+   * If the set has cardinality > 1, then (set.choose A) will deterministically
    * return an element in A.
    *
    * Parameters:
@@ -2274,6 +2274,21 @@ enum Kind : int32_t
    *   - `Solver::mkTerm(Kind kind, const Term& child) const`
    */
   SET_IS_SINGLETON,
+  /**
+   * set.map operator applies the first argument, a function of type (-> T1 T2),
+   * to every element of the second argument, a set of type (Set T1),
+   * and returns a set of type (Set T2).
+   *
+   * Parameters:
+   *   - 1: a function of type (-> T1 T2)
+   *   - 2: a set of type (Set T1)
+   *
+   * Create with:
+   *   - `Solver::mkTerm(Kind kind, const Term& child1, const Term& child2)
+   * const`
+   *   - `Solver::mkTerm(Kind kind, const std::vector<Term>& children) const`
+   */
+   SET_MAP,
 
   /* Relations ------------------------------------------------------------- */
 
diff --git a/src/parser/smt2/smt2.cpp b/src/parser/smt2/smt2.cpp
index 30dcd06cc..3c423674a 100644
--- a/src/parser/smt2/smt2.cpp
+++ b/src/parser/smt2/smt2.cpp
@@ -602,6 +602,7 @@ Command* Smt2::setLogic(std::string name, bool fromCommand)
     addOperator(api::SET_COMPLEMENT, "set.complement");
     addOperator(api::SET_CHOOSE, "set.choose");
     addOperator(api::SET_IS_SINGLETON, "set.is_singleton");
+    addOperator(api::SET_MAP, "set.map");
     addOperator(api::RELATION_JOIN, "rel.join");
     addOperator(api::RELATION_PRODUCT, "rel.product");
     addOperator(api::RELATION_TRANSPOSE, "rel.transpose");
diff --git a/src/printer/smt2/smt2_printer.cpp b/src/printer/smt2/smt2_printer.cpp
index 07ad0b2c2..30433ec1a 100644
--- a/src/printer/smt2/smt2_printer.cpp
+++ b/src/printer/smt2/smt2_printer.cpp
@@ -1073,6 +1073,7 @@ std::string Smt2Printer::smtKindString(Kind k, Variant v)
   case kind::SET_COMPREHENSION: return "set.comprehension";
   case kind::SET_CHOOSE: return "set.choose";
   case kind::SET_IS_SINGLETON: return "set.is_singleton";
+  case kind::SET_MAP: return "set.map";
   case kind::RELATION_JOIN: return "rel.join";
   case kind::RELATION_PRODUCT: return "rel.product";
   case kind::RELATION_TRANSPOSE: return "rel.transpose";
diff --git a/src/theory/sets/kinds b/src/theory/sets/kinds
index 71d5fce4a..0c1cc0452 100644
--- a/src/theory/sets/kinds
+++ b/src/theory/sets/kinds
@@ -80,11 +80,15 @@ operator SET_CHOOSE        1  "return an element in the set given as a parameter
 # The operator is_singleton returns whether the given set is a singleton
 operator SET_IS_SINGLETON  1  "return whether the given set is a singleton"
 
+# The set.map operator applies the first argument, a function of type (-> T1 T2), to every element
+# of the second argument, a set of type (Set T1), and returns a set of type (Set T2).
+operator SET_MAP           2  "set map function"
+
 operator RELATION_JOIN 		   2  "relation join"
 operator RELATION_PRODUCT 	   2  "relation cartesian product"
-operator RELATION_TRANSPOSE 	   1  "relation transpose"
-operator RELATION_TCLOSURE      1  "relation transitive closure"
-operator RELATION_JOIN_IMAGE	   2  "relation join image"
+operator RELATION_TRANSPOSE    1  "relation transpose"
+operator RELATION_TCLOSURE     1  "relation transitive closure"
+operator RELATION_JOIN_IMAGE   2  "relation join image"
 operator RELATION_IDEN 	   	   1  "relation identity"
 
 typerule SET_UNION          ::cvc5::theory::sets::SetsBinaryOperatorTypeRule
@@ -102,6 +106,7 @@ typerule SET_UNIVERSE       ::cvc5::theory::sets::UniverseSetTypeRule
 typerule SET_COMPREHENSION  ::cvc5::theory::sets::ComprehensionTypeRule
 typerule SET_CHOOSE         ::cvc5::theory::sets::ChooseTypeRule
 typerule SET_IS_SINGLETON   ::cvc5::theory::sets::IsSingletonTypeRule
+typerule SET_MAP            ::cvc5::theory::sets::SetMapTypeRule
 
 typerule RELATION_JOIN 			::cvc5::theory::sets::RelBinaryOperatorTypeRule
 typerule RELATION_PRODUCT 		::cvc5::theory::sets::RelBinaryOperatorTypeRule
diff --git a/src/theory/sets/theory_sets_rewriter.cpp b/src/theory/sets/theory_sets_rewriter.cpp
index d63368f19..d603946c4 100644
--- a/src/theory/sets/theory_sets_rewriter.cpp
+++ b/src/theory/sets/theory_sets_rewriter.cpp
@@ -330,6 +330,8 @@ RewriteResponse TheorySetsRewriter::postRewrite(TNode node) {
     break;
   }  // kind::SET_IS_SINGLETON
 
+  case SET_MAP: return postRewriteMap(node);
+
   case kind::RELATION_TRANSPOSE:
   {
     if (node[0].getKind() == kind::RELATION_TRANSPOSE)
@@ -630,6 +632,40 @@ RewriteResponse TheorySetsRewriter::preRewrite(TNode node) {
   return RewriteResponse(REWRITE_DONE, node);
 }
 
+RewriteResponse TheorySetsRewriter::postRewriteMap(TNode n)
+{
+  Assert(n.getKind() == kind::SET_MAP);
+  NodeManager* nm = NodeManager::currentNM();
+  Kind k = n[1].getKind();
+  TypeNode rangeType = n[0].getType().getRangeType();
+  switch (k)
+  {
+    case SET_EMPTY:
+    {
+      // (set.map f (as set.empty (Set T1)) = (as set.empty (Set T2))
+      Node ret = nm->mkConst(EmptySet(nm->mkSetType(rangeType)));
+      return RewriteResponse(REWRITE_DONE, ret);
+    }
+    case SET_SINGLETON:
+    {
+      // (set.map f (set.singleton x)) = (set.singleton (f x))
+      Node mappedElement = nm->mkNode(APPLY_UF, n[0], n[1][0]);
+      Node ret = nm->mkSingleton(rangeType, mappedElement);
+      return RewriteResponse(REWRITE_AGAIN_FULL, ret);
+    }
+    case SET_UNION:
+    {
+      // (set.map f (set.union A B)) = (set.union (set.map f A) (set.map f B))
+      Node a = nm->mkNode(SET_MAP, n[0], n[1][0]);
+      Node b = nm->mkNode(SET_MAP, n[0], n[1][1]);
+      Node ret = nm->mkNode(SET_UNION, a, b);
+      return RewriteResponse(REWRITE_AGAIN_FULL, ret);
+    }
+
+    default: return RewriteResponse(REWRITE_DONE, n);
+  }
+}
+
 }  // namespace sets
 }  // namespace theory
 }  // namespace cvc5
diff --git a/src/theory/sets/theory_sets_rewriter.h b/src/theory/sets/theory_sets_rewriter.h
index 693731862..1d2f36255 100644
--- a/src/theory/sets/theory_sets_rewriter.h
+++ b/src/theory/sets/theory_sets_rewriter.h
@@ -74,6 +74,15 @@ private:
    * Returns true if elementTerm is in setTerm, where both terms are constants.
    */
   bool checkConstantMembership(TNode elementTerm, TNode setTerm);
+ /**
+  *  rewrites for n include:
+  *  - (set.map f (as set.empty (Set T1)) = (as set.empty (Set T2))
+  *  - (set.map f (set.singleton x)) = (set.singleton (apply f x))
+  *  - (set.map f (set.union A B)) =
+  *       (set.union (set.map f A) (set.map f B))
+  *  where f: T1 -> T2
+  */
+ RewriteResponse postRewriteMap(TNode n);
 }; /* class TheorySetsRewriter */
 
 }  // namespace sets
diff --git a/src/theory/sets/theory_sets_type_rules.cpp b/src/theory/sets/theory_sets_type_rules.cpp
index 6602e7d03..8101661bd 100644
--- a/src/theory/sets/theory_sets_type_rules.cpp
+++ b/src/theory/sets/theory_sets_type_rules.cpp
@@ -293,6 +293,48 @@ TypeNode InsertTypeRule::computeType(NodeManager* nodeManager,
   return setType;
 }
 
+TypeNode SetMapTypeRule::computeType(NodeManager* nodeManager,
+                                     TNode n,
+                                     bool check)
+{
+  Assert(n.getKind() == kind::SET_MAP);
+  TypeNode functionType = n[0].getType(check);
+  TypeNode setType = n[1].getType(check);
+  if (check)
+  {
+    if (!setType.isSet())
+    {
+      throw TypeCheckingExceptionPrivate(
+          n,
+          "set.map operator expects a set in the second argument, "
+          "a non-set is found");
+    }
+
+    TypeNode elementType = setType.getSetElementType();
+
+    if (!(functionType.isFunction()))
+    {
+      std::stringstream ss;
+      ss << "Operator " << n.getKind() << " expects a function of type  (-> "
+         << elementType << " *) as a first argument. "
+         << "Found a term of type '" << functionType << "'.";
+      throw TypeCheckingExceptionPrivate(n, ss.str());
+    }
+    std::vector<TypeNode> argTypes = functionType.getArgTypes();
+    if (!(argTypes.size() == 1 && argTypes[0] == elementType))
+    {
+      std::stringstream ss;
+      ss << "Operator " << n.getKind() << " expects a function of type  (-> "
+         << elementType << " *). "
+         << "Found a function of type '" << functionType << "'.";
+      throw TypeCheckingExceptionPrivate(n, ss.str());
+    }
+  }
+  TypeNode rangeType = n[0].getType().getRangeType();
+  TypeNode retType = nodeManager->mkSetType(rangeType);
+  return retType;
+}
+
 TypeNode RelBinaryOperatorTypeRule::computeType(NodeManager* nodeManager,
                                                 TNode n,
                                                 bool check)
diff --git a/src/theory/sets/theory_sets_type_rules.h b/src/theory/sets/theory_sets_type_rules.h
index f0d8991ce..b793205c0 100644
--- a/src/theory/sets/theory_sets_type_rules.h
+++ b/src/theory/sets/theory_sets_type_rules.h
@@ -89,6 +89,15 @@ struct InsertTypeRule
   static TypeNode computeType(NodeManager* nodeManager, TNode n, bool check);
 };
 
+/**
+ * Type rule for (set.map f S) to make sure f is a unary function of type
+ * (-> T1 T2) where S is a bag of type (Set T1)
+ */
+struct SetMapTypeRule
+{
+  static TypeNode computeType(NodeManager* nodeManager, TNode n, bool check);
+}; /* struct SetMapTypeRule */
+
 struct RelBinaryOperatorTypeRule
 {
   static TypeNode computeType(NodeManager* nodeManager, TNode n, bool check);
diff --git a/test/unit/theory/CMakeLists.txt b/test/unit/theory/CMakeLists.txt
index 15c5e8570..9a0a3e26d 100644
--- a/test/unit/theory/CMakeLists.txt
+++ b/test/unit/theory/CMakeLists.txt
@@ -35,6 +35,7 @@ cvc5_add_unit_test_white(theory_int_opt_white theory)
 cvc5_add_unit_test_white(theory_opt_multigoal_white theory)
 cvc5_add_unit_test_white(theory_quantifiers_bv_instantiator_white theory)
 cvc5_add_unit_test_white(theory_quantifiers_bv_inverter_white theory)
+cvc5_add_unit_test_white(theory_sets_rewriter_white theory)
 cvc5_add_unit_test_white(theory_sets_type_enumerator_white theory)
 cvc5_add_unit_test_white(theory_sets_type_rules_white theory)
 cvc5_add_unit_test_white(theory_strings_skolem_cache_black theory)
diff --git a/test/unit/theory/theory_sets_rewriter_white.cpp b/test/unit/theory/theory_sets_rewriter_white.cpp
new file mode 100644
index 000000000..49dbe0c73
--- /dev/null
+++ b/test/unit/theory/theory_sets_rewriter_white.cpp
@@ -0,0 +1,93 @@
+/******************************************************************************
+ * Top contributors (to current version):
+ *   Mudathir Mohamed
+ *
+ * This file is part of the cvc5 project.
+ *
+ * Copyright (c) 2009-2021 by the authors listed in the file AUTHORS
+ * in the top-level source directory and their institutional affiliations.
+ * All rights reserved.  See the file COPYING in the top-level source
+ * directory for licensing information.
+ * ****************************************************************************
+ *
+ * White box testing of sets rewriter
+ */
+
+#include "expr/dtype.h"
+#include "expr/emptyset.h"
+#include "test_smt.h"
+#include "theory/sets/theory_sets_rewriter.h"
+#include "util/rational.h"
+#include "util/string.h"
+
+namespace cvc5 {
+
+using namespace theory;
+using namespace kind;
+using namespace theory::sets;
+
+namespace test {
+
+typedef expr::Attribute<Node, Node> attribute;
+
+class TestTheoryWhiteSetsRewriter : public TestSmt
+{
+ protected:
+  void SetUp() override
+  {
+    TestSmt::SetUp();
+    d_rewriter.reset(new TheorySetsRewriter());
+  }
+  std::unique_ptr<TheorySetsRewriter> d_rewriter;
+};
+
+TEST_F(TestTheoryWhiteSetsRewriter, map)
+{
+  Node one = d_nodeManager->mkConst(CONST_RATIONAL, Rational(1));
+  TypeNode stringType = d_nodeManager->stringType();
+  TypeNode integerType = d_nodeManager->integerType();
+  Node emptysetInteger =
+      d_nodeManager->mkConst(EmptySet(d_nodeManager->mkSetType(integerType)));
+  Node emptysetString =
+      d_nodeManager->mkConst(EmptySet(d_nodeManager->mkSetType(stringType)));
+  Node x = d_nodeManager->mkBoundVar("x", stringType);
+  Node bound = d_nodeManager->mkNode(kind::BOUND_VAR_LIST, x);
+  Node lambda = d_nodeManager->mkNode(LAMBDA, bound, one);
+
+  // (set.map (lambda ((x U))  t) (as set.empty (Set String)) =
+  // (as set.empty (Set Int))
+  Node n1 = d_nodeManager->mkNode(SET_MAP, lambda, emptysetString);
+  RewriteResponse response1 = d_rewriter->postRewrite(n1);
+  ASSERT_TRUE(response1.d_node == emptysetInteger
+              && response1.d_status == REWRITE_DONE);
+
+  Node a = d_nodeManager->mkConst(String("a"));
+  Node b = d_nodeManager->mkConst(String("b"));
+  Node A = d_nodeManager->mkSingleton(d_nodeManager->stringType(), a);
+  Node B = d_nodeManager->mkSingleton(d_nodeManager->stringType(), b);
+  Node unionAB = d_nodeManager->mkNode(SET_UNION, A, B);
+
+  // (set.map
+  //   (lambda ((x String)) 1)
+  //   (set.union (set.singleton "a") (set.singleton "b"))) = (set.singleton 1))
+  Node n2 = d_nodeManager->mkNode(SET_MAP, lambda, unionAB);
+  Node rewritten2 = Rewriter::rewrite(n2);
+  Node bag = d_nodeManager->mkSingleton(d_nodeManager->integerType(), one);
+  ASSERT_TRUE(rewritten2 == bag);
+
+  //  - (set.map f (set.union K1 K2)) =
+  //      (set.union (set.map f K1) (set.map f K2))
+  Node k1 = d_skolemManager->mkDummySkolem("K1", A.getType());
+  Node k2 = d_skolemManager->mkDummySkolem("K2", A.getType());
+  Node f = d_skolemManager->mkDummySkolem("f", lambda.getType());
+  Node unionK1K2 = d_nodeManager->mkNode(SET_UNION, k1, k2);
+  Node n3 = d_nodeManager->mkNode(SET_MAP, f, unionK1K2);
+  Node rewritten3 = Rewriter::rewrite(n3);
+  Node mapK1 = d_nodeManager->mkNode(SET_MAP, f, k1);
+  Node mapK2 = d_nodeManager->mkNode(SET_MAP, f, k2);
+  Node unionMapK1K2 = d_nodeManager->mkNode(SET_UNION, mapK1, mapK2);
+  ASSERT_TRUE(rewritten3 == unionMapK1K2);
+}
+
+}  // namespace test
+}  // namespace cvc5