From a11de769885cf9ac4b2c2f06409976080b326fe6 Mon Sep 17 00:00:00 2001
From: Gereon Kremer <nafur42@gmail.com>
Date: Mon, 16 Aug 2021 17:20:27 -0700
Subject: [PATCH] Push Env class into TheoryState (#7012)

This PR is a follow-up to #7011, making the Env object available in the TheoryState base class.
---
 src/theory/arith/arith_state.cpp              |  6 ++----
 src/theory/arith/arith_state.h                |  3 ++-
 src/theory/arith/theory_arith.cpp             |  2 +-
 src/theory/arrays/theory_arrays.cpp           |  2 +-
 src/theory/bags/solver_state.cpp              |  5 +----
 src/theory/bags/solver_state.h                |  3 ++-
 src/theory/bags/theory_bags.cpp               |  2 +-
 src/theory/builtin/theory_builtin.cpp         |  2 +-
 src/theory/bv/theory_bv.cpp                   |  2 +-
 src/theory/datatypes/theory_datatypes.cpp     |  2 +-
 src/theory/fp/theory_fp.cpp                   |  2 +-
 src/theory/quantifiers/quantifiers_state.cpp  |  7 ++++---
 src/theory/quantifiers/quantifiers_state.h    |  3 +--
 src/theory/quantifiers/theory_quantifiers.cpp |  2 +-
 src/theory/sep/theory_sep.cpp                 |  2 +-
 src/theory/sets/solver_state.cpp              |  7 ++-----
 src/theory/sets/solver_state.h                |  3 +--
 src/theory/sets/theory_sets.cpp               |  2 +-
 src/theory/strings/solver_state.cpp           | 11 ++++++-----
 src/theory/strings/solver_state.h             |  3 +--
 src/theory/strings/theory_strings.cpp         |  2 +-
 src/theory/theory_state.cpp                   | 19 +++++++++++--------
 src/theory/theory_state.h                     | 12 +++++++-----
 src/theory/uf/theory_uf.cpp                   |  2 +-
 test/unit/test_smt.h                          |  2 +-
 25 files changed, 53 insertions(+), 55 deletions(-)

diff --git a/src/theory/arith/arith_state.cpp b/src/theory/arith/arith_state.cpp
index 93d410bf8..3d43077e5 100644
--- a/src/theory/arith/arith_state.cpp
+++ b/src/theory/arith/arith_state.cpp
@@ -21,10 +21,8 @@ namespace cvc5 {
 namespace theory {
 namespace arith {
 
-ArithState::ArithState(context::Context* c,
-                       context::UserContext* u,
-                       Valuation val)
-    : TheoryState(c, u, val), d_parent(nullptr)
+ArithState::ArithState(Env& env, Valuation val)
+    : TheoryState(env, val), d_parent(nullptr)
 {
 }
 
diff --git a/src/theory/arith/arith_state.h b/src/theory/arith/arith_state.h
index a54ae6bf0..0f0f02f02 100644
--- a/src/theory/arith/arith_state.h
+++ b/src/theory/arith/arith_state.h
@@ -38,7 +38,8 @@ class TheoryArithPrivate;
 class ArithState : public TheoryState
 {
  public:
-  ArithState(context::Context* c, context::UserContext* u, Valuation val);
+  ArithState(Env& env,
+             Valuation val);
   ~ArithState() {}
   /** Are we currently in conflict? */
   bool isInConflict() const override;
diff --git a/src/theory/arith/theory_arith.cpp b/src/theory/arith/theory_arith.cpp
index 8cb607dd7..37069d8b8 100644
--- a/src/theory/arith/theory_arith.cpp
+++ b/src/theory/arith/theory_arith.cpp
@@ -39,7 +39,7 @@ TheoryArith::TheoryArith(Env& env, OutputChannel& out, Valuation valuation)
     : Theory(THEORY_ARITH, env, out, valuation),
       d_ppRewriteTimer(smtStatisticsRegistry().registerTimer(
           "theory::arith::ppRewriteTimer")),
-      d_astate(getSatContext(), getUserContext(), valuation),
+      d_astate(env, valuation),
       d_im(*this, d_astate, d_pnm),
       d_ppre(getSatContext(), d_pnm),
       d_bab(d_astate, d_im, d_ppre, d_pnm),
diff --git a/src/theory/arrays/theory_arrays.cpp b/src/theory/arrays/theory_arrays.cpp
index eaa9c9294..1a6dfedbb 100644
--- a/src/theory/arrays/theory_arrays.cpp
+++ b/src/theory/arrays/theory_arrays.cpp
@@ -83,7 +83,7 @@ TheoryArrays::TheoryArrays(Env& env,
       d_ppEqualityEngine(getUserContext(), name + "pp", true),
       d_ppFacts(getUserContext()),
       d_rewriter(d_pnm),
-      d_state(getSatContext(), getUserContext(), valuation),
+      d_state(env, valuation),
       d_im(*this, d_state, d_pnm),
       d_literalsToPropagate(getSatContext()),
       d_literalsToPropagateIndex(getSatContext(), 0),
diff --git a/src/theory/bags/solver_state.cpp b/src/theory/bags/solver_state.cpp
index 6c80e00bd..50c6919fa 100644
--- a/src/theory/bags/solver_state.cpp
+++ b/src/theory/bags/solver_state.cpp
@@ -27,10 +27,7 @@ namespace cvc5 {
 namespace theory {
 namespace bags {
 
-SolverState::SolverState(context::Context* c,
-                         context::UserContext* u,
-                         Valuation val)
-    : TheoryState(c, u, val)
+SolverState::SolverState(Env& env, Valuation val) : TheoryState(env, val)
 {
   d_true = NodeManager::currentNM()->mkConst(true);
   d_false = NodeManager::currentNM()->mkConst(false);
diff --git a/src/theory/bags/solver_state.h b/src/theory/bags/solver_state.h
index 9c2222e95..68fffacbd 100644
--- a/src/theory/bags/solver_state.h
+++ b/src/theory/bags/solver_state.h
@@ -29,7 +29,8 @@ namespace bags {
 class SolverState : public TheoryState
 {
  public:
-  SolverState(context::Context* c, context::UserContext* u, Valuation val);
+  SolverState(Env& env,
+              Valuation val);
 
   /**
    * This function adds the bag representative n to the set d_bags if it is not
diff --git a/src/theory/bags/theory_bags.cpp b/src/theory/bags/theory_bags.cpp
index 42a1e2c6b..f8483064d 100644
--- a/src/theory/bags/theory_bags.cpp
+++ b/src/theory/bags/theory_bags.cpp
@@ -29,7 +29,7 @@ namespace bags {
 
 TheoryBags::TheoryBags(Env& env, OutputChannel& out, Valuation valuation)
     : Theory(THEORY_BAGS, env, out, valuation),
-      d_state(getSatContext(), getUserContext(), valuation),
+      d_state(env, valuation),
       d_im(*this, d_state, d_pnm),
       d_ig(&d_state, &d_im),
       d_notify(*this, d_im),
diff --git a/src/theory/builtin/theory_builtin.cpp b/src/theory/builtin/theory_builtin.cpp
index f12d2caf9..092bcc9ab 100644
--- a/src/theory/builtin/theory_builtin.cpp
+++ b/src/theory/builtin/theory_builtin.cpp
@@ -27,7 +27,7 @@ namespace builtin {
 
 TheoryBuiltin::TheoryBuiltin(Env& env, OutputChannel& out, Valuation valuation)
     : Theory(THEORY_BUILTIN, env, out, valuation),
-      d_state(getSatContext(), getUserContext(), valuation),
+      d_state(env, valuation),
       d_im(*this, d_state, d_pnm, "theory::builtin::")
 {
   // indicate we are using the default theory state and inference managers
diff --git a/src/theory/bv/theory_bv.cpp b/src/theory/bv/theory_bv.cpp
index 1976177b7..d4926a785 100644
--- a/src/theory/bv/theory_bv.cpp
+++ b/src/theory/bv/theory_bv.cpp
@@ -37,7 +37,7 @@ TheoryBV::TheoryBV(Env& env,
     : Theory(THEORY_BV, env, out, valuation, name),
       d_internal(nullptr),
       d_rewriter(),
-      d_state(getSatContext(), getUserContext(), valuation),
+      d_state(env, valuation),
       d_im(*this, d_state, nullptr, "theory::bv::"),
       d_notify(d_im),
       d_invalidateModelCache(getSatContext(), true),
diff --git a/src/theory/datatypes/theory_datatypes.cpp b/src/theory/datatypes/theory_datatypes.cpp
index 759cc4c87..3dfd9b458 100644
--- a/src/theory/datatypes/theory_datatypes.cpp
+++ b/src/theory/datatypes/theory_datatypes.cpp
@@ -60,7 +60,7 @@ TheoryDatatypes::TheoryDatatypes(Env& env,
       d_singleton_eq(getUserContext()),
       d_lemmas_produced_c(getUserContext()),
       d_sygusExtension(nullptr),
-      d_state(getSatContext(), getUserContext(), valuation),
+      d_state(env, valuation),
       d_im(*this, d_state, d_pnm),
       d_notify(d_im, *this)
 {
diff --git a/src/theory/fp/theory_fp.cpp b/src/theory/fp/theory_fp.cpp
index a016c7897..bd5662cdd 100644
--- a/src/theory/fp/theory_fp.cpp
+++ b/src/theory/fp/theory_fp.cpp
@@ -70,7 +70,7 @@ TheoryFp::TheoryFp(Env& env, OutputChannel& out, Valuation valuation)
       d_floatToRealMap(getUserContext()),
       d_abstractionMap(getUserContext()),
       d_rewriter(getUserContext()),
-      d_state(getSatContext(), getUserContext(), valuation),
+      d_state(env, valuation),
       d_im(*this, d_state, d_pnm, "theory::fp::", false),
       d_wbFactsCache(getUserContext())
 {
diff --git a/src/theory/quantifiers/quantifiers_state.cpp b/src/theory/quantifiers/quantifiers_state.cpp
index 7654a6326..39af9c2c9 100644
--- a/src/theory/quantifiers/quantifiers_state.cpp
+++ b/src/theory/quantifiers/quantifiers_state.cpp
@@ -22,11 +22,12 @@ namespace cvc5 {
 namespace theory {
 namespace quantifiers {
 
-QuantifiersState::QuantifiersState(context::Context* c,
-                                   context::UserContext* u,
+QuantifiersState::QuantifiersState(Env& env,
                                    Valuation val,
                                    const LogicInfo& logicInfo)
-    : TheoryState(c, u, val), d_ierCounterc(c), d_logicInfo(logicInfo)
+    : TheoryState(env, val),
+      d_ierCounterc(env.getContext()),
+      d_logicInfo(logicInfo)
 {
   // allow theory combination to go first, once initially
   d_ierCounter = options::instWhenTcFirst() ? 0 : 1;
diff --git a/src/theory/quantifiers/quantifiers_state.h b/src/theory/quantifiers/quantifiers_state.h
index e4a05b078..92b744cd0 100644
--- a/src/theory/quantifiers/quantifiers_state.h
+++ b/src/theory/quantifiers/quantifiers_state.h
@@ -32,8 +32,7 @@ namespace quantifiers {
 class QuantifiersState : public TheoryState
 {
  public:
-  QuantifiersState(context::Context* c,
-                   context::UserContext* u,
+  QuantifiersState(Env& env,
                    Valuation val,
                    const LogicInfo& logicInfo);
   ~QuantifiersState() {}
diff --git a/src/theory/quantifiers/theory_quantifiers.cpp b/src/theory/quantifiers/theory_quantifiers.cpp
index 3adf53b57..a108d4614 100644
--- a/src/theory/quantifiers/theory_quantifiers.cpp
+++ b/src/theory/quantifiers/theory_quantifiers.cpp
@@ -34,7 +34,7 @@ TheoryQuantifiers::TheoryQuantifiers(Env& env,
                                      OutputChannel& out,
                                      Valuation valuation)
     : Theory(THEORY_QUANTIFIERS, env, out, valuation),
-      d_qstate(getSatContext(), getUserContext(), valuation, getLogicInfo()),
+      d_qstate(env, valuation, getLogicInfo()),
       d_qreg(),
       d_treg(d_qstate, d_qreg),
       d_qim(*this, d_qstate, d_qreg, d_treg, d_pnm),
diff --git a/src/theory/sep/theory_sep.cpp b/src/theory/sep/theory_sep.cpp
index 0b16ddd27..92d15653e 100644
--- a/src/theory/sep/theory_sep.cpp
+++ b/src/theory/sep/theory_sep.cpp
@@ -46,7 +46,7 @@ TheorySep::TheorySep(Env& env, OutputChannel& out, Valuation valuation)
     : Theory(THEORY_SEP, env, out, valuation),
       d_lemmas_produced_c(getUserContext()),
       d_bounds_init(false),
-      d_state(getSatContext(), getUserContext(), valuation),
+      d_state(env, valuation),
       d_im(*this, d_state, d_pnm, "theory::sep::"),
       d_notify(*this),
       d_reduce(getUserContext()),
diff --git a/src/theory/sets/solver_state.cpp b/src/theory/sets/solver_state.cpp
index 2aaa82706..6f8976f4d 100644
--- a/src/theory/sets/solver_state.cpp
+++ b/src/theory/sets/solver_state.cpp
@@ -26,11 +26,8 @@ namespace cvc5 {
 namespace theory {
 namespace sets {
 
-SolverState::SolverState(context::Context* c,
-                         context::UserContext* u,
-                         Valuation val,
-                         SkolemCache& skc)
-    : TheoryState(c, u, val), d_skCache(skc), d_members(c)
+SolverState::SolverState(Env& env, Valuation val, SkolemCache& skc)
+    : TheoryState(env, val), d_skCache(skc), d_members(env.getContext())
 {
   d_true = NodeManager::currentNM()->mkConst(true);
   d_false = NodeManager::currentNM()->mkConst(false);
diff --git a/src/theory/sets/solver_state.h b/src/theory/sets/solver_state.h
index 63039eddd..ff9bc8bf9 100644
--- a/src/theory/sets/solver_state.h
+++ b/src/theory/sets/solver_state.h
@@ -46,8 +46,7 @@ class SolverState : public TheoryState
   typedef context::CDHashMap<Node, size_t> NodeIntMap;
 
  public:
-  SolverState(context::Context* c,
-              context::UserContext* u,
+  SolverState(Env& env,
               Valuation val,
               SkolemCache& skc);
   //-------------------------------- initialize per check
diff --git a/src/theory/sets/theory_sets.cpp b/src/theory/sets/theory_sets.cpp
index 5ea13ea4d..23ac08749 100644
--- a/src/theory/sets/theory_sets.cpp
+++ b/src/theory/sets/theory_sets.cpp
@@ -30,7 +30,7 @@ namespace sets {
 TheorySets::TheorySets(Env& env, OutputChannel& out, Valuation valuation)
     : Theory(THEORY_SETS, env, out, valuation),
       d_skCache(),
-      d_state(getSatContext(), getUserContext(), valuation, d_skCache),
+      d_state(env, valuation, d_skCache),
       d_im(*this, d_state, d_pnm),
       d_internal(new TheorySetsPrivate(*this, d_state, d_im, d_skCache, d_pnm)),
       d_notify(*d_internal.get(), d_im)
diff --git a/src/theory/strings/solver_state.cpp b/src/theory/strings/solver_state.cpp
index 1ddf2af58..32ed6896c 100644
--- a/src/theory/strings/solver_state.cpp
+++ b/src/theory/strings/solver_state.cpp
@@ -28,10 +28,11 @@ namespace cvc5 {
 namespace theory {
 namespace strings {
 
-SolverState::SolverState(context::Context* c,
-                         context::UserContext* u,
-                         Valuation& v)
-    : TheoryState(c, u, v), d_eeDisequalities(c), d_pendingConflictSet(c, false), d_pendingConflict(InferenceId::UNKNOWN)
+SolverState::SolverState(Env& env, Valuation& v)
+    : TheoryState(env, v),
+      d_eeDisequalities(env.getContext()),
+      d_pendingConflictSet(env.getContext(), false),
+      d_pendingConflict(InferenceId::UNKNOWN)
 {
   d_zero = NodeManager::currentNM()->mkConst(Rational(0));
   d_false = NodeManager::currentNM()->mkConst(false);
@@ -64,7 +65,7 @@ EqcInfo* SolverState::getOrMakeEqcInfo(Node eqc, bool doMake)
   }
   if (doMake)
   {
-    EqcInfo* ei = new EqcInfo(d_context);
+    EqcInfo* ei = new EqcInfo(d_env.getContext());
     d_eqcInfo[eqc] = ei;
     return ei;
   }
diff --git a/src/theory/strings/solver_state.h b/src/theory/strings/solver_state.h
index 422c29760..bbeb470f7 100644
--- a/src/theory/strings/solver_state.h
+++ b/src/theory/strings/solver_state.h
@@ -48,8 +48,7 @@ class SolverState : public TheoryState
   typedef context::CDList<Node> NodeList;
 
  public:
-  SolverState(context::Context* c,
-              context::UserContext* u,
+  SolverState(Env& env,
               Valuation& v);
   ~SolverState();
   //-------------------------------------- disequality information
diff --git a/src/theory/strings/theory_strings.cpp b/src/theory/strings/theory_strings.cpp
index 19713d0a1..8b71df31b 100644
--- a/src/theory/strings/theory_strings.cpp
+++ b/src/theory/strings/theory_strings.cpp
@@ -54,7 +54,7 @@ TheoryStrings::TheoryStrings(Env& env, OutputChannel& out, Valuation valuation)
     : Theory(THEORY_STRINGS, env, out, valuation),
       d_notify(*this),
       d_statistics(),
-      d_state(getSatContext(), getUserContext(), d_valuation),
+      d_state(env, d_valuation),
       d_eagerSolver(d_state),
       d_termReg(d_state, d_statistics, d_pnm),
       d_extTheoryCb(),
diff --git a/src/theory/theory_state.cpp b/src/theory/theory_state.cpp
index 5ac5e6ae9..72ab24a7e 100644
--- a/src/theory/theory_state.cpp
+++ b/src/theory/theory_state.cpp
@@ -20,22 +20,25 @@
 namespace cvc5 {
 namespace theory {
 
-TheoryState::TheoryState(context::Context* c,
-                         context::UserContext* u,
-                         Valuation val)
-    : d_context(c),
-      d_ucontext(u),
+TheoryState::TheoryState(Env& env, Valuation val)
+    : d_env(env),
       d_valuation(val),
       d_ee(nullptr),
-      d_conflict(c, false)
+      d_conflict(d_env.getContext(), false)
 {
 }
 
 void TheoryState::setEqualityEngine(eq::EqualityEngine* ee) { d_ee = ee; }
 
-context::Context* TheoryState::getSatContext() const { return d_context; }
+context::Context* TheoryState::getSatContext() const
+{
+  return d_env.getContext();
+}
 
-context::UserContext* TheoryState::getUserContext() const { return d_ucontext; }
+context::UserContext* TheoryState::getUserContext() const
+{
+  return d_env.getUserContext();
+}
 
 bool TheoryState::hasTerm(TNode a) const
 {
diff --git a/src/theory/theory_state.h b/src/theory/theory_state.h
index 2c7bad60b..58a66ef46 100644
--- a/src/theory/theory_state.h
+++ b/src/theory/theory_state.h
@@ -20,6 +20,7 @@
 
 #include "context/cdo.h"
 #include "expr/node.h"
+#include "smt/env.h"
 #include "theory/valuation.h"
 
 namespace cvc5 {
@@ -32,7 +33,8 @@ class EqualityEngine;
 class TheoryState
 {
  public:
-  TheoryState(context::Context* c, context::UserContext* u, Valuation val);
+  TheoryState(Env& env,
+              Valuation val);
   virtual ~TheoryState() {}
   /**
    * Set equality engine, where ee is a pointer to the official equality engine
@@ -43,6 +45,8 @@ class TheoryState
   context::Context* getSatContext() const;
   /** Get the user context */
   context::UserContext* getUserContext() const;
+  /** Get the environment */
+  Env& getEnv() const { return d_env; }
   //-------------------------------------- equality information
   /** Is t registered as a term in the equality engine of this class? */
   virtual bool hasTerm(TNode a) const;
@@ -111,10 +115,8 @@ class TheoryState
   Valuation& getValuation();
 
  protected:
-  /** Pointer to the SAT context object used by the theory. */
-  context::Context* d_context;
-  /** Pointer to the user context object used by the theory. */
-  context::UserContext* d_ucontext;
+  /** Reference to the environment. */
+  Env& d_env;
   /**
    * The valuation proxy for the Theory to communicate back with the
    * theory engine (and other theories).
diff --git a/src/theory/uf/theory_uf.cpp b/src/theory/uf/theory_uf.cpp
index 633934f61..3525786d5 100644
--- a/src/theory/uf/theory_uf.cpp
+++ b/src/theory/uf/theory_uf.cpp
@@ -49,7 +49,7 @@ TheoryUF::TheoryUF(Env& env,
       d_ho(nullptr),
       d_functionsTerms(getSatContext()),
       d_symb(getUserContext(), instanceName),
-      d_state(getSatContext(), getUserContext(), valuation),
+      d_state(env, valuation),
       d_im(*this, d_state, d_pnm, "theory::uf::" + instanceName, false),
       d_notify(d_im, *this)
 {
diff --git a/test/unit/test_smt.h b/test/unit/test_smt.h
index c27c02925..f1644dfcd 100644
--- a/test/unit/test_smt.h
+++ b/test/unit/test_smt.h
@@ -208,7 +208,7 @@ class DummyTheory : public theory::Theory
  public:
   DummyTheory(Env& env, theory::OutputChannel& out, theory::Valuation valuation)
       : Theory(theoryId, env, out, valuation),
-        d_state(getSatContext(), getUserContext(), valuation)
+        d_state(env, valuation)
   {
     // use a default theory state object
     d_theoryState = &d_state;
-- 
2.30.2