More flexible design for model manager distributed (#4976)

This PR makes it so that the distributed model manager can be used independently of the architecture for equality engine management for theories, meaning the choice of using a separate equality engine for the model can be done in both the current distributed architecture, or the proposed central architecture (where there would be 2 equality engines, the central one and the model one). The "central model manager" on the other hand will only be combined with the central architecture. This will make it easier to test the centralized equality engine manager, since all things related to model construction can be preserved when using a central architecture.

This PR moves some of the responsibilities from the (distributed) equality engine manager to the distributed model manager, including the management of the context of the model equality engine and the allocation of its equality engine.

This PR is not intended to make any behavior changes to the current architecture.

11 files changed, 118 insertions, 155 deletions

diff --git a/src/theory/combination_engine.cpp b/src/theory/combination_engine.cpp
index 5c9e6713b..7ff2848df 100644
--- a/src/theory/combination_engine.cpp
+++ b/src/theory/combination_engine.cpp
@@ -18,8 +18,6 @@
 #include "theory/care_graph.h"
 #include "theory/ee_manager_distributed.h"
 #include "theory/model_manager_distributed.h"
-#include "theory/shared_terms_database.h"
-#include "theory/term_registration_visitor.h"
 #include "theory/theory_engine.h"
 
 namespace CVC4 {
@@ -46,11 +44,9 @@ void CombinationEngine::finishInit()
   if (options::eeMode() == options::EqEngineMode::DISTRIBUTED)
   {
     // make the distributed equality engine manager
-    std::unique_ptr<EqEngineManagerDistributed> eeDistributed(
-        new EqEngineManagerDistributed(d_te));
+    d_eemanager.reset(new EqEngineManagerDistributed(d_te));
     // make the distributed model manager
-    d_mmanager.reset(new ModelManagerDistributed(d_te, *eeDistributed.get()));
-    d_eemanager = std::move(eeDistributed);
+    d_mmanager.reset(new ModelManagerDistributed(d_te, *d_eemanager.get()));
   }
   else
   {
@@ -65,13 +61,9 @@ void CombinationEngine::finishInit()
   d_eemanager->initializeTheories();
 
   Assert(d_mmanager != nullptr);
-  // initialize the model manager
-  d_mmanager->finishInit();
-
-  // initialize equality engine of the model using the equality engine manager
-  TheoryModel* m = d_mmanager->getModel();
+  // initialize the model manager, based on the notify object of this class
   eq::EqualityEngineNotify* meen = getModelEqualityEngineNotify();
-  d_eemanager->initializeModel(m, meen);
+  d_mmanager->finishInit(meen);
 }
 
 const EeTheoryInfo* CombinationEngine::getEeTheoryInfo(TheoryId tid) const
@@ -118,7 +110,7 @@ void CombinationEngine::sendLemma(TrustNode trn, TheoryId atomsTo)
 
 void CombinationEngine::resetRound()
 {
-  // do nothing
+  // compute the relevant terms?
 }
 
 }  // namespace theory
diff --git a/src/theory/ee_manager.cpp b/src/theory/ee_manager.cpp
index bec355e7d..ad6892e31 100644
--- a/src/theory/ee_manager.cpp
+++ b/src/theory/ee_manager.cpp
@@ -14,9 +14,13 @@
 
 #include "theory/ee_manager.h"
 
+#include "theory/theory_model.h"
+
 namespace CVC4 {
 namespace theory {
 
+EqEngineManager::EqEngineManager(TheoryEngine& te) : d_te(te) {}
+
 const EeTheoryInfo* EqEngineManager::getEeTheoryInfo(TheoryId tid) const
 {
   std::map<TheoryId, EeTheoryInfo>::const_iterator it = d_einfo.find(tid);
@@ -27,5 +31,17 @@ const EeTheoryInfo* EqEngineManager::getEeTheoryInfo(TheoryId tid) const
   return nullptr;
 }
 
+eq::EqualityEngine* EqEngineManager::allocateEqualityEngine(EeSetupInfo& esi,
+                                                            context::Context* c)
+{
+  if (esi.d_notify != nullptr)
+  {
+    return new eq::EqualityEngine(
+        *esi.d_notify, c, esi.d_name, esi.d_constantsAreTriggers);
+  }
+  // the theory doesn't care about explicit notifications
+  return new eq::EqualityEngine(c, esi.d_name, esi.d_constantsAreTriggers);
+}
+
 }  // namespace theory
 }  // namespace CVC4
diff --git a/src/theory/ee_manager.h b/src/theory/ee_manager.h
index aa4bcc66c..14175037e 100644
--- a/src/theory/ee_manager.h
+++ b/src/theory/ee_manager.h
@@ -25,6 +25,9 @@
 #include "theory/uf/equality_engine.h"
 
 namespace CVC4 {
+
+class TheoryEngine;
+
 namespace theory {
 
 /**
@@ -48,9 +51,10 @@ struct EeTheoryInfo
 class EqEngineManager
 {
  public:
+  EqEngineManager(TheoryEngine& te);
   virtual ~EqEngineManager() {}
   /**
-   * Finish initialize, called by TheoryEngine::finishInit after theory
+   * Initialize theories, called during TheoryEngine::finishInit after theory
    * objects have been created but prior to their final initialization. This
    * sets up equality engines for all theories.
    *
@@ -59,16 +63,6 @@ class EqEngineManager
    */
   virtual void initializeTheories() = 0;
   /**
-   * Finish initialize, called by TheoryEngine::finishInit after theory
-   * objects have been created but prior to their final initialization. This
-   * sets up equality engines for all theories.
-   *
-   * This method is context-independent, and is applied once during
-   * the lifetime of TheoryEngine (during finishInit).
-   */
-  virtual void initializeModel(TheoryModel* m,
-                               eq::EqualityEngineNotify* notify) = 0;
-  /**
    * Get the equality engine theory information for theory with the given id.
    */
   const EeTheoryInfo* getEeTheoryInfo(TheoryId tid) const;
@@ -80,7 +74,13 @@ class EqEngineManager
    */
   virtual eq::EqualityEngine* getCoreEqualityEngine() = 0;
 
+  /** Allocate equality engine that is context-dependent on c with info esi */
+  eq::EqualityEngine* allocateEqualityEngine(EeSetupInfo& esi,
+                                             context::Context* c);
+
  protected:
+  /** Reference to the theory engine */
+  TheoryEngine& d_te;
   /** Information related to the equality engine, per theory. */
   std::map<TheoryId, EeTheoryInfo> d_einfo;
 };
diff --git a/src/theory/ee_manager_distributed.cpp b/src/theory/ee_manager_distributed.cpp
index ea618fcae..7abd60219 100644
--- a/src/theory/ee_manager_distributed.cpp
+++ b/src/theory/ee_manager_distributed.cpp
@@ -21,14 +21,12 @@ namespace CVC4 {
 namespace theory {
 
 EqEngineManagerDistributed::EqEngineManagerDistributed(TheoryEngine& te)
-    : d_te(te), d_masterEENotify(nullptr)
+    : EqEngineManager(te), d_masterEENotify(nullptr)
 {
 }
 
 EqEngineManagerDistributed::~EqEngineManagerDistributed()
 {
-  // pop the model context which we pushed on initialization
-  d_modelEeContext.pop();
 }
 
 void EqEngineManagerDistributed::initializeTheories()
@@ -96,62 +94,16 @@ void EqEngineManagerDistributed::initializeTheories()
   }
 }
 
-void EqEngineManagerDistributed::initializeModel(
-    TheoryModel* m, eq::EqualityEngineNotify* notify)
-{
-  Assert(m != nullptr);
-  // initialize the model equality engine
-  EeSetupInfo esim;
-  if (m->needsEqualityEngine(esim))
-  {
-    // use the provided notification object
-    esim.d_notify = notify;
-    d_modelEqualityEngine.reset(
-        allocateEqualityEngine(esim, &d_modelEeContext));
-    m->setEqualityEngine(d_modelEqualityEngine.get());
-  }
-  else
-  {
-    AlwaysAssert(false) << "Expected model to use equality engine";
-  }
-  m->finishInit();
-  // We push a context during initialization since the model is cleared during
-  // collectModelInfo using pop/push.
-  d_modelEeContext.push();
-}
-
 void EqEngineManagerDistributed::MasterNotifyClass::eqNotifyNewClass(TNode t)
 {
   // adds t to the quantifiers term database
   d_quantEngine->eqNotifyNewClass(t);
 }
 
-context::Context* EqEngineManagerDistributed::getModelEqualityEngineContext()
-{
-  return &d_modelEeContext;
-}
-
-eq::EqualityEngine* EqEngineManagerDistributed::getModelEqualityEngine()
-{
-  return d_modelEqualityEngine.get();
-}
-
 eq::EqualityEngine* EqEngineManagerDistributed::getCoreEqualityEngine()
 {
   return d_masterEqualityEngine.get();
 }
 
-eq::EqualityEngine* EqEngineManagerDistributed::allocateEqualityEngine(
-    EeSetupInfo& esi, context::Context* c)
-{
-  if (esi.d_notify != nullptr)
-  {
-    return new eq::EqualityEngine(
-        *esi.d_notify, c, esi.d_name, esi.d_constantsAreTriggers);
-  }
-  // the theory doesn't care about explicit notifications
-  return new eq::EqualityEngine(c, esi.d_name, esi.d_constantsAreTriggers);
-}
-
 }  // namespace theory
 }  // namespace CVC4
diff --git a/src/theory/ee_manager_distributed.h b/src/theory/ee_manager_distributed.h
index dd4608c9a..bbdd99881 100644
--- a/src/theory/ee_manager_distributed.h
+++ b/src/theory/ee_manager_distributed.h
@@ -25,9 +25,6 @@
 #include "theory/uf/equality_engine.h"
 
 namespace CVC4 {
-
-class TheoryEngine;
-
 namespace theory {
 
 /**
@@ -55,27 +52,9 @@ class EqEngineManagerDistributed : public EqEngineManager
    * per theories and connects them to a master equality engine.
    */
   void initializeTheories() override;
-  /**
-   * Initialize model. This method allocates a new equality engine for the
-   * model.
-   */
-  void initializeModel(TheoryModel* m,
-                       eq::EqualityEngineNotify* notify) override;
-  /**
-   * Get the model equality engine context. This is a dummy context that is
-   * used for clearing the contents of the model's equality engine via
-   * pop/push.
-   */
-  context::Context* getModelEqualityEngineContext();
-  /** get the model equality engine */
-  eq::EqualityEngine* getModelEqualityEngine();
   /** get the core equality engine */
   eq::EqualityEngine* getCoreEqualityEngine() override;
-
  private:
-  /** Allocate equality engine that is context-dependent on c with info esi */
-  eq::EqualityEngine* allocateEqualityEngine(EeSetupInfo& esi,
-                                             context::Context* c);
   /** notify class for master equality engine */
   class MasterNotifyClass : public theory::eq::EqualityEngineNotify
   {
@@ -106,21 +85,10 @@ class EqEngineManagerDistributed : public EqEngineManager
     /** Pointer to quantifiers engine */
     QuantifiersEngine* d_quantEngine;
   };
-  /** Reference to the theory engine */
-  TheoryEngine& d_te;
   /** The master equality engine notify class */
   std::unique_ptr<MasterNotifyClass> d_masterEENotify;
   /** The master equality engine. */
   std::unique_ptr<eq::EqualityEngine> d_masterEqualityEngine;
-  /**
-   * A dummy context for the model equality engine, so we can clear it
-   * independently of search context.
-   */
-  context::Context d_modelEeContext;
-  /**
-   * The equality engine of the model.
-   */
-  std::unique_ptr<eq::EqualityEngine> d_modelEqualityEngine;
 };
 
 }  // namespace theory
diff --git a/src/theory/model_manager.cpp b/src/theory/model_manager.cpp
index 89daa922c..530aec190 100644
--- a/src/theory/model_manager.cpp
+++ b/src/theory/model_manager.cpp
@@ -21,9 +21,12 @@
 namespace CVC4 {
 namespace theory {
 
-ModelManager::ModelManager(TheoryEngine& te)
+ModelManager::ModelManager(TheoryEngine& te, EqEngineManager& eem)
     : d_te(te),
       d_logicInfo(te.getLogicInfo()),
+      d_eem(eem),
+      d_modelEqualityEngine(nullptr),
+      d_modelEqualityEngineAlloc(nullptr),
       d_model(nullptr),
       d_modelBuilder(nullptr),
       d_modelBuilt(false),
@@ -33,7 +36,7 @@ ModelManager::ModelManager(TheoryEngine& te)
 
 ModelManager::~ModelManager() {}
 
-void ModelManager::finishInit()
+void ModelManager::finishInit(eq::EqualityEngineNotify* notify)
 {
   // construct the model
   const LogicInfo& logicInfo = d_te.getLogicInfo();
@@ -61,6 +64,7 @@ void ModelManager::finishInit()
     d_modelBuilder = d_alocModelBuilder.get();
   }
   // notice that the equality engine of the model has yet to be assigned.
+  initializeModelEqEngine(notify);
 }
 
 void ModelManager::resetModel()
@@ -209,7 +213,6 @@ void ModelManager::collectTerms(TheoryId tid,
     // only add to term set if a relevant kind
     if (irrKinds.find(k) == irrKinds.end())
     {
-      Assert(Theory::theoryOf(cur) == tid);
       termSet.insert(cur);
     }
     // traverse owned terms, don't go under quantifiers
diff --git a/src/theory/model_manager.h b/src/theory/model_manager.h
index 9511f2779..b84cd29d7 100644
--- a/src/theory/model_manager.h
+++ b/src/theory/model_manager.h
@@ -19,6 +19,7 @@
 
 #include <memory>
 
+#include "theory/ee_manager.h"
 #include "theory/logic_info.h"
 #include "theory/theory_model.h"
 #include "theory/theory_model_builder.h"
@@ -39,10 +40,16 @@ namespace theory {
 class ModelManager
 {
  public:
-  ModelManager(TheoryEngine& te);
+  ModelManager(TheoryEngine& te, EqEngineManager& eem);
   virtual ~ModelManager();
-  /** Finish initializing this class. */
-  void finishInit();
+  /**
+   * Finish initializing this class, which allocates the model, the model
+   * builder as well as the equality engine of the model. The equality engine
+   * to use is determined by the virtual method initializeModelEqEngine.
+   *
+   * @param notify The object that wants to be notified for callbacks occurring
+   */
+  void finishInit(eq::EqualityEngineNotify* notify);
   /** Reset model, called during full effort check before the model is built */
   void resetModel();
   /**
@@ -85,6 +92,12 @@ class ModelManager
   //------------------------ end finer grained control over model building
  protected:
   /**
+   * Initialize model equality engine. This is called at the end of finish
+   * init, after we have created a model object but before we have assigned it
+   * an equality engine.
+   */
+  virtual void initializeModelEqEngine(eq::EqualityEngineNotify* notify) = 0;
+  /**
    * Collect model Boolean variables.
    * This asserts the values of all boolean variables to the equality engine of
    * the model, based on their value in the prop engine.
@@ -112,6 +125,17 @@ class ModelManager
   TheoryEngine& d_te;
   /** Logic info of theory engine (cached) */
   const LogicInfo& d_logicInfo;
+  /** The equality engine manager */
+  EqEngineManager& d_eem;
+  /**
+   * A dummy context for the model equality engine, so we can clear it
+   * independently of search context.
+   */
+  context::Context d_modelEeContext;
+  /** Pointer to the equality engine of the model */
+  eq::EqualityEngine* d_modelEqualityEngine;
+  /** The equality engine of the model, if we allocated it */
+  std::unique_ptr<eq::EqualityEngine> d_modelEqualityEngineAlloc;
   /** The model object we are using */
   theory::TheoryModel* d_model;
   /** The model object we have allocated (if one exists) */
diff --git a/src/theory/model_manager_distributed.cpp b/src/theory/model_manager_distributed.cpp
index c2e1d6521..4d47329b5 100644
--- a/src/theory/model_manager_distributed.cpp
+++ b/src/theory/model_manager_distributed.cpp
@@ -20,13 +20,36 @@
 namespace CVC4 {
 namespace theory {
 
-ModelManagerDistributed::ModelManagerDistributed(
-    TheoryEngine& te, EqEngineManagerDistributed& eem)
-    : ModelManager(te), d_eem(eem)
+ModelManagerDistributed::ModelManagerDistributed(TheoryEngine& te,
+                                                 EqEngineManager& eem)
+    : ModelManager(te, eem)
 {
 }
 
-ModelManagerDistributed::~ModelManagerDistributed() {}
+ModelManagerDistributed::~ModelManagerDistributed()
+{
+  // pop the model context which we pushed on initialization
+  d_modelEeContext.pop();
+}
+
+void ModelManagerDistributed::initializeModelEqEngine(
+    eq::EqualityEngineNotify* notify)
+{
+  // initialize the model equality engine, use the provided notification object,
+  // which belongs e.g. to CombinationModelBased
+  EeSetupInfo esim;
+  esim.d_notify = notify;
+  esim.d_name = d_model->getName() + "::ee";
+  esim.d_constantsAreTriggers = false;
+  d_modelEqualityEngineAlloc.reset(
+      d_eem.allocateEqualityEngine(esim, &d_modelEeContext));
+  d_modelEqualityEngine = d_modelEqualityEngineAlloc.get();
+  // finish initializing the model
+  d_model->finishInit(d_modelEqualityEngine);
+  // We push a context during initialization since the model is cleared during
+  // collectModelInfo using pop/push.
+  d_modelEeContext.push();
+}
 
 bool ModelManagerDistributed::prepareModel()
 {
@@ -34,9 +57,8 @@ bool ModelManagerDistributed::prepareModel()
                          << std::endl;
 
   // push/pop to clear the equality engine of the model
-  context::Context* meec = d_eem.getModelEqualityEngineContext();
-  meec->pop();
-  meec->push();
+  d_modelEeContext.pop();
+  d_modelEeContext.push();
 
   // Collect model info from the theories
   Trace("model-builder") << "ModelManagerDistributed: Collect model info..."
@@ -76,6 +98,7 @@ bool ModelManagerDistributed::prepareModel()
 
 bool ModelManagerDistributed::finishBuildModel() const
 {
+  // do not use relevant terms
   if (!d_modelBuilder->buildModel(d_model))
   {
     Trace("model-builder") << "ModelManager: fail build model" << std::endl;
diff --git a/src/theory/model_manager_distributed.h b/src/theory/model_manager_distributed.h
index 2f86c0b00..2cf35d039 100644
--- a/src/theory/model_manager_distributed.h
+++ b/src/theory/model_manager_distributed.h
@@ -19,7 +19,7 @@
 
 #include <memory>
 
-#include "theory/ee_manager_distributed.h"
+#include "theory/ee_manager.h"
 #include "theory/model_manager.h"
 
 namespace CVC4 {
@@ -29,17 +29,20 @@ class TheoryEngine;
 namespace theory {
 
 /**
- * Manager for building models in a distributed architecture. Its prepare
- * model method uses collectModelInfo to assert all equalities from the equality
- * engine of each theory into the equality engine of the model. It additionally
- * uses the model equality engine context to clear the information from
- * the model's equality engine, as maintained by the distributed equality
- * engine manager.
+ * Manager for building models where the equality engine of the model is
+ * a separate instance. Notice that this manager can be used regardless of the
+ * method for managing the equality engines of the theories (which is the
+ * responsibility of the equality engine manager eem referenced by this class).
+ *
+ * Its prepare model method uses collectModelInfo to assert all equalities from
+ * the equality engine of each theory into the equality engine of the model. It
+ * additionally uses the model equality engine context to clear the information
+ * from the model's equality engine, which is maintained by this class.
  */
 class ModelManagerDistributed : public ModelManager
 {
  public:
-  ModelManagerDistributed(TheoryEngine& te, EqEngineManagerDistributed& eem);
+  ModelManagerDistributed(TheoryEngine& te, EqEngineManager& eem);
   ~ModelManagerDistributed();
 
   /** Prepare the model, as described above. */
@@ -49,13 +52,9 @@ class ModelManagerDistributed : public ModelManager
    * model, return true if successful.
    */
   bool finishBuildModel() const override;
-
  protected:
-  /**
-   * Distributed equality engine manager, which maintains the context of the
-   * model's equality engine.
-   */
-  EqEngineManagerDistributed& d_eem;
+  /** Initialize model equality engine */
+  void initializeModelEqEngine(eq::EqualityEngineNotify* notify) override;
 };
 
 }  // namespace theory
diff --git a/src/theory/theory_model.cpp b/src/theory/theory_model.cpp
index da9f28d01..53267727b 100644
--- a/src/theory/theory_model.cpp
+++ b/src/theory/theory_model.cpp
@@ -44,22 +44,10 @@ TheoryModel::TheoryModel(context::Context* c,
 
 TheoryModel::~TheoryModel() {}
 
-void TheoryModel::setEqualityEngine(eq::EqualityEngine* ee)
+void TheoryModel::finishInit(eq::EqualityEngine* ee)
 {
+  Assert(ee != nullptr);
   d_equalityEngine = ee;
-}
-
-bool TheoryModel::needsEqualityEngine(EeSetupInfo& esi)
-{
-  // no notifications
-  esi.d_name = d_name;
-  esi.d_constantsAreTriggers = false;
-  return true;
-}
-
-void TheoryModel::finishInit()
-{
-  Assert(d_equalityEngine != nullptr);
   // The kinds we are treating as function application in congruence
   d_equalityEngine->addFunctionKind(kind::APPLY_UF, false, options::ufHo());
   d_equalityEngine->addFunctionKind(kind::HO_APPLY);
@@ -772,5 +760,7 @@ std::vector< Node > TheoryModel::getFunctionsToAssign() {
   return funcs_to_assign;
 }
 
+const std::string& TheoryModel::getName() const { return d_name; }
+
 } /* namespace CVC4::theory */
 } /* namespace CVC4 */
diff --git a/src/theory/theory_model.h b/src/theory/theory_model.h
index aea3463c8..0de52b5ff 100644
--- a/src/theory/theory_model.h
+++ b/src/theory/theory_model.h
@@ -82,17 +82,10 @@ class TheoryModel : public Model
 public:
   TheoryModel(context::Context* c, std::string name, bool enableFuncModels);
   ~TheoryModel() override;
-  //---------------------------- initialization
-  /** Called to set the equality engine. */
-  void setEqualityEngine(eq::EqualityEngine* ee);
   /**
-   * Returns true if we need an equality engine, this has the same contract
-   * as Theory::needsEqualityEngine.
+   * Finish init, where ee is the equality engine the model should use.
    */
-  bool needsEqualityEngine(EeSetupInfo& esi);
-  /** Finish init */
-  void finishInit();
-  //---------------------------- end initialization
+  void finishInit(eq::EqualityEngine* ee);
 
   /** reset the model */
   virtual void reset();
@@ -355,6 +348,9 @@ public:
   */
   std::vector< Node > getFunctionsToAssign();
   //---------------------------- end function values
+  /** Get the name of this model */
+  const std::string& getName() const;
+
  protected:
   /** Unique name of this model */
   std::string d_name;