From 2d99d409922531c149a61db291f2b781cf2c4450 Mon Sep 17 00:00:00 2001
From: Daniel Balasubramanian <danielbala@gmail.com>
Date: Fri, 13 Sep 2024 13:44:18 -0500
Subject: [PATCH] Assert equality between all elements of intermediate
 comprehension terms in symbolic counting constraints. Fixes #64

---
 Src/Core/Common/Symbols/OpLibrary.cs      |  5 +++++
 Src/Core/Solver/Execution/TermEncIndex.cs | 19 +++++++++++++++----
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/Src/Core/Common/Symbols/OpLibrary.cs b/Src/Core/Common/Symbols/OpLibrary.cs
index 7c90250..c474579 100644
--- a/Src/Core/Common/Symbols/OpLibrary.cs
+++ b/Src/Core/Common/Symbols/OpLibrary.cs
@@ -1603,6 +1603,11 @@ internal static Term SymEvaluator_Count(SymExecuter facts, Bindable[] values)
             bool wasAdded;
             var res = facts.Query(values[0].Binding, out nResults);
 
+            if (nResults == 0)
+            {
+                return facts.Index.MkCnst(new Rational(nResults), out wasAdded);
+            }
+
             int baseCount = 0; // always 0 for now
             List<Term> realTerms = new List<Term>();
             List<Term> fakeTerms = new List<Term>();
diff --git a/Src/Core/Solver/Execution/TermEncIndex.cs b/Src/Core/Solver/Execution/TermEncIndex.cs
index a3f057f..40701d3 100644
--- a/Src/Core/Solver/Execution/TermEncIndex.cs
+++ b/Src/Core/Solver/Execution/TermEncIndex.cs
@@ -340,14 +340,19 @@ private Z3Expr GetSymCountExpr(SymExecuter facts, Term x, IEnumerable<Z3Expr> ch
             Term comprTerms = facts.GetSymbolicCountTerm(x.Args[2], index);
 
             int numTerms = comprTerms.Args.Count() - 2;
-            Z3Expr[] termExprs = new Z3Expr[numTerms];
+            List<Z3Expr>[] termExprs = new List<Z3Expr>[numTerms];
             Z3BoolExpr[] boolExprs = new Z3BoolExpr[numTerms];
             Term normalized;
 
             for (int i = 0; i < comprTerms.Args.Count() - 2; i++)
             {
-                var e1Term = comprTerms.Args[i + 2].Args[0];
-                termExprs[i] = GetTerm(e1Term, out normalized, facts);
+                termExprs[i] = new List<Z3Expr>();
+                for (int j = 0; j < comprTerms.Args[i + 2].Args.Length; j++)
+                {
+                    var t = comprTerms.Args[i + 2].Args[j];
+                    termExprs[i].Add(GetTerm(t, out normalized, facts));
+                }
+
                 boolExprs[i] = facts.GetSideConstraints(comprTerms.Args[i + 2]); // TODO: check if we need comprTerms.Args[i + 2].Args[0] here
             }
 
@@ -369,7 +374,13 @@ private Z3Expr GetSymCountExpr(SymExecuter facts, Term x, IEnumerable<Z3Expr> ch
                 List<Z3BoolExpr> currExprs = new List<Z3BoolExpr>();
                 for (int j = i + 1; j < numTerms; j++)
                 {
-                    currExpr = facts.Solver.Context.MkEq(termExprs[i], termExprs[j]);
+                    currExpr = facts.Solver.Context.MkEq(termExprs[i].ElementAt(0), termExprs[j].ElementAt(0));
+                    for (int k = 1; k < termExprs[i].Count; k++)
+                    {
+                        var tempExpr = facts.Solver.Context.MkEq(termExprs[i].ElementAt(k), termExprs[j].ElementAt(k));
+                        currExpr = facts.Solver.Context.MkAnd(currExpr, tempExpr);
+                    }
+
                     currExpr = facts.Solver.Context.MkAnd(currExpr, boolExprs[i]);
                     currExpr = facts.Solver.Context.MkAnd(currExpr, boolExprs[j]);
                     currExprs.Add(currExpr);