From ec1abb0ba86ac06c955848f718fa70d3ffe8e40d Mon Sep 17 00:00:00 2001
From: Andrew Reynolds <andrew.j.reynolds@gmail.com>
Date: Sun, 25 Jul 2021 09:54:28 -0500
Subject: [PATCH] Refactor equality engine setup for arithmetic congruence
 manager (#6902)

Allows congruence manager to use its own (non-official) equality engine if both it and the arithmetic equality solver are enabled.
---
 src/theory/arith/congruence_manager.cpp   | 34 ++++++++++++++++++-----
 src/theory/arith/congruence_manager.h     |  8 ++++--
 src/theory/arith/theory_arith_private.cpp | 14 +++++++---
 3 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/src/theory/arith/congruence_manager.cpp b/src/theory/arith/congruence_manager.cpp
index 9e7202f1d..3a35319ed 100644
--- a/src/theory/arith/congruence_manager.cpp
+++ b/src/theory/arith/congruence_manager.cpp
@@ -71,24 +71,44 @@ ArithCongruenceManager::~ArithCongruenceManager() {}
 
 bool ArithCongruenceManager::needsEqualityEngine(EeSetupInfo& esi)
 {
+  Assert(!options::arithEqSolver());
   esi.d_notify = &d_notify;
-  esi.d_name = "theory::arith::ArithCongruenceManager";
+  esi.d_name = "arithCong::ee";
   return true;
 }
 
-void ArithCongruenceManager::finishInit(eq::EqualityEngine* ee,
-                                        eq::ProofEqEngine* pfee)
+void ArithCongruenceManager::finishInit(eq::EqualityEngine* ee)
 {
-  Assert(ee != nullptr);
-  d_ee = ee;
+  if (options::arithEqSolver())
+  {
+    // use our own copy
+    d_allocEe.reset(
+        new eq::EqualityEngine(d_notify, d_satContext, "arithCong::ee", true));
+    d_ee = d_allocEe.get();
+    if (d_pnm != nullptr)
+    {
+      // allocate an internal proof equality engine
+      d_allocPfee.reset(
+          new eq::ProofEqEngine(d_satContext, d_userContext, *d_ee, d_pnm));
+      d_ee->setProofEqualityEngine(d_allocPfee.get());
+    }
+  }
+  else
+  {
+    Assert(ee != nullptr);
+    // otherwise, we use the official one
+    d_ee = ee;
+  }
+  // set the congruence kinds on the separate equality engine
   d_ee->addFunctionKind(kind::NONLINEAR_MULT);
   d_ee->addFunctionKind(kind::EXPONENTIAL);
   d_ee->addFunctionKind(kind::SINE);
   d_ee->addFunctionKind(kind::IAND);
   d_ee->addFunctionKind(kind::POW2);
+  // the proof equality engine is the one from the equality engine
+  d_pfee = d_ee->getProofEqualityEngine();
   // have proof equality engine only if proofs are enabled
-  Assert(isProofEnabled() == (pfee != nullptr));
-  d_pfee = pfee;
+  Assert(isProofEnabled() == (d_pfee != nullptr));
 }
 
 ArithCongruenceManager::Statistics::Statistics()
diff --git a/src/theory/arith/congruence_manager.h b/src/theory/arith/congruence_manager.h
index 4ab0b313b..2c59a405f 100644
--- a/src/theory/arith/congruence_manager.h
+++ b/src/theory/arith/congruence_manager.h
@@ -110,6 +110,8 @@ private:
 
   /** The equality engine being used by this class */
   eq::EqualityEngine* d_ee;
+  /** The equality engine we allocated */
+  std::unique_ptr<eq::EqualityEngine> d_allocEe;
   /** The sat context */
   context::Context* d_satContext;
   /** The user context */
@@ -143,6 +145,8 @@ private:
 
   /** Pointer to the proof equality engine of TheoryArith */
   theory::eq::ProofEqEngine* d_pfee;
+  /** The proof equality engine we allocated */
+  std::unique_ptr<eq::ProofEqEngine> d_allocPfee;
 
   /** Raise a conflict node `conflict` to the theory of arithmetic.
    *
@@ -240,9 +244,9 @@ private:
   bool needsEqualityEngine(EeSetupInfo& esi);
   /**
    * Finish initialize. This class is instructed by TheoryArithPrivate to use
-   * the equality engine ee and proof equality engine pfee.
+   * the equality engine ee.
    */
-  void finishInit(eq::EqualityEngine* ee, eq::ProofEqEngine* pfee);
+  void finishInit(eq::EqualityEngine* ee);
   //--------------------------------- end initialization
 
   /**
diff --git a/src/theory/arith/theory_arith_private.cpp b/src/theory/arith/theory_arith_private.cpp
index fb9de9641..d7ae08379 100644
--- a/src/theory/arith/theory_arith_private.cpp
+++ b/src/theory/arith/theory_arith_private.cpp
@@ -181,14 +181,20 @@ TheoryArithPrivate::~TheoryArithPrivate(){
 
 bool TheoryArithPrivate::needsEqualityEngine(EeSetupInfo& esi)
 {
+  if (!d_cmEnabled)
+  {
+    return false;
+  }
   return d_congruenceManager.needsEqualityEngine(esi);
 }
 void TheoryArithPrivate::finishInit()
 {
-  eq::EqualityEngine* ee = d_containing.getEqualityEngine();
-  eq::ProofEqEngine* pfee = d_containing.getProofEqEngine();
-  Assert(ee != nullptr);
-  d_congruenceManager.finishInit(ee, pfee);
+  if (d_cmEnabled)
+  {
+    eq::EqualityEngine* ee = d_containing.getEqualityEngine();
+    Assert(ee != nullptr);
+    d_congruenceManager.finishInit(ee);
+  }
 }
 
 static bool contains(const ConstraintCPVec& v, ConstraintP con){
-- 
2.30.2