Add kind BAG_MAP and its type rule to bags (#6503)

This PR adds kind BAG_MAP to bags.

2 files changed, 82 insertions, 0 deletions

diff --git a/test/unit/theory/theory_bags_rewriter_white.cpp b/test/unit/theory/theory_bags_rewriter_white.cpp
index f70ff0c5d..e63fb3b20 100644
--- a/test/unit/theory/theory_bags_rewriter_white.cpp
+++ b/test/unit/theory/theory_bags_rewriter_white.cpp
@@ -694,5 +694,52 @@ TEST_F(TestTheoryWhiteBagsRewriter, to_set)
   ASSERT_TRUE(response.d_node == singleton
               && response.d_status == REWRITE_AGAIN_FULL);
 }
+
+TEST_F(TestTheoryWhiteBagsRewriter, map)
+{
+  Node emptybagString =
+      d_nodeManager->mkConst(EmptyBag(d_nodeManager->stringType()));
+
+  Node one = d_nodeManager->mkConst(Rational(1));
+  Node x = d_nodeManager->mkBoundVar("x", d_nodeManager->integerType());
+  std::vector<Node> args;
+  args.push_back(x);
+  Node bound = d_nodeManager->mkNode(kind::BOUND_VAR_LIST, args);
+  Node lambda = d_nodeManager->mkNode(LAMBDA, bound, one);
+
+  // (bag.map (lambda ((x U))  t) emptybag) = emptybag
+  Node n1 = d_nodeManager->mkNode(BAG_MAP, lambda, emptybagString);
+  RewriteResponse response1 = d_rewriter->postRewrite(n1);
+  TypeNode type = d_nodeManager->mkBagType(d_nodeManager->integerType());
+  Node emptybagInteger = d_nodeManager->mkConst(EmptyBag(type));
+  ASSERT_TRUE(response1.d_node == emptybagInteger
+              && response1.d_status == REWRITE_AGAIN_FULL);
+
+  std::vector<Node> elements = getNStrings(2);
+  Node a = d_nodeManager->mkConst(String("a"));
+  Node b = d_nodeManager->mkConst(String("b"));
+  Node A = d_nodeManager->mkBag(d_nodeManager->stringType(),
+                                a,
+                                d_nodeManager->mkConst(Rational(3)));
+  Node B = d_nodeManager->mkBag(d_nodeManager->stringType(),
+                                b,
+                                d_nodeManager->mkConst(Rational(4)));
+  Node unionDisjointAB = d_nodeManager->mkNode(UNION_DISJOINT, A, B);
+
+  ASSERT_TRUE(unionDisjointAB.isConst());
+  // - (bag.map (lambda ((x Int)) 1) (union_disjoint (bag "a" 3) (bag "b" 4))) =
+  //        (bag 1 7))
+  Node n2 = d_nodeManager->mkNode(BAG_MAP, lambda, unionDisjointAB);
+
+  std::cout << n2 << std::endl;
+
+  Node rewritten = Rewriter:: rewrite(n2);
+  std::cout << rewritten << std::endl;
+
+  Node bag = d_nodeManager->mkBag(d_nodeManager->integerType(),
+                                  one,               d_nodeManager->mkConst(Rational(7)));
+  ASSERT_TRUE(rewritten == bag);
+}
+
 }  // namespace test
 }  // namespace cvc5
diff --git a/test/unit/theory/theory_bags_type_rules_white.cpp b/test/unit/theory/theory_bags_type_rules_white.cpp
index 8013d06ea..eace59c96 100644
--- a/test/unit/theory/theory_bags_type_rules_white.cpp
+++ b/test/unit/theory/theory_bags_type_rules_white.cpp
@@ -111,5 +111,40 @@ TEST_F(TestTheoryWhiteBagsTypeRule, to_set_operator)
   ASSERT_NO_THROW(d_nodeManager->mkNode(BAG_TO_SET, bag));
   ASSERT_TRUE(d_nodeManager->mkNode(BAG_TO_SET, bag).getType().isSet());
 }
+
+TEST_F(TestTheoryWhiteBagsTypeRule, map_operator)
+{
+  std::vector<Node> elements = getNStrings(1);
+  Node bag = d_nodeManager->mkBag(d_nodeManager->stringType(),
+                                  elements[0],
+                                  d_nodeManager->mkConst(Rational(10)));
+  Node set =
+      d_nodeManager->mkSingleton(d_nodeManager->stringType(), elements[0]);
+
+  Node x1 = d_nodeManager->mkBoundVar("x", d_nodeManager->stringType());
+  Node length = d_nodeManager->mkNode(STRING_LENGTH, x1);
+  std::vector<Node> args1;
+  args1.push_back(x1);
+  Node bound1 = d_nodeManager->mkNode(kind::BOUND_VAR_LIST, args1);
+  Node lambda1 = d_nodeManager->mkNode(LAMBDA, bound1, length);
+
+  ASSERT_NO_THROW(d_nodeManager->mkNode(BAG_MAP, lambda1, bag));
+  Node mappedBag = d_nodeManager->mkNode(BAG_MAP, lambda1, bag);
+  ASSERT_TRUE(mappedBag.getType().isBag());
+  ASSERT_EQ(d_nodeManager->integerType(),
+            mappedBag.getType().getBagElementType());
+
+  Node one = d_nodeManager->mkConst(Rational(1));
+  Node x2 = d_nodeManager->mkBoundVar("x", d_nodeManager->integerType());
+  std::vector<Node> args2;
+  args2.push_back(x2);
+  Node bound2 = d_nodeManager->mkNode(kind::BOUND_VAR_LIST, args2);
+  Node lambda2 = d_nodeManager->mkNode(LAMBDA, bound2, one);
+  ASSERT_THROW(d_nodeManager->mkNode(BAG_MAP, lambda2, bag).getType(true),
+               TypeCheckingExceptionPrivate);
+  ASSERT_THROW(d_nodeManager->mkNode(BAG_MAP, lambda2, set).getType(true),
+               TypeCheckingExceptionPrivate);
+}
+
 }  // namespace test
 }  // namespace cvc5