Add operators bag.from_set, bag.to_set to the theory of bags (#5186)

7 files changed, 93 insertions, 1 deletions

diff --git a/src/theory/bags/bags_rewriter.cpp b/src/theory/bags/bags_rewriter.cpp
index dd6565a5b..1faaf55c0 100644
--- a/src/theory/bags/bags_rewriter.cpp
+++ b/src/theory/bags/bags_rewriter.cpp
@@ -71,6 +71,8 @@ RewriteResponse BagsRewriter::postRewrite(TNode n)
       case BAG_CHOOSE: response = rewriteChoose(n); break;
       case BAG_CARD: response = rewriteCard(n); break;
       case BAG_IS_SINGLETON: response = rewriteIsSingleton(n); break;
+      case BAG_FROM_SET: response = rewriteFromSet(n); break;
+      case BAG_TO_SET: response = rewriteToSet(n); break;
       default: response = BagsRewriteResponse(n, Rewrite::NONE); break;
     }
   }
@@ -429,6 +431,33 @@ BagsRewriteResponse BagsRewriter::rewriteIsSingleton(const TNode& n) const
   return BagsRewriteResponse(n, Rewrite::NONE);
 }
 
+BagsRewriteResponse BagsRewriter::rewriteFromSet(const TNode& n) const
+{
+  Assert(n.getKind() == BAG_FROM_SET);
+  if (n[0].getKind() == SINGLETON)
+  {
+    // (bag.from_set (singleton (singleton_op Int) x)) = (mkBag x 1)
+    Node one = d_nm->mkConst(Rational(1));
+    Node bag = d_nm->mkNode(MK_BAG, n[0][0], one);
+    return BagsRewriteResponse(bag, Rewrite::FROM_SINGLETON);
+  }
+  return BagsRewriteResponse(n, Rewrite::NONE);
+}
+
+BagsRewriteResponse BagsRewriter::rewriteToSet(const TNode& n) const
+{
+  Assert(n.getKind() == BAG_TO_SET);
+  if (n[0].getKind() == MK_BAG && n[0][1].isConst()
+      && n[0][1].getConst<Rational>().sgn() == 1)
+  {
+    // (bag.to_set (mkBag x n)) = (singleton (singleton_op T) x)
+    // where n is a positive constant and T is the type of the bag's elements
+    Node bag = d_nm->mkSingleton(n[0][0].getType(), n[0][0]);
+    return BagsRewriteResponse(bag, Rewrite::TO_SINGLETON);
+  }
+  return BagsRewriteResponse(n, Rewrite::NONE);
+}
+
 }  // namespace bags
 }  // namespace theory
 }  // namespace CVC4
diff --git a/src/theory/bags/bags_rewriter.h b/src/theory/bags/bags_rewriter.h
index d165be910..d36a21ccf 100644
--- a/src/theory/bags/bags_rewriter.h
+++ b/src/theory/bags/bags_rewriter.h
@@ -181,6 +181,19 @@ class BagsRewriter : public TheoryRewriter
    */
   BagsRewriteResponse rewriteIsSingleton(const TNode& n) const;
 
+  /**
+   *  rewrites for n include:
+   *  - (bag.from_set (singleton (singleton_op Int) x)) = (mkBag x 1)
+   */
+  BagsRewriteResponse rewriteFromSet(const TNode& n) const;
+
+  /**
+   *  rewrites for n include:
+   *  - (bag.to_set (mkBag x n)) = (singleton (singleton_op T) x)
+   *     where n is a positive constant and T is the type of the bag's elements
+   */
+  BagsRewriteResponse rewriteToSet(const TNode& n) const;
+
  private:
   /** Reference to the rewriter statistics. */
   NodeManager* d_nm;
diff --git a/src/theory/bags/kinds b/src/theory/bags/kinds
index cdbef58de..72326de08 100644
--- a/src/theory/bags/kinds
+++ b/src/theory/bags/kinds
@@ -53,6 +53,8 @@ operator MK_BAG            2  "constructs a bag from one element along with its
 operator BAG_IS_SINGLETON  1  "return whether the given bag is a singleton"
 
 operator BAG_CARD          1  "bag cardinality operator"
+operator BAG_FROM_SET      1  "converts a set to a bag"
+operator BAG_TO_SET        1  "converts a bag to a set"
 
 # The operator choose returns an element from a given bag.
 # If bag A = {("a", 1)}, then the term (choose A) is equivalent to the term a.
@@ -72,6 +74,8 @@ typerule EMPTYBAG            ::CVC4::theory::bags::EmptyBagTypeRule
 typerule BAG_CARD            ::CVC4::theory::bags::CardTypeRule
 typerule BAG_CHOOSE          ::CVC4::theory::bags::ChooseTypeRule
 typerule BAG_IS_SINGLETON    ::CVC4::theory::bags::IsSingletonTypeRule
+typerule BAG_FROM_SET        ::CVC4::theory::bags::FromSetTypeRule
+typerule BAG_TO_SET          ::CVC4::theory::bags::ToSetTypeRule
 
 construle UNION_DISJOINT     ::CVC4::theory::bags::BinaryOperatorTypeRule
 construle MK_BAG             ::CVC4::theory::bags::MkBagTypeRule
diff --git a/src/theory/bags/rewrites.cpp b/src/theory/bags/rewrites.cpp
index 758f8a6e6..be3b3cc71 100644
--- a/src/theory/bags/rewrites.cpp
+++ b/src/theory/bags/rewrites.cpp
@@ -31,6 +31,7 @@ const char* toString(Rewrite r)
     case Rewrite::CONSTANT_EVALUATION: return "CONSTANT_EVALUATION";
     case Rewrite::COUNT_EMPTY: return "COUNT_EMPTY";
     case Rewrite::COUNT_MK_BAG: return "COUNT_MK_BAG";
+    case Rewrite::FROM_SINGLETON: return "FROM_SINGLETON";
     case Rewrite::IDENTICAL_NODES: return "IDENTICAL_NODES";
     case Rewrite::INTERSECTION_EMPTY_LEFT: return "INTERSECTION_EMPTY_LEFT";
     case Rewrite::INTERSECTION_EMPTY_RIGHT: return "INTERSECTION_EMPTY_RIGHT";
@@ -53,6 +54,7 @@ const char* toString(Rewrite r)
     case Rewrite::SUBTRACT_RETURN_LEFT: return "SUBTRACT_RETURN_LEFT";
     case Rewrite::SUBTRACT_SAME: return "SUBTRACT_SAME";
     case Rewrite::UNION_DISJOINT_EMPTY_LEFT: return "UNION_DISJOINT_EMPTY_LEFT";
+    case Rewrite::TO_SINGLETON: return "TO_SINGLETON";
     case Rewrite::UNION_DISJOINT_EMPTY_RIGHT:
       return "UNION_DISJOINT_EMPTY_RIGHT";
     case Rewrite::UNION_DISJOINT_MAX_MIN: return "UNION_DISJOINT_MAX_MIN";
diff --git a/src/theory/bags/rewrites.h b/src/theory/bags/rewrites.h
index 13b0ff202..dc1921e24 100644
--- a/src/theory/bags/rewrites.h
+++ b/src/theory/bags/rewrites.h
@@ -29,13 +29,14 @@ namespace bags {
  */
 enum class Rewrite : uint32_t
 {
-  NONE, // no rewrite happened
+  NONE,  // no rewrite happened
   CARD_DISJOINT,
   CARD_MK_BAG,
   CHOOSE_MK_BAG,
   CONSTANT_EVALUATION,
   COUNT_EMPTY,
   COUNT_MK_BAG,
+  FROM_SINGLETON,
   IDENTICAL_NODES,
   INTERSECTION_EMPTY_LEFT,
   INTERSECTION_EMPTY_RIGHT,
@@ -55,6 +56,7 @@ enum class Rewrite : uint32_t
   SUBTRACT_MIN,
   SUBTRACT_RETURN_LEFT,
   SUBTRACT_SAME,
+  TO_SINGLETON,
   UNION_DISJOINT_EMPTY_LEFT,
   UNION_DISJOINT_EMPTY_RIGHT,
   UNION_DISJOINT_MAX_MIN,
diff --git a/src/theory/bags/theory_bags.cpp b/src/theory/bags/theory_bags.cpp
index e4cd64b48..9dcad9bb1 100644
--- a/src/theory/bags/theory_bags.cpp
+++ b/src/theory/bags/theory_bags.cpp
@@ -65,6 +65,8 @@ void TheoryBags::finishInit()
   d_equalityEngine->addFunctionKind(BAG_COUNT);
   d_equalityEngine->addFunctionKind(MK_BAG);
   d_equalityEngine->addFunctionKind(BAG_CARD);
+  d_equalityEngine->addFunctionKind(BAG_FROM_SET);
+  d_equalityEngine->addFunctionKind(BAG_TO_SET);
 }
 
 void TheoryBags::postCheck(Effort level) {}
diff --git a/src/theory/bags/theory_bags_type_rules.h b/src/theory/bags/theory_bags_type_rules.h
index e4279479d..67293e222 100644
--- a/src/theory/bags/theory_bags_type_rules.h
+++ b/src/theory/bags/theory_bags_type_rules.h
@@ -232,6 +232,46 @@ struct ChooseTypeRule
   }
 }; /* struct ChooseTypeRule */
 
+struct FromSetTypeRule
+{
+  static TypeNode computeType(NodeManager* nodeManager, TNode n, bool check)
+  {
+    Assert(n.getKind() == kind::BAG_FROM_SET);
+    TypeNode setType = n[0].getType(check);
+    if (check)
+    {
+      if (!setType.isSet())
+      {
+        throw TypeCheckingExceptionPrivate(
+            n, "bag.from_set operator expects a set, a non-set is found");
+      }
+    }
+    TypeNode elementType = setType.getSetElementType();
+    TypeNode bagType = nodeManager->mkBagType(elementType);
+    return bagType;
+  }
+}; /* struct FromSetTypeRule */
+
+struct ToSetTypeRule
+{
+  static TypeNode computeType(NodeManager* nodeManager, TNode n, bool check)
+  {
+    Assert(n.getKind() == kind::BAG_TO_SET);
+    TypeNode bagType = n[0].getType(check);
+    if (check)
+    {
+      if (!bagType.isBag())
+      {
+        throw TypeCheckingExceptionPrivate(
+            n, "bag.to_set operator expects a bag, a non-bag is found");
+      }
+    }
+    TypeNode elementType = bagType.getBagElementType();
+    TypeNode setType = nodeManager->mkSetType(elementType);
+    return setType;
+  }
+}; /* struct ToSetTypeRule */
+
 struct BagsProperties
 {
   static Cardinality computeCardinality(TypeNode type)