Connect the equality solver to theory arith (#6894)

--arith-eq-solver is a new option to enable the equality solver in arithmetic, disabled by default.

This PR connects the equality solver to TheoryArith when this option is enabled.

This is in preparation for the central equality engine.

4 files changed, 55 insertions, 4 deletions

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 8053ba25e..b6bf49ed8 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -375,6 +375,8 @@ libcvc5_add_sources(
   theory/arith/dio_solver.h
   theory/arith/dual_simplex.cpp
   theory/arith/dual_simplex.h
+  theory/arith/equality_solver.cpp
+  theory/arith/equality_solver.h
   theory/arith/error_set.cpp
   theory/arith/error_set.h
   theory/arith/fc_simplex.cpp
diff --git a/src/options/arith_options.toml b/src/options/arith_options.toml
index 43614fd60..8766672a5 100644
--- a/src/options/arith_options.toml
+++ b/src/options/arith_options.toml
@@ -588,3 +588,10 @@ name   = "Arithmetic Theory"
   default    = "false"
   help       = "whether to use ICP-style propagations for non-linear arithmetic"
 
+[[option]]
+  name       = "arithEqSolver"
+  category   = "regular"
+  long       = "arith-eq-solver"
+  type       = "bool"
+  default    = "false"
+  help       = "whether to use the equality solver in the theory of arithmetic"
diff --git a/src/theory/arith/theory_arith.cpp b/src/theory/arith/theory_arith.cpp
index 980714d53..2c7187089 100644
--- a/src/theory/arith/theory_arith.cpp
+++ b/src/theory/arith/theory_arith.cpp
@@ -20,6 +20,7 @@
 #include "proof/proof_rule.h"
 #include "smt/smt_statistics_registry.h"
 #include "theory/arith/arith_rewriter.h"
+#include "theory/arith/equality_solver.h"
 #include "theory/arith/infer_bounds.h"
 #include "theory/arith/nl/nonlinear_extension.h"
 #include "theory/arith/theory_arith_private.h"
@@ -47,6 +48,7 @@ TheoryArith::TheoryArith(context::Context* c,
       d_im(*this, d_astate, pnm),
       d_ppre(c, pnm),
       d_bab(d_astate, d_im, d_ppre, pnm),
+      d_eqSolver(nullptr),
       d_internal(new TheoryArithPrivate(*this, c, u, d_bab, pnm)),
       d_nonlinearExtension(nullptr),
       d_opElim(pnm, logicInfo),
@@ -58,6 +60,11 @@ TheoryArith::TheoryArith(context::Context* c,
   // indicate we are using the theory state object and inference manager
   d_theoryState = &d_astate;
   d_inferManager = &d_im;
+
+  if (options::arithEqSolver())
+  {
+    d_eqSolver.reset(new EqualitySolver(d_astate, d_im));
+  }
 }
 
 TheoryArith::~TheoryArith(){
@@ -73,6 +80,14 @@ ProofRuleChecker* TheoryArith::getProofChecker()
 
 bool TheoryArith::needsEqualityEngine(EeSetupInfo& esi)
 {
+  // if the equality solver is enabled, then it is responsible for setting
+  // up the equality engine
+  if (d_eqSolver != nullptr)
+  {
+    return d_eqSolver->needsEqualityEngine(esi);
+  }
+  // otherwise, the linear arithmetic solver is responsible for setting up
+  // the equality engine
   return d_internal->needsEqualityEngine(esi);
 }
 void TheoryArith::finishInit()
@@ -94,7 +109,11 @@ void TheoryArith::finishInit()
     d_nonlinearExtension.reset(
         new nl::NonlinearExtension(*this, d_astate, d_equalityEngine, d_pnm));
   }
-  // finish initialize internally
+  if (d_eqSolver != nullptr)
+  {
+    d_eqSolver->finishInit();
+  }
+  // finish initialize in the old linear solver
   d_internal->finishInit();
 }
 
@@ -184,10 +203,18 @@ bool TheoryArith::preNotifyFact(
   Trace("arith-check") << "TheoryArith::preNotifyFact: " << fact
                        << ", isPrereg=" << isPrereg
                        << ", isInternal=" << isInternal << std::endl;
-  d_internal->preNotifyFact(atom, pol, fact);
   // We do not assert to the equality engine of arithmetic in the standard way,
   // hence we return "true" to indicate we are finished with this fact.
-  return true;
+  bool ret = true;
+  if (d_eqSolver != nullptr)
+  {
+    // the equality solver may indicate ret = false, after which the assertion
+    // will be asserted to the equality engine in the default way.
+    ret = d_eqSolver->preNotifyFact(atom, pol, fact, isPrereg, isInternal);
+  }
+  // we also always also notify the internal solver
+  d_internal->preNotifyFact(atom, pol, fact);
+  return ret;
 }
 
 bool TheoryArith::needsCheckLastEffort() {
@@ -198,7 +225,19 @@ bool TheoryArith::needsCheckLastEffort() {
   return false;
 }
 
-TrustNode TheoryArith::explain(TNode n) { return d_internal->explain(n); }
+TrustNode TheoryArith::explain(TNode n)
+{
+  if (d_eqSolver != nullptr)
+  {
+    // if the equality solver has an explanation for it, use it
+    TrustNode texp = d_eqSolver->explain(n);
+    if (!texp.isNull())
+    {
+      return texp;
+    }
+  }
+  return d_internal->explain(n);
+}
 
 void TheoryArith::propagate(Effort e) {
   d_internal->propagate(e);
diff --git a/src/theory/arith/theory_arith.h b/src/theory/arith/theory_arith.h
index decbf93bd..ee99f44e8 100644
--- a/src/theory/arith/theory_arith.h
+++ b/src/theory/arith/theory_arith.h
@@ -33,6 +33,7 @@ namespace nl {
 class NonlinearExtension;
 }
 
+class EqualitySolver;
 class TheoryArithPrivate;
 
 /**
@@ -141,6 +142,8 @@ class TheoryArith : public Theory {
   PreprocessRewriteEq d_ppre;
   /** The branch and bound utility */
   BranchAndBound d_bab;
+  /** The equality solver */
+  std::unique_ptr<EqualitySolver> d_eqSolver;
   /** The (old) linear arithmetic solver */
   TheoryArithPrivate* d_internal;