2 files changed, 31 insertions, 32 deletions

diff --git a/src/theory/sets/theory_sets.cpp b/src/theory/sets/theory_sets.cpp
index bf81099a7..fd9af488f 100644
--- a/src/theory/sets/theory_sets.cpp
+++ b/src/theory/sets/theory_sets.cpp
@@ -35,8 +35,7 @@ TheorySets::TheorySets(context::Context* c,
                        ProofNodeManager* pnm)
     : Theory(THEORY_SETS, c, u, out, valuation, logicInfo, pnm),
       d_internal(new TheorySetsPrivate(*this, c, u)),
-      d_notify(*d_internal.get()),
-      d_equalityEngine(d_notify, c, "theory::sets::ee", true)
+      d_notify(*d_internal.get())
 {
   // Do not move me to the header.
   // The constructor + destructor are not in the header as d_internal is a
@@ -54,29 +53,38 @@ TheoryRewriter* TheorySets::getTheoryRewriter()
   return d_internal->getTheoryRewriter();
 }
 
+bool TheorySets::needsEqualityEngine(EeSetupInfo& esi)
+{
+  esi.d_notify = &d_notify;
+  esi.d_name = "theory::sets::ee";
+  return true;
+}
+
 void TheorySets::finishInit()
 {
+  Assert(d_equalityEngine != nullptr);
+
   d_valuation.setUnevaluatedKind(COMPREHENSION);
   // choice is used to eliminate witness
   d_valuation.setUnevaluatedKind(WITNESS);
 
   // functions we are doing congruence over
-  d_equalityEngine.addFunctionKind(kind::SINGLETON);
-  d_equalityEngine.addFunctionKind(kind::UNION);
-  d_equalityEngine.addFunctionKind(kind::INTERSECTION);
-  d_equalityEngine.addFunctionKind(kind::SETMINUS);
-  d_equalityEngine.addFunctionKind(kind::MEMBER);
-  d_equalityEngine.addFunctionKind(kind::SUBSET);
+  d_equalityEngine->addFunctionKind(kind::SINGLETON);
+  d_equalityEngine->addFunctionKind(kind::UNION);
+  d_equalityEngine->addFunctionKind(kind::INTERSECTION);
+  d_equalityEngine->addFunctionKind(kind::SETMINUS);
+  d_equalityEngine->addFunctionKind(kind::MEMBER);
+  d_equalityEngine->addFunctionKind(kind::SUBSET);
   // relation operators
-  d_equalityEngine.addFunctionKind(PRODUCT);
-  d_equalityEngine.addFunctionKind(JOIN);
-  d_equalityEngine.addFunctionKind(TRANSPOSE);
-  d_equalityEngine.addFunctionKind(TCLOSURE);
-  d_equalityEngine.addFunctionKind(JOIN_IMAGE);
-  d_equalityEngine.addFunctionKind(IDEN);
-  d_equalityEngine.addFunctionKind(APPLY_CONSTRUCTOR);
+  d_equalityEngine->addFunctionKind(PRODUCT);
+  d_equalityEngine->addFunctionKind(JOIN);
+  d_equalityEngine->addFunctionKind(TRANSPOSE);
+  d_equalityEngine->addFunctionKind(TCLOSURE);
+  d_equalityEngine->addFunctionKind(JOIN_IMAGE);
+  d_equalityEngine->addFunctionKind(IDEN);
+  d_equalityEngine->addFunctionKind(APPLY_CONSTRUCTOR);
   // we do congruence over cardinality
-  d_equalityEngine.addFunctionKind(CARD);
+  d_equalityEngine->addFunctionKind(CARD);
 
   // finish initialization internally
   d_internal->finishInit();
@@ -198,16 +206,6 @@ bool TheorySets::isEntailed( Node n, bool pol ) {
   return d_internal->isEntailed( n, pol );
 }
 
-eq::EqualityEngine* TheorySets::getEqualityEngine() 
-{
-  return &d_equalityEngine;
-}
-
-void TheorySets::setMasterEqualityEngine(eq::EqualityEngine* eq)
-{
-  d_equalityEngine.setMasterEqualityEngine(eq);
-}
-
 /**************************** eq::NotifyClass *****************************/
 
 bool TheorySets::NotifyClass::eqNotifyTriggerEquality(TNode equality,
diff --git a/src/theory/sets/theory_sets.h b/src/theory/sets/theory_sets.h
index 84291346b..cb8fdfbc3 100644
--- a/src/theory/sets/theory_sets.h
+++ b/src/theory/sets/theory_sets.h
@@ -48,6 +48,12 @@ class TheorySets : public Theory
   //--------------------------------- initialization
   /** get the official theory rewriter of this theory */
   TheoryRewriter* getTheoryRewriter() override;
+  /**
+   * Returns true if we need an equality engine. If so, we initialize the
+   * information regarding how it should be setup. For details, see the
+   * documentation in Theory::needsEqualityEngine.
+   */
+  bool needsEqualityEngine(EeSetupInfo& esi) override;
   /** finish initialization */
   void finishInit() override;
   //--------------------------------- end initialization
@@ -65,10 +71,7 @@ class TheorySets : public Theory
   PPAssertStatus ppAssert(TNode in, SubstitutionMap& outSubstitutions) override;
   void presolve() override;
   void propagate(Effort) override;
-  void setMasterEqualityEngine(eq::EqualityEngine* eq) override;
   bool isEntailed(Node n, bool pol);
-  /* equality engine */
-  virtual eq::EqualityEngine* getEqualityEngine() override;
  private:
   /** Functions to handle callbacks from equality engine */
   class NotifyClass : public eq::EqualityEngineNotify
@@ -92,9 +95,7 @@ class TheorySets : public Theory
   /** The internal theory */
   std::unique_ptr<TheorySetsPrivate> d_internal;
   /** Instance of the above class */
-  NotifyClass d_notify;  
-  /** Equality engine */
-  eq::EqualityEngine d_equalityEngine;
+  NotifyClass d_notify;
 }; /* class TheorySets */
 
 }/* CVC4::theory::sets namespace */