Merge branch 'master' into issue7504issue7504

7 files changed, 275 insertions, 18 deletions

diff --git a/src/theory/bags/bag_solver.cpp b/src/theory/bags/bag_solver.cpp
index 203903d88..341798eb2 100644
--- a/src/theory/bags/bag_solver.cpp
+++ b/src/theory/bags/bag_solver.cpp
@@ -31,8 +31,16 @@ namespace cvc5 {
 namespace theory {
 namespace bags {
 
-BagSolver::BagSolver(SolverState& s, InferenceManager& im, TermRegistry& tr)
-    : d_state(s), d_ig(&s, &im), d_im(im), d_termReg(tr)
+BagSolver::BagSolver(Env& env,
+                     SolverState& s,
+                     InferenceManager& im,
+                     TermRegistry& tr)
+    : EnvObj(env),
+      d_state(s),
+      d_ig(&s, &im),
+      d_im(im),
+      d_termReg(tr),
+      d_mapCache(userContext())
 {
   d_zero = NodeManager::currentNM()->mkConst(Rational(0));
   d_one = NodeManager::currentNM()->mkConst(Rational(1));
@@ -69,6 +77,7 @@ void BagSolver::postCheck()
         case kind::DIFFERENCE_SUBTRACT: checkDifferenceSubtract(n); break;
         case kind::DIFFERENCE_REMOVE: checkDifferenceRemove(n); break;
         case kind::DUPLICATE_REMOVAL: checkDuplicateRemoval(n); break;
+        case kind::BAG_MAP: checkMap(n); break;
         default: break;
       }
       it++;
@@ -210,6 +219,34 @@ void BagSolver::checkDisequalBagTerms()
   }
 }
 
+void BagSolver::checkMap(Node n)
+{
+  Assert(n.getKind() == BAG_MAP);
+  const set<Node>& downwards = d_state.getElements(n);
+  const set<Node>& upwards = d_state.getElements(n[1]);
+  for (const Node& y : downwards)
+  {
+    if (d_mapCache.count(n) && d_mapCache[n].get()->contains(y))
+    {
+      continue;
+    }
+    auto [downInference, uf, preImageSize] = d_ig.mapDownwards(n, y);
+    d_im.lemmaTheoryInference(&downInference);
+    for (const Node& x : upwards)
+    {
+      InferInfo upInference = d_ig.mapUpwards(n, uf, preImageSize, y, x);
+      d_im.lemmaTheoryInference(&upInference);
+    }
+    if (!d_mapCache.count(n))
+    {
+      std::shared_ptr<context::CDHashSet<Node> > set =
+          std::make_shared<context::CDHashSet<Node> >(userContext());
+      d_mapCache.insert(n, set);
+    }
+    d_mapCache[n].get()->insert(y);
+  }
+}
+
 }  // namespace bags
 }  // namespace theory
 }  // namespace cvc5
diff --git a/src/theory/bags/bag_solver.h b/src/theory/bags/bag_solver.h
index 5fb49fab7..9155c93d0 100644
--- a/src/theory/bags/bag_solver.h
+++ b/src/theory/bags/bag_solver.h
@@ -13,7 +13,10 @@
  * Solver for the theory of bags.
  */
 
+#include "context/cdhashmap.h"
+#include "context/cdhashset.h"
 #include "cvc5_private.h"
+#include "smt/env_obj.h"
 
 #ifndef CVC5__THEORY__BAG__SOLVER_H
 #define CVC5__THEORY__BAG__SOLVER_H
@@ -31,10 +34,10 @@ class TermRegistry;
 /** The solver for the theory of bags
  *
  */
-class BagSolver
+class BagSolver : protected EnvObj
 {
  public:
-  BagSolver(SolverState& s, InferenceManager& im, TermRegistry& tr);
+  BagSolver(Env& env, SolverState& s, InferenceManager& im, TermRegistry& tr);
   ~BagSolver();
 
   void postCheck();
@@ -73,6 +76,8 @@ class BagSolver
   void checkNonNegativeCountTerms(const Node& bag, const Node& element);
   /** apply inference rules for disequal bag terms */
   void checkDisequalBagTerms();
+  /** apply inference rules for map operator */
+  void checkMap(Node n);
 
   /** The solver state object */
   SolverState& d_state;
@@ -82,6 +87,14 @@ class BagSolver
   InferenceManager& d_im;
   /** Reference to the term registry of theory of bags */
   TermRegistry& d_termReg;
+
+  /** a cache that stores bags of kind BAG_MAP and those element representatives
+   * which we generated their inferences.
+   */
+  using BagElementsMap =
+      context::CDHashMap<Node, std::shared_ptr<context::CDHashSet<Node> > >;
+  BagElementsMap d_mapCache;
+
   /** Commonly used constants */
   Node d_true;
   Node d_false;
diff --git a/src/theory/bags/bags_rewriter.cpp b/src/theory/bags/bags_rewriter.cpp
index 646bb28cf..0be83fb13 100644
--- a/src/theory/bags/bags_rewriter.cpp
+++ b/src/theory/bags/bags_rewriter.cpp
@@ -532,7 +532,8 @@ BagsRewriteResponse BagsRewriter::postRewriteMap(const TNode& n) const
     case MK_BAG:
     {
       Node mappedElement = d_nm->mkNode(APPLY_UF, n[0], n[1][0]);
-      Node ret = d_nm->mkNode(MK_BAG, mappedElement, n[1][0]);
+      Node ret =
+          d_nm->mkBag(n[0].getType().getRangeType(), mappedElement, n[1][1]);
       return BagsRewriteResponse(ret, Rewrite::MAP_MK_BAG);
     }
 
diff --git a/src/theory/bags/inference_generator.cpp b/src/theory/bags/inference_generator.cpp
index 556dc76d6..734572f7c 100644
--- a/src/theory/bags/inference_generator.cpp
+++ b/src/theory/bags/inference_generator.cpp
@@ -20,6 +20,7 @@
 #include "expr/skolem_manager.h"
 #include "theory/bags/inference_manager.h"
 #include "theory/bags/solver_state.h"
+#include "theory/quantifiers/fmf/bounded_integers.h"
 #include "theory/uf/equality_engine.h"
 #include "util/rational.h"
 
@@ -42,7 +43,7 @@ InferInfo InferenceGenerator::nonNegativeCount(Node n, Node e)
   Assert(n.getType().isBag());
   Assert(e.getType() == n.getType().getBagElementType());
 
-  InferInfo inferInfo(d_im, InferenceId::BAG_NON_NEGATIVE_COUNT);
+  InferInfo inferInfo(d_im, InferenceId::BAGS_NON_NEGATIVE_COUNT);
   Node count = d_nm->mkNode(kind::BAG_COUNT, e, n);
 
   Node gte = d_nm->mkNode(kind::GEQ, count, d_zero);
@@ -59,7 +60,7 @@ InferInfo InferenceGenerator::mkBag(Node n, Node e)
   {
     // TODO issue #78: refactor this with BagRewriter
     // (=> true (= (bag.count e (bag e c)) c))
-    InferInfo inferInfo(d_im, InferenceId::BAG_MK_BAG_SAME_ELEMENT);
+    InferInfo inferInfo(d_im, InferenceId::BAGS_MK_BAG_SAME_ELEMENT);
     Node skolem = getSkolem(n, inferInfo);
     Node count = getMultiplicityTerm(e, skolem);
     inferInfo.d_conclusion = count.eqNode(n[1]);
@@ -70,7 +71,7 @@ InferInfo InferenceGenerator::mkBag(Node n, Node e)
     // (=>
     //   true
     //   (= (bag.count e (bag x c)) (ite (= e x) c 0)))
-    InferInfo inferInfo(d_im, InferenceId::BAG_MK_BAG);
+    InferInfo inferInfo(d_im, InferenceId::BAGS_MK_BAG);
     Node skolem = getSkolem(n, inferInfo);
     Node count = getMultiplicityTerm(e, skolem);
     Node same = d_nm->mkNode(kind::EQUAL, n[0], e);
@@ -81,6 +82,27 @@ InferInfo InferenceGenerator::mkBag(Node n, Node e)
   }
 }
 
+/**
+ * A bound variable corresponding to the universally quantified integer
+ * variable used to range over the distinct elements in a bag, used
+ * for axiomatizing the behavior of some term.
+ */
+struct FirstIndexVarAttributeId
+{
+};
+typedef expr::Attribute<FirstIndexVarAttributeId, Node> FirstIndexVarAttribute;
+
+/**
+ * A bound variable corresponding to the universally quantified integer
+ * variable used to range over the distinct elements in a bag, used
+ * for axiomatizing the behavior of some term.
+ */
+struct SecondIndexVarAttributeId
+{
+};
+typedef expr::Attribute<SecondIndexVarAttributeId, Node>
+    SecondIndexVarAttribute;
+
 struct BagsDeqAttributeId
 {
 };
@@ -93,7 +115,7 @@ InferInfo InferenceGenerator::bagDisequality(Node n)
   Node A = n[0];
   Node B = n[1];
 
-  InferInfo inferInfo(d_im, InferenceId::BAG_DISEQUALITY);
+  InferInfo inferInfo(d_im, InferenceId::BAGS_DISEQUALITY);
 
   TypeNode elementType = A.getType().getBagElementType();
   BoundVarManager* bvm = d_nm->getBoundVarManager();
@@ -126,7 +148,7 @@ InferInfo InferenceGenerator::empty(Node n, Node e)
   Assert(n.getKind() == kind::EMPTYBAG);
   Assert(e.getType() == n.getType().getBagElementType());
 
-  InferInfo inferInfo(d_im, InferenceId::BAG_EMPTY);
+  InferInfo inferInfo(d_im, InferenceId::BAGS_EMPTY);
   Node skolem = getSkolem(n, inferInfo);
   Node count = getMultiplicityTerm(e, skolem);
 
@@ -142,7 +164,7 @@ InferInfo InferenceGenerator::unionDisjoint(Node n, Node e)
 
   Node A = n[0];
   Node B = n[1];
-  InferInfo inferInfo(d_im, InferenceId::BAG_UNION_DISJOINT);
+  InferInfo inferInfo(d_im, InferenceId::BAGS_UNION_DISJOINT);
 
   Node countA = getMultiplicityTerm(e, A);
   Node countB = getMultiplicityTerm(e, B);
@@ -164,7 +186,7 @@ InferInfo InferenceGenerator::unionMax(Node n, Node e)
 
   Node A = n[0];
   Node B = n[1];
-  InferInfo inferInfo(d_im, InferenceId::BAG_UNION_MAX);
+  InferInfo inferInfo(d_im, InferenceId::BAGS_UNION_MAX);
 
   Node countA = getMultiplicityTerm(e, A);
   Node countB = getMultiplicityTerm(e, B);
@@ -187,7 +209,7 @@ InferInfo InferenceGenerator::intersection(Node n, Node e)
 
   Node A = n[0];
   Node B = n[1];
-  InferInfo inferInfo(d_im, InferenceId::BAG_INTERSECTION_MIN);
+  InferInfo inferInfo(d_im, InferenceId::BAGS_INTERSECTION_MIN);
 
   Node countA = getMultiplicityTerm(e, A);
   Node countB = getMultiplicityTerm(e, B);
@@ -208,7 +230,7 @@ InferInfo InferenceGenerator::differenceSubtract(Node n, Node e)
 
   Node A = n[0];
   Node B = n[1];
-  InferInfo inferInfo(d_im, InferenceId::BAG_DIFFERENCE_SUBTRACT);
+  InferInfo inferInfo(d_im, InferenceId::BAGS_DIFFERENCE_SUBTRACT);
 
   Node countA = getMultiplicityTerm(e, A);
   Node countB = getMultiplicityTerm(e, B);
@@ -230,7 +252,7 @@ InferInfo InferenceGenerator::differenceRemove(Node n, Node e)
 
   Node A = n[0];
   Node B = n[1];
-  InferInfo inferInfo(d_im, InferenceId::BAG_DIFFERENCE_REMOVE);
+  InferInfo inferInfo(d_im, InferenceId::BAGS_DIFFERENCE_REMOVE);
 
   Node countA = getMultiplicityTerm(e, A);
   Node countB = getMultiplicityTerm(e, B);
@@ -251,7 +273,7 @@ InferInfo InferenceGenerator::duplicateRemoval(Node n, Node e)
   Assert(e.getType() == n[0].getType().getBagElementType());
 
   Node A = n[0];
-  InferInfo inferInfo(d_im, InferenceId::BAG_DUPLICATE_REMOVAL);
+  InferInfo inferInfo(d_im, InferenceId::BAGS_DUPLICATE_REMOVAL);
 
   Node countA = getMultiplicityTerm(e, A);
   Node skolem = getSkolem(n, inferInfo);
@@ -270,6 +292,137 @@ Node InferenceGenerator::getMultiplicityTerm(Node element, Node bag)
   return count;
 }
 
+std::tuple<InferInfo, Node, Node> InferenceGenerator::mapDownwards(Node n,
+                                                                   Node e)
+{
+  Assert(n.getKind() == kind::BAG_MAP && n[1].getType().isBag());
+  Assert(n[0].getType().isFunction()
+         && n[0].getType().getArgTypes().size() == 1);
+  Assert(e.getType() == n[0].getType().getRangeType());
+
+  InferInfo inferInfo(d_im, InferenceId::BAGS_MAP);
+
+  Node f = n[0];
+  Node A = n[1];
+  // declare an uninterpreted function uf: Int -> T
+  TypeNode domainType = f.getType().getArgTypes()[0];
+  TypeNode ufType = d_nm->mkFunctionType(d_nm->integerType(), domainType);
+  Node uf =
+      d_sm->mkSkolemFunction(SkolemFunId::BAGS_MAP_PREIMAGE, ufType, {n, e});
+
+  // declare uninterpreted function sum: Int -> Int
+  TypeNode sumType =
+      d_nm->mkFunctionType(d_nm->integerType(), d_nm->integerType());
+  Node sum = d_sm->mkSkolemFunction(SkolemFunId::BAGS_MAP_SUM, sumType, {n, e});
+
+  // (= (sum 0) 0)
+  Node sum_zero = d_nm->mkNode(kind::APPLY_UF, sum, d_zero);
+  Node baseCase = d_nm->mkNode(Kind::EQUAL, sum_zero, d_zero);
+
+  // guess the size of the preimage of e
+  Node preImageSize = d_sm->mkDummySkolem("preImageSize", d_nm->integerType());
+
+  // (= (sum preImageSize) (bag.count e skolem))
+  Node mapSkolem = getSkolem(n, inferInfo);
+  Node countE = getMultiplicityTerm(e, mapSkolem);
+  Node totalSum = d_nm->mkNode(kind::APPLY_UF, sum, preImageSize);
+  Node totalSumEqualCountE = d_nm->mkNode(kind::EQUAL, totalSum, countE);
+
+  // (forall ((i Int))
+  //        (let ((uf_i (uf i)))
+  //          (let ((count_uf_i (bag.count uf_i A)))
+  //            (=>
+  //             (and (>= i 1) (<= i preImageSize))
+  //             (and
+  //               (= (f uf_i) e)
+  //               (>= count_uf_i 1)
+  //               (= (sum i) (+ (sum (- i 1)) count_uf_i))
+  //               (forall ((j Int))
+  //                 (=>
+  //                  (and (< i j) (<= j preImageSize))
+  //                  (not (= (uf i) (uf j))))))
+  //               )))))
+
+  BoundVarManager* bvm = d_nm->getBoundVarManager();
+  Node i = bvm->mkBoundVar<FirstIndexVarAttribute>(n, "i", d_nm->integerType());
+  Node j =
+      bvm->mkBoundVar<SecondIndexVarAttribute>(n, "j", d_nm->integerType());
+  Node iList = d_nm->mkNode(kind::BOUND_VAR_LIST, i);
+  Node jList = d_nm->mkNode(kind::BOUND_VAR_LIST, j);
+  Node iPlusOne = d_nm->mkNode(kind::PLUS, i, d_one);
+  Node iMinusOne = d_nm->mkNode(kind::MINUS, i, d_one);
+  Node uf_i = d_nm->mkNode(kind::APPLY_UF, uf, i);
+  Node uf_j = d_nm->mkNode(kind::APPLY_UF, uf, j);
+  Node f_uf_i = d_nm->mkNode(kind::APPLY_UF, f, uf_i);
+  Node uf_iPlusOne = d_nm->mkNode(kind::APPLY_UF, uf, iPlusOne);
+  Node uf_iMinusOne = d_nm->mkNode(kind::APPLY_UF, uf, iMinusOne);
+  Node interval_i = d_nm->mkNode(kind::AND,
+                                 d_nm->mkNode(kind::GEQ, i, d_one),
+                                 d_nm->mkNode(kind::LEQ, i, preImageSize));
+  Node sum_i = d_nm->mkNode(kind::APPLY_UF, sum, i);
+  Node sum_iPlusOne = d_nm->mkNode(kind::APPLY_UF, sum, iPlusOne);
+  Node sum_iMinusOne = d_nm->mkNode(kind::APPLY_UF, sum, iMinusOne);
+  Node count_iMinusOne = d_nm->mkNode(kind::BAG_COUNT, uf_iMinusOne, A);
+  Node count_uf_i = d_nm->mkNode(kind::BAG_COUNT, uf_i, A);
+  Node inductiveCase = d_nm->mkNode(
+      Kind::EQUAL, sum_i, d_nm->mkNode(kind::PLUS, sum_iMinusOne, count_uf_i));
+  Node f_iEqualE = d_nm->mkNode(kind::EQUAL, f_uf_i, e);
+  Node geqOne = d_nm->mkNode(kind::GEQ, count_uf_i, d_one);
+
+  // i < j <= preImageSize
+  Node interval_j = d_nm->mkNode(kind::AND,
+                                 d_nm->mkNode(kind::LT, i, j),
+                                 d_nm->mkNode(kind::LEQ, j, preImageSize));
+  // uf(i) != uf(j)
+  Node uf_i_equals_uf_j = d_nm->mkNode(kind::EQUAL, uf_i, uf_j);
+  Node notEqual = d_nm->mkNode(kind::EQUAL, uf_i, uf_j).negate();
+  Node body_j = d_nm->mkNode(kind::OR, interval_j.negate(), notEqual);
+  Node forAll_j = quantifiers::BoundedIntegers::mkBoundedForall(jList, body_j);
+  Node andNode =
+      d_nm->mkNode(kind::AND, {f_iEqualE, geqOne, inductiveCase, forAll_j});
+  Node body_i = d_nm->mkNode(kind::OR, interval_i.negate(), andNode);
+  Node forAll_i = quantifiers::BoundedIntegers::mkBoundedForall(iList, body_i);
+  Node preImageGTE_zero = d_nm->mkNode(kind::GEQ, preImageSize, d_zero);
+  Node conclusion = d_nm->mkNode(
+      kind::AND, {baseCase, totalSumEqualCountE, forAll_i, preImageGTE_zero});
+  inferInfo.d_conclusion = conclusion;
+
+  std::map<Node, Node> m;
+  m[e] = conclusion;
+  return std::tuple(inferInfo, uf, preImageSize);
+}
+
+InferInfo InferenceGenerator::mapUpwards(
+    Node n, Node uf, Node preImageSize, Node y, Node x)
+{
+  Assert(n.getKind() == kind::BAG_MAP && n[1].getType().isBag());
+  Assert(n[0].getType().isFunction()
+         && n[0].getType().getArgTypes().size() == 1);
+
+  InferInfo inferInfo(d_im, InferenceId::BAGS_MAP);
+  Node f = n[0];
+  Node A = n[1];
+
+  Node countA = getMultiplicityTerm(x, A);
+  Node xInA = d_nm->mkNode(kind::GEQ, countA, d_one);
+  Node notEqual =
+      d_nm->mkNode(kind::EQUAL, d_nm->mkNode(kind::APPLY_UF, f, x), y).negate();
+
+  Node k = d_sm->mkDummySkolem("k", d_nm->integerType());
+  Node inRange = d_nm->mkNode(kind::AND,
+                              d_nm->mkNode(kind::GEQ, k, d_one),
+                              d_nm->mkNode(kind::LEQ, k, preImageSize));
+  Node equal =
+      d_nm->mkNode(kind::EQUAL, d_nm->mkNode(kind::APPLY_UF, uf, k), x);
+  Node andNode = d_nm->mkNode(kind::AND, inRange, equal);
+  Node orNode = d_nm->mkNode(kind::OR, notEqual, andNode);
+  Node implies = d_nm->mkNode(kind::IMPLIES, xInA, orNode);
+  inferInfo.d_conclusion = implies;
+  std::cout << "Upwards conclusion: " << inferInfo.d_conclusion << std::endl
+            << std::endl;
+  return inferInfo;
+}
+
 }  // namespace bags
 }  // namespace theory
 }  // namespace cvc5
diff --git a/src/theory/bags/inference_generator.h b/src/theory/bags/inference_generator.h
index 252b9641e..ab3a84b29 100644
--- a/src/theory/bags/inference_generator.h
+++ b/src/theory/bags/inference_generator.h
@@ -164,6 +164,59 @@ class InferenceGenerator
    * where skolem is a fresh variable equals (duplicate_removal A)
    */
   InferInfo duplicateRemoval(Node n, Node e);
+  /**
+   * @param n is (bag.map f A) where f is a function (-> E T), A a bag of type
+   * (Bag E)
+   * @param e is a node of Type E
+   * @return an inference that represents the following implication
+   * (and
+   *   (= (sum 0) 0)
+   *   (= (sum preImageSize) (bag.count e skolem))
+   *   (>= preImageSize 0)
+   *   (forall ((i Int))
+   *          (let ((uf_i (uf i)))
+   *            (let ((count_uf_i (bag.count uf_i A)))
+   *              (=>
+   *               (and (>= i 1) (<= i preImageSize))
+   *               (and
+   *                 (= (f uf_i) e)
+   *                 (>= count_uf_i 1)
+   *                 (= (sum i) (+ (sum (- i 1)) count_uf_i))
+   *                 (forall ((j Int))
+   *                   (or
+   *                     (not (and (< i j) (<= j preImageSize)))
+   *                     (not (= (uf i) (uf j)))) )
+   *                 ))))))
+   * where uf: Int -> E is an uninterpreted function from integers to the
+   * type of the elements of A
+   * preImageSize is the cardinality of the distinct elements in A that are
+   * mapped to e by function f (i.e., preimage of {e})
+   * sum: Int -> Int is a function that aggregates the multiplicities of the
+   * preimage of e,
+   * and skolem is a fresh variable equals (bag.map f A))
+   */
+  std::tuple<InferInfo, Node, Node> mapDownwards(Node n, Node e);
+
+  /**
+   * @param n is (bag.map f A) where f is a function (-> E T), A a bag of type
+   * (Bag E)
+   * @param uf is an uninterpreted function Int -> E
+   * @param preImageSize is the cardinality of the distinct elements in A that
+   * are mapped to y by function f (i.e., preimage of {y})
+   * @param y is an element of type T
+   * @param e is an element of type E
+   * @return an inference that represents the following implication
+   * (=>
+   *   (>= (bag.count x A) 1)
+   *   (or
+   *     (not (= (f x) y)
+   *     (and
+   *       (>= skolem 1)
+   *       (<= skolem preImageSize)
+   *       (= (uf skolem) x)))))
+   * where skolem is a fresh variable
+   */
+  InferInfo mapUpwards(Node n, Node uf, Node preImageSize, Node y, Node x);
 
   /**
    * @param element of type T
diff --git a/src/theory/bags/normal_form.cpp b/src/theory/bags/normal_form.cpp
index 58445de59..69401b3fa 100644
--- a/src/theory/bags/normal_form.cpp
+++ b/src/theory/bags/normal_form.cpp
@@ -159,7 +159,7 @@ Node NormalForm::evaluateBinaryOperation(const TNode& n,
   // handle the remaining elements from A
   remainderOfA(elements, elementsA, itA);
   // handle the remaining elements from B
-  remainderOfA(elements, elementsB, itB);
+  remainderOfB(elements, elementsB, itB);
 
   Trace("bags-evaluate") << "elements: " << elements << std::endl;
   Node bag = constructConstantBagFromElements(n.getType(), elements);
diff --git a/src/theory/bags/theory_bags.cpp b/src/theory/bags/theory_bags.cpp
index 4a6d345e9..f10144255 100644
--- a/src/theory/bags/theory_bags.cpp
+++ b/src/theory/bags/theory_bags.cpp
@@ -36,7 +36,7 @@ TheoryBags::TheoryBags(Env& env, OutputChannel& out, Valuation valuation)
       d_statistics(),
       d_rewriter(&d_statistics.d_rewrites),
       d_termReg(env, d_state, d_im),
-      d_solver(d_state, d_im, d_termReg)
+      d_solver(env, d_state, d_im, d_termReg)
 {
   // use the official theory state and inference manager objects
   d_theoryState = &d_state;