Fix soundess issue for bags with negative multiplicity (#7539)

This PR fixes soundness issues found by cvc5 fuzzy sygus when it answers unsat for sat problems.
One error was with the rewrite rule (bag.count x (bag x c) = c which did not account for cases when c is negative.
This conflicts with fact that all multiplicities are non-negative.
Another error was that the difference_remove inference rule didn't consider the negative case for (count e B)

(= (count e (difference_remove A B)) (ite (=  (count e B) 0) (count e A) 0))) ; not enough
(= (count e (difference_remove A B)) (ite (<= (count e B) 0) (count e A) 0))) ; the right rule

10 files changed, 162 insertions, 108 deletions

diff --git a/src/theory/bags/bags_rewriter.cpp b/src/theory/bags/bags_rewriter.cpp
index 0be83fb13..2170175b4 100644
--- a/src/theory/bags/bags_rewriter.cpp
+++ b/src/theory/bags/bags_rewriter.cpp
@@ -173,8 +173,11 @@ BagsRewriteResponse BagsRewriter::rewriteBagCount(const TNode& n) const
   }
   if (n[1].getKind() == MK_BAG && n[0] == n[1][0])
   {
-    // (bag.count x (mkBag x c) = c
-    return BagsRewriteResponse(n[1][1], Rewrite::COUNT_MK_BAG);
+    // (bag.count x (mkBag x c)) = (ite (>= c 1) c 0)
+    Node c = n[1][1];
+    Node geq = d_nm->mkNode(GEQ, c, d_one);
+    Node ite = d_nm->mkNode(ITE, geq, c, d_zero);
+    return BagsRewriteResponse(ite, Rewrite::COUNT_MK_BAG);
   }
   return BagsRewriteResponse(n, Rewrite::NONE);
 }
diff --git a/src/theory/bags/bags_rewriter.h b/src/theory/bags/bags_rewriter.h
index eb5c9f9ab..3edd5af32 100644
--- a/src/theory/bags/bags_rewriter.h
+++ b/src/theory/bags/bags_rewriter.h
@@ -56,6 +56,7 @@ class BagsRewriter : public TheoryRewriter
    * See the rewrite rules for these kinds below.
    */
   RewriteResponse preRewrite(TNode n) override;
+
  private:
   /**
    * rewrites for n include:
@@ -81,7 +82,7 @@ class BagsRewriter : public TheoryRewriter
   /**
    * rewrites for n include:
    * - (bag.count x emptybag) = 0
-   * - (bag.count x (bag x c) = c
+   * - (bag.count x (bag x c)) = (ite (>= c 1) c 0)
    * - otherwise = n
    */
   BagsRewriteResponse rewriteBagCount(const TNode& n) const;
@@ -214,7 +215,8 @@ class BagsRewriter : public TheoryRewriter
   /**
    *  rewrites for n include:
    *  - (bag.map (lambda ((x U)) t) emptybag) = emptybag
-   *  - (bag.map (lambda ((x U)) t) (bag y z)) = (bag (apply (lambda ((x U)) t) y) z)
+   *  - (bag.map (lambda ((x U)) t) (bag y z)) = (bag (apply (lambda ((x U)) t)
+   * y) z)
    *  - (bag.map (lambda ((x U)) t) (union_disjoint A B)) =
    *       (union_disjoint
    *          (bag ((lambda ((x U)) t) "a") 3)
diff --git a/src/theory/bags/inference_generator.cpp b/src/theory/bags/inference_generator.cpp
index 734572f7c..e88a7e0ca 100644
--- a/src/theory/bags/inference_generator.cpp
+++ b/src/theory/bags/inference_generator.cpp
@@ -24,6 +24,8 @@
 #include "theory/uf/equality_engine.h"
 #include "util/rational.h"
 
+using namespace cvc5::kind;
+
 namespace cvc5 {
 namespace theory {
 namespace bags {
@@ -44,38 +46,49 @@ InferInfo InferenceGenerator::nonNegativeCount(Node n, Node e)
   Assert(e.getType() == n.getType().getBagElementType());
 
   InferInfo inferInfo(d_im, InferenceId::BAGS_NON_NEGATIVE_COUNT);
-  Node count = d_nm->mkNode(kind::BAG_COUNT, e, n);
+  Node count = d_nm->mkNode(BAG_COUNT, e, n);
 
-  Node gte = d_nm->mkNode(kind::GEQ, count, d_zero);
+  Node gte = d_nm->mkNode(GEQ, count, d_zero);
   inferInfo.d_conclusion = gte;
   return inferInfo;
 }
 
 InferInfo InferenceGenerator::mkBag(Node n, Node e)
 {
-  Assert(n.getKind() == kind::MK_BAG);
+  Assert(n.getKind() == MK_BAG);
   Assert(e.getType() == n.getType().getBagElementType());
 
-  if (n[0] == e)
+  Node x = n[0];
+  Node c = n[1];
+  Node geq = d_nm->mkNode(GEQ, c, d_one);
+  if (d_state->areEqual(e, x))
+  {
+    // (= (bag.count e skolem) (ite (>= c 1) c 0)))
+    InferInfo inferInfo(d_im, InferenceId::BAGS_MK_BAG_SAME_ELEMENT);
+    Node skolem = getSkolem(n, inferInfo);
+    Node count = getMultiplicityTerm(e, skolem);
+    Node ite = d_nm->mkNode(ITE, geq, c, d_zero);
+    inferInfo.d_conclusion = count.eqNode(ite);
+    return inferInfo;
+  }
+  if (d_state->areDisequal(e, x))
   {
-    // TODO issue #78: refactor this with BagRewriter
-    // (=> true (= (bag.count e (bag e c)) c))
+    //(= (bag.count e skolem) 0))
     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]);
+    inferInfo.d_conclusion = count.eqNode(d_zero);
     return inferInfo;
   }
   else
   {
-    // (=>
-    //   true
-    //   (= (bag.count e (bag x c)) (ite (= e x) c 0)))
+    // (= (bag.count e skolem) (ite (and (= e x) (>= c 1)) c 0)))
     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);
-    Node ite = d_nm->mkNode(kind::ITE, same, n[1], d_zero);
+    Node same = d_nm->mkNode(EQUAL, e, x);
+    Node andNode = same.andNode(geq);
+    Node ite = d_nm->mkNode(ITE, andNode, c, d_zero);
     Node equal = count.eqNode(ite);
     inferInfo.d_conclusion = equal;
     return inferInfo;
@@ -110,7 +123,7 @@ typedef expr::Attribute<BagsDeqAttributeId, Node> BagsDeqAttribute;
 
 InferInfo InferenceGenerator::bagDisequality(Node n)
 {
-  Assert(n.getKind() == kind::EQUAL && n[0].getType().isBag());
+  Assert(n.getKind() == EQUAL && n[0].getType().isBag());
 
   Node A = n[0];
   Node B = n[1];
@@ -145,7 +158,7 @@ Node InferenceGenerator::getSkolem(Node& n, InferInfo& inferInfo)
 
 InferInfo InferenceGenerator::empty(Node n, Node e)
 {
-  Assert(n.getKind() == kind::EMPTYBAG);
+  Assert(n.getKind() == EMPTYBAG);
   Assert(e.getType() == n.getType().getBagElementType());
 
   InferInfo inferInfo(d_im, InferenceId::BAGS_EMPTY);
@@ -159,7 +172,7 @@ InferInfo InferenceGenerator::empty(Node n, Node e)
 
 InferInfo InferenceGenerator::unionDisjoint(Node n, Node e)
 {
-  Assert(n.getKind() == kind::UNION_DISJOINT && n[0].getType().isBag());
+  Assert(n.getKind() == UNION_DISJOINT && n[0].getType().isBag());
   Assert(e.getType() == n[0].getType().getBagElementType());
 
   Node A = n[0];
@@ -172,7 +185,7 @@ InferInfo InferenceGenerator::unionDisjoint(Node n, Node e)
   Node skolem = getSkolem(n, inferInfo);
   Node count = getMultiplicityTerm(e, skolem);
 
-  Node sum = d_nm->mkNode(kind::PLUS, countA, countB);
+  Node sum = d_nm->mkNode(PLUS, countA, countB);
   Node equal = count.eqNode(sum);
 
   inferInfo.d_conclusion = equal;
@@ -181,7 +194,7 @@ InferInfo InferenceGenerator::unionDisjoint(Node n, Node e)
 
 InferInfo InferenceGenerator::unionMax(Node n, Node e)
 {
-  Assert(n.getKind() == kind::UNION_MAX && n[0].getType().isBag());
+  Assert(n.getKind() == UNION_MAX && n[0].getType().isBag());
   Assert(e.getType() == n[0].getType().getBagElementType());
 
   Node A = n[0];
@@ -194,8 +207,8 @@ InferInfo InferenceGenerator::unionMax(Node n, Node e)
   Node skolem = getSkolem(n, inferInfo);
   Node count = getMultiplicityTerm(e, skolem);
 
-  Node gt = d_nm->mkNode(kind::GT, countA, countB);
-  Node max = d_nm->mkNode(kind::ITE, gt, countA, countB);
+  Node gt = d_nm->mkNode(GT, countA, countB);
+  Node max = d_nm->mkNode(ITE, gt, countA, countB);
   Node equal = count.eqNode(max);
 
   inferInfo.d_conclusion = equal;
@@ -204,7 +217,7 @@ InferInfo InferenceGenerator::unionMax(Node n, Node e)
 
 InferInfo InferenceGenerator::intersection(Node n, Node e)
 {
-  Assert(n.getKind() == kind::INTERSECTION_MIN && n[0].getType().isBag());
+  Assert(n.getKind() == INTERSECTION_MIN && n[0].getType().isBag());
   Assert(e.getType() == n[0].getType().getBagElementType());
 
   Node A = n[0];
@@ -216,8 +229,8 @@ InferInfo InferenceGenerator::intersection(Node n, Node e)
   Node skolem = getSkolem(n, inferInfo);
   Node count = getMultiplicityTerm(e, skolem);
 
-  Node lt = d_nm->mkNode(kind::LT, countA, countB);
-  Node min = d_nm->mkNode(kind::ITE, lt, countA, countB);
+  Node lt = d_nm->mkNode(LT, countA, countB);
+  Node min = d_nm->mkNode(ITE, lt, countA, countB);
   Node equal = count.eqNode(min);
   inferInfo.d_conclusion = equal;
   return inferInfo;
@@ -225,7 +238,7 @@ InferInfo InferenceGenerator::intersection(Node n, Node e)
 
 InferInfo InferenceGenerator::differenceSubtract(Node n, Node e)
 {
-  Assert(n.getKind() == kind::DIFFERENCE_SUBTRACT && n[0].getType().isBag());
+  Assert(n.getKind() == DIFFERENCE_SUBTRACT && n[0].getType().isBag());
   Assert(e.getType() == n[0].getType().getBagElementType());
 
   Node A = n[0];
@@ -237,9 +250,9 @@ InferInfo InferenceGenerator::differenceSubtract(Node n, Node e)
   Node skolem = getSkolem(n, inferInfo);
   Node count = getMultiplicityTerm(e, skolem);
 
-  Node subtract = d_nm->mkNode(kind::MINUS, countA, countB);
-  Node gte = d_nm->mkNode(kind::GEQ, countA, countB);
-  Node difference = d_nm->mkNode(kind::ITE, gte, subtract, d_zero);
+  Node subtract = d_nm->mkNode(MINUS, countA, countB);
+  Node gte = d_nm->mkNode(GEQ, countA, countB);
+  Node difference = d_nm->mkNode(ITE, gte, subtract, d_zero);
   Node equal = count.eqNode(difference);
   inferInfo.d_conclusion = equal;
   return inferInfo;
@@ -247,7 +260,7 @@ InferInfo InferenceGenerator::differenceSubtract(Node n, Node e)
 
 InferInfo InferenceGenerator::differenceRemove(Node n, Node e)
 {
-  Assert(n.getKind() == kind::DIFFERENCE_REMOVE && n[0].getType().isBag());
+  Assert(n.getKind() == DIFFERENCE_REMOVE && n[0].getType().isBag());
   Assert(e.getType() == n[0].getType().getBagElementType());
 
   Node A = n[0];
@@ -260,8 +273,8 @@ InferInfo InferenceGenerator::differenceRemove(Node n, Node e)
   Node skolem = getSkolem(n, inferInfo);
   Node count = getMultiplicityTerm(e, skolem);
 
-  Node notInB = d_nm->mkNode(kind::EQUAL, countB, d_zero);
-  Node difference = d_nm->mkNode(kind::ITE, notInB, countA, d_zero);
+  Node notInB = d_nm->mkNode(LEQ, countB, d_zero);
+  Node difference = d_nm->mkNode(ITE, notInB, countA, d_zero);
   Node equal = count.eqNode(difference);
   inferInfo.d_conclusion = equal;
   return inferInfo;
@@ -269,7 +282,7 @@ InferInfo InferenceGenerator::differenceRemove(Node n, Node e)
 
 InferInfo InferenceGenerator::duplicateRemoval(Node n, Node e)
 {
-  Assert(n.getKind() == kind::DUPLICATE_REMOVAL && n[0].getType().isBag());
+  Assert(n.getKind() == DUPLICATE_REMOVAL && n[0].getType().isBag());
   Assert(e.getType() == n[0].getType().getBagElementType());
 
   Node A = n[0];
@@ -279,8 +292,8 @@ InferInfo InferenceGenerator::duplicateRemoval(Node n, Node e)
   Node skolem = getSkolem(n, inferInfo);
   Node count = getMultiplicityTerm(e, skolem);
 
-  Node gte = d_nm->mkNode(kind::GEQ, countA, d_one);
-  Node ite = d_nm->mkNode(kind::ITE, gte, d_one, d_zero);
+  Node gte = d_nm->mkNode(GEQ, countA, d_one);
+  Node ite = d_nm->mkNode(ITE, gte, d_one, d_zero);
   Node equal = count.eqNode(ite);
   inferInfo.d_conclusion = equal;
   return inferInfo;
@@ -288,14 +301,14 @@ InferInfo InferenceGenerator::duplicateRemoval(Node n, Node e)
 
 Node InferenceGenerator::getMultiplicityTerm(Node element, Node bag)
 {
-  Node count = d_nm->mkNode(kind::BAG_COUNT, element, bag);
+  Node count = d_nm->mkNode(BAG_COUNT, element, 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.getKind() == BAG_MAP && n[1].getType().isBag());
   Assert(n[0].getType().isFunction()
          && n[0].getType().getArgTypes().size() == 1);
   Assert(e.getType() == n[0].getType().getRangeType());
@@ -316,8 +329,8 @@ std::tuple<InferInfo, Node, Node> InferenceGenerator::mapDownwards(Node n,
   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);
+  Node sum_zero = d_nm->mkNode(APPLY_UF, sum, d_zero);
+  Node baseCase = d_nm->mkNode(EQUAL, sum_zero, d_zero);
 
   // guess the size of the preimage of e
   Node preImageSize = d_sm->mkDummySkolem("preImageSize", d_nm->integerType());
@@ -325,8 +338,8 @@ std::tuple<InferInfo, Node, Node> InferenceGenerator::mapDownwards(Node n,
   // (= (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);
+  Node totalSum = d_nm->mkNode(APPLY_UF, sum, preImageSize);
+  Node totalSumEqualCountE = d_nm->mkNode(EQUAL, totalSum, countE);
 
   // (forall ((i Int))
   //        (let ((uf_i (uf i)))
@@ -347,44 +360,42 @@ std::tuple<InferInfo, Node, Node> InferenceGenerator::mapDownwards(Node n,
   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);
+  Node iList = d_nm->mkNode(BOUND_VAR_LIST, i);
+  Node jList = d_nm->mkNode(BOUND_VAR_LIST, j);
+  Node iPlusOne = d_nm->mkNode(PLUS, i, d_one);
+  Node iMinusOne = d_nm->mkNode(MINUS, i, d_one);
+  Node uf_i = d_nm->mkNode(APPLY_UF, uf, i);
+  Node uf_j = d_nm->mkNode(APPLY_UF, uf, j);
+  Node f_uf_i = d_nm->mkNode(APPLY_UF, f, uf_i);
+  Node uf_iPlusOne = d_nm->mkNode(APPLY_UF, uf, iPlusOne);
+  Node uf_iMinusOne = d_nm->mkNode(APPLY_UF, uf, iMinusOne);
+  Node interval_i = d_nm->mkNode(
+      AND, d_nm->mkNode(GEQ, i, d_one), d_nm->mkNode(LEQ, i, preImageSize));
+  Node sum_i = d_nm->mkNode(APPLY_UF, sum, i);
+  Node sum_iPlusOne = d_nm->mkNode(APPLY_UF, sum, iPlusOne);
+  Node sum_iMinusOne = d_nm->mkNode(APPLY_UF, sum, iMinusOne);
+  Node count_iMinusOne = d_nm->mkNode(BAG_COUNT, uf_iMinusOne, A);
+  Node count_uf_i = d_nm->mkNode(BAG_COUNT, uf_i, A);
+  Node inductiveCase =
+      d_nm->mkNode(EQUAL, sum_i, d_nm->mkNode(PLUS, sum_iMinusOne, count_uf_i));
+  Node f_iEqualE = d_nm->mkNode(EQUAL, f_uf_i, e);
+  Node geqOne = d_nm->mkNode(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));
+  Node interval_j = d_nm->mkNode(
+      AND, d_nm->mkNode(LT, i, j), d_nm->mkNode(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 uf_i_equals_uf_j = d_nm->mkNode(EQUAL, uf_i, uf_j);
+  Node notEqual = d_nm->mkNode(EQUAL, uf_i, uf_j).negate();
+  Node body_j = d_nm->mkNode(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);
+      d_nm->mkNode(AND, {f_iEqualE, geqOne, inductiveCase, forAll_j});
+  Node body_i = d_nm->mkNode(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 preImageGTE_zero = d_nm->mkNode(GEQ, preImageSize, d_zero);
   Node conclusion = d_nm->mkNode(
-      kind::AND, {baseCase, totalSumEqualCountE, forAll_i, preImageGTE_zero});
+      AND, {baseCase, totalSumEqualCountE, forAll_i, preImageGTE_zero});
   inferInfo.d_conclusion = conclusion;
 
   std::map<Node, Node> m;
@@ -395,7 +406,7 @@ std::tuple<InferInfo, Node, Node> InferenceGenerator::mapDownwards(Node n,
 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.getKind() == BAG_MAP && n[1].getType().isBag());
   Assert(n[0].getType().isFunction()
          && n[0].getType().getArgTypes().size() == 1);
 
@@ -404,19 +415,16 @@ InferInfo InferenceGenerator::mapUpwards(
   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 xInA = d_nm->mkNode(GEQ, countA, d_one);
+  Node notEqual = d_nm->mkNode(EQUAL, d_nm->mkNode(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);
+  Node inRange = d_nm->mkNode(
+      AND, d_nm->mkNode(GEQ, k, d_one), d_nm->mkNode(LEQ, k, preImageSize));
+  Node equal = d_nm->mkNode(EQUAL, d_nm->mkNode(APPLY_UF, uf, k), x);
+  Node andNode = d_nm->mkNode(AND, inRange, equal);
+  Node orNode = d_nm->mkNode(OR, notEqual, andNode);
+  Node implies = d_nm->mkNode(IMPLIES, xInA, orNode);
   inferInfo.d_conclusion = implies;
   std::cout << "Upwards conclusion: " << inferInfo.d_conclusion << std::endl
             << std::endl;
diff --git a/src/theory/bags/inference_generator.h b/src/theory/bags/inference_generator.h
index ab3a84b29..3f38d05b9 100644
--- a/src/theory/bags/inference_generator.h
+++ b/src/theory/bags/inference_generator.h
@@ -50,15 +50,14 @@ class InferenceGenerator
   /**
    * @param n is (bag x c) of type (Bag E)
    * @param e is a node of type E
-   * @return an inference that represents the following implication
-   * (=>
-   *   true
-   *   (= (bag.count e skolem) c))
-   *   if e is exactly node x. Node skolem is a fresh variable equals (bag x c).
-   *   Otherwise the following inference is returned
-   * (=>
-   *   true
-   *   (= (bag.count e skolem) (ite (= e x) c 0)))
+   * @return an inference that represents the following cases:
+   * 1- e, x are in the same equivalent class, then we infer:
+   *    (= (bag.count e skolem) (ite (>= c 1) c 0)))
+   * 2- e, x are known to be disequal, then we infer:
+   *    (= (bag.count e skolem) 0))
+   * 3- if neither holds, we infer:
+   *    (= (bag.count e skolem) (ite (and (= e x) (>= c 1)) c 0)))
+   * where skolem = (bag x c) is a fresh variable
    */
   InferInfo mkBag(Node n, Node e);
   /**
@@ -146,7 +145,7 @@ class InferenceGenerator
    *   (=
    *     (count e skolem)
    *     (ite
-   *       (= (count e B) 0)
+   *       (<= (count e B) 0)
    *       (count e A)
    *       0))))
    * where skolem is a fresh variable equals (difference_remove A B)
diff --git a/src/theory/inference_id.cpp b/src/theory/inference_id.cpp
index b50522834..20c1f0cbd 100644
--- a/src/theory/inference_id.cpp
+++ b/src/theory/inference_id.cpp
@@ -104,6 +104,7 @@ const char* toString(InferenceId i)
     case InferenceId::ARRAYS_EQ_TAUTOLOGY: return "ARRAYS_EQ_TAUTOLOGY";
 
     case InferenceId::BAGS_NON_NEGATIVE_COUNT: return "BAGS_NON_NEGATIVE_COUNT";
+    case InferenceId::BAGS_MK_BAG_DIFFERENT_ELEMENT: return "BAGS_MK_BAG_DIFFERENT_ELEMENT";
     case InferenceId::BAGS_MK_BAG_SAME_ELEMENT: return "BAGS_MK_BAG_SAME_ELEMENT";
     case InferenceId::BAGS_MK_BAG: return "BAGS_MK_BAG";
     case InferenceId::BAGS_EQUALITY: return "BAGS_EQUALITY";
diff --git a/src/theory/inference_id.h b/src/theory/inference_id.h
index 3c644e545..1bc32d913 100644
--- a/src/theory/inference_id.h
+++ b/src/theory/inference_id.h
@@ -169,6 +169,7 @@ enum class InferenceId
 
   // ---------------------------------- bags theory
   BAGS_NON_NEGATIVE_COUNT,
+  BAGS_MK_BAG_DIFFERENT_ELEMENT,
   BAGS_MK_BAG_SAME_ELEMENT,
   BAGS_MK_BAG,
   BAGS_EQUALITY,
diff --git a/test/regress/CMakeLists.txt b/test/regress/CMakeLists.txt
index 3f099cf67..5d1ff9b49 100644
--- a/test/regress/CMakeLists.txt
+++ b/test/regress/CMakeLists.txt
@@ -1591,6 +1591,8 @@ set(regress_1_tests
   regress1/bags/fuzzy1.smt2
   regress1/bags/fuzzy2.smt2
   regress1/bags/fuzzy3.smt2
+  regress1/bags/fuzzy4.smt2
+  regress1/bags/fuzzy5.smt2
   regress1/bags/intersection_min1.smt2
   regress1/bags/intersection_min2.smt2
   regress1/bags/issue5759.smt2
diff --git a/test/regress/regress1/bags/fuzzy4.smt2 b/test/regress/regress1/bags/fuzzy4.smt2
new file mode 100644
index 000000000..b733a4862
--- /dev/null
+++ b/test/regress/regress1/bags/fuzzy4.smt2
@@ -0,0 +1,15 @@
+(set-logic ALL)
+(set-option :produce-models true)
+(set-info :status sat)
+(declare-fun A () (Bag (Tuple Int Int)))
+(declare-fun c () Int)
+(declare-fun d () (Tuple Int Int))
+(assert
+ (not
+  (=
+   (= A (bag d (+ c (bag.count d (union_disjoint A A)))))
+   (= A (difference_remove (bag d c) A)))))
+(assert (= A (bag (tuple 0 0) 5)))
+(assert (= c (- 5)))
+(assert (= d (tuple 0 0)))
+(check-sat)
diff --git a/test/regress/regress1/bags/fuzzy5.smt2 b/test/regress/regress1/bags/fuzzy5.smt2
new file mode 100644
index 000000000..2dea236a5
--- /dev/null
+++ b/test/regress/regress1/bags/fuzzy5.smt2
@@ -0,0 +1,21 @@
+(set-logic ALL)
+(set-option :produce-models true)
+(set-info :status sat)
+(declare-fun A () (Bag (Tuple Int Int)))
+(declare-fun c () Int)
+(declare-fun d () (Tuple Int Int))
+
+(assert
+ (let ((c_plus_1 (+ c 1)))
+   (and
+    (not
+     (= (= A (bag (tuple 0 c) (+ c c)))
+        (= A (difference_remove (bag d c) A))))
+    (not
+     (= (= A (bag (tuple 0 1) c_plus_1))
+        (= A (bag (tuple c 1) c_plus_1)))))))
+
+;(assert (= A (bag (tuple 0 1) 2)))
+;(assert (= c 1))
+;(assert (= d (tuple 0 1)))
+(check-sat)
\ No newline at end of file
diff --git a/test/unit/theory/theory_bags_rewriter_white.cpp b/test/unit/theory/theory_bags_rewriter_white.cpp
index 025d5aec7..7250f581c 100644
--- a/test/unit/theory/theory_bags_rewriter_white.cpp
+++ b/test/unit/theory/theory_bags_rewriter_white.cpp
@@ -165,25 +165,29 @@ TEST_F(TestTheoryWhiteBagsRewriter, mkBag_variable_element)
 
 TEST_F(TestTheoryWhiteBagsRewriter, bag_count)
 {
-  int n = 3;
+  Node zero = d_nodeManager->mkConst(Rational(0));
+  Node one = d_nodeManager->mkConst(Rational(1));
+  Node three = d_nodeManager->mkConst(Rational(3));
   Node skolem =
       d_skolemManager->mkDummySkolem("x", d_nodeManager->stringType());
   Node emptyBag = d_nodeManager->mkConst(
       EmptyBag(d_nodeManager->mkBagType(skolem.getType())));
-  Node bag = d_nodeManager->mkBag(
-      d_nodeManager->stringType(), skolem, d_nodeManager->mkConst(Rational(n)));
 
   // (bag.count x emptybag) = 0
   Node n1 = d_nodeManager->mkNode(BAG_COUNT, skolem, emptyBag);
   RewriteResponse response1 = d_rewriter->postRewrite(n1);
   ASSERT_TRUE(response1.d_status == REWRITE_AGAIN_FULL
-              && response1.d_node == d_nodeManager->mkConst(Rational(0)));
+              && response1.d_node == zero);
 
-  // (bag.count x (mkBag x c) = c where c > 0 is a constant
+  // (bag.count x (mkBag x c) = (ite (>= c 1) c 0)
+  Node bag = d_nodeManager->mkBag(d_nodeManager->stringType(), skolem, three);
   Node n2 = d_nodeManager->mkNode(BAG_COUNT, skolem, bag);
   RewriteResponse response2 = d_rewriter->postRewrite(n2);
+
+  Node geq = d_nodeManager->mkNode(GEQ, three, one);
+  Node ite = d_nodeManager->mkNode(ITE, geq, three, zero);
   ASSERT_TRUE(response2.d_status == REWRITE_AGAIN_FULL
-              && response2.d_node == d_nodeManager->mkConst(Rational(n)));
+              && response2.d_node == ite);
 }
 
 TEST_F(TestTheoryWhiteBagsRewriter, duplicate_removal)
@@ -715,12 +719,10 @@ TEST_F(TestTheoryWhiteBagsRewriter, map)
   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 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());
@@ -729,7 +731,7 @@ TEST_F(TestTheoryWhiteBagsRewriter, map)
   //   (bag "" 7))
   Node n2 = d_nodeManager->mkNode(BAG_MAP, lambda, unionDisjointAB);
 
-  Node rewritten = Rewriter:: rewrite(n2);
+  Node rewritten = Rewriter::rewrite(n2);
   Node bag = d_nodeManager->mkBag(
       d_nodeManager->stringType(), empty, d_nodeManager->mkConst(Rational(7)));
   ASSERT_TRUE(rewritten == bag);