Refactor term registration visitors (#5875)

This refactors the term registration visitors (PreRegisterVisitor and SharedTermsVisitor), which are responsible for calling Theory::preRegisterTerm and Theory::addSharedTerm.

The purpose of this PR is to resolve a subtle bug in the implementation of PreRegisterVisitor (see explanation below), and to improve performance in particular on the Facebook benchmarks (where preregistering terms is currently 25-35% of runtime on certain challenge benchmarks).

This PR makes it so that that SharedTermsVisitor now subsumes PreRegisterVisitor and should be run when sharing is enabled only. Previously, we ran both PreRegisterVisitor and SharedTermsVisitor, and moreover the former misclassified when a literal had multiple theories since the d_theories field of PreRegisterVisitor is never reset. This meant we would consider a literal to have multiple theories as soon as we saw any literal that had multipleTheories. After this PR, we run SharedTermsVisitor only which also handles calling preRegisterTerm, since it has the ability to anyways (it computes the same information as PreRegisterVisitor regarding which theories care about a term).

The refactoring merges the blocks of code that were copy and pasted in term_registration_visitor.cpp and makes them more optimal, by reducing our calls to Theory::theoryOf.

FYI @barrettcw

6 files changed, 221 insertions, 294 deletions

diff --git a/src/theory/shared_solver.cpp b/src/theory/shared_solver.cpp
index 01dd37f19..4bee75247 100644
--- a/src/theory/shared_solver.cpp
+++ b/src/theory/shared_solver.cpp
@@ -30,7 +30,8 @@ SharedSolver::SharedSolver(TheoryEngine& te, ProofNodeManager* pnm)
     : d_te(te),
       d_logicInfo(te.getLogicInfo()),
       d_sharedTerms(&d_te, d_te.getSatContext(), d_te.getUserContext(), pnm),
-      d_sharedTermsVisitor(d_sharedTerms)
+      d_preRegistrationVisitor(&te, d_te.getSatContext()),
+      d_sharedTermsVisitor(&te, d_sharedTerms)
 {
 }
 
@@ -39,19 +40,31 @@ bool SharedSolver::needsEqualityEngine(theory::EeSetupInfo& esi)
   return false;
 }
 
-void SharedSolver::preRegisterShared(TNode t, bool multipleTheories)
+void SharedSolver::preRegister(TNode atom)
 {
-  // register it with the equality engine manager if sharing is enabled
+  // This method uses two different implementations for preregistering terms,
+  // which depends on whether sharing is enabled.
+  // If sharing is disabled, we use PreRegisterVisitor, which keeps a global
+  // SAT-context dependent cache of terms visited.
+  // If sharing is disabled, we use SharedTermsVisitor which does *not* keep a
+  // global cache. This is because shared terms must be associated with the
+  // given atom, and thus it must traverse the set of subterms in each atom.
+  // See term_registration_visitor.h for more details.
   if (d_logicInfo.isSharingEnabled())
   {
-    preRegisterSharedInternal(t);
+    // register it with the shared terms database if sharing is enabled
+    preRegisterSharedInternal(atom);
+    // Collect the shared terms in atom, as well as calling preregister on the
+    // appropriate theories in atom.
+    // This calls Theory::preRegisterTerm and Theory::addSharedTerm, possibly
+    // multiple times.
+    NodeVisitor<SharedTermsVisitor>::run(d_sharedTermsVisitor, atom);
   }
-  // if multiple theories are present in t
-  if (multipleTheories)
+  else
   {
-    // Collect the shared terms if there are multiple theories
-    // This calls Theory::addSharedTerm, possibly multiple times
-    NodeVisitor<SharedTermsVisitor>::run(d_sharedTermsVisitor, t);
+    // just use the normal preregister visitor, which calls
+    // Theory::preRegisterTerm possibly multiple times.
+    NodeVisitor<PreRegisterVisitor>::run(d_preRegistrationVisitor, atom);
   }
 }
 
diff --git a/src/theory/shared_solver.h b/src/theory/shared_solver.h
index 46009e23c..43117d994 100644
--- a/src/theory/shared_solver.h
+++ b/src/theory/shared_solver.h
@@ -58,16 +58,18 @@ class SharedSolver
   virtual void setEqualityEngine(eq::EqualityEngine* ee) = 0;
   //------------------------------------- end initialization
   /**
-   * Called when the given term t is pre-registered in TheoryEngine.
+   * Called when the given atom is pre-registered in TheoryEngine.
    *
-   * This adds t as an equality to propagate in the shared terms database
-   * if it is an equality, or adds its shared terms if it involves multiple
-   * theories.
+   * This calls Theory::preRegisterTerm on all subterms of atom for the
+   * appropriate theories.
    *
-   * @param t The term to preregister
-   * @param multipleTheories Whether multiple theories are present in t.
+   * Also, if sharing is enabled, this adds atom as an equality to propagate in
+   * the shared terms database if it is an equality, and adds its shared terms
+   * to the appropariate theories.
+   *
+   * @param atom The atom to preregister
    */
-  void preRegisterShared(TNode t, bool multipleTheories);
+  void preRegister(TNode atom);
   /**
    * Pre-notify assertion fact with the given atom. This is called when any
    * fact is asserted in TheoryEngine, just before it is dispatched to the
@@ -124,6 +126,8 @@ class SharedSolver
   const LogicInfo& d_logicInfo;
   /** The database of shared terms.*/
   SharedTermsDatabase d_sharedTerms;
+  /** Default visitor for pre-registration */
+  PreRegisterVisitor d_preRegistrationVisitor;
   /** Visitor for collecting shared terms */
   SharedTermsVisitor d_sharedTermsVisitor;
 };
diff --git a/src/theory/term_registration_visitor.cpp b/src/theory/term_registration_visitor.cpp
index b4ad87f1e..82463367e 100644
--- a/src/theory/term_registration_visitor.cpp
+++ b/src/theory/term_registration_visitor.cpp
@@ -18,9 +18,9 @@
 #include "options/quantifiers_options.h"
 #include "theory/theory_engine.h"
 
-using namespace std;
-using namespace CVC4;
-using namespace theory;
+using namespace CVC4::theory;
+
+namespace CVC4 {
 
 std::string PreRegisterVisitor::toString() const {
   std::stringstream ss;
@@ -32,6 +32,47 @@ std::string PreRegisterVisitor::toString() const {
   return ss.str();
 }
 
+/**
+ * Return true if we already visited the term current with the given parent,
+ * assuming that the set of theories in visitedTheories has already processed
+ * current. This method is used by PreRegisterVisitor and SharedTermsVisitor
+ * below.
+ */
+bool isAlreadyVisited(TheoryIdSet visitedTheories, TNode current, TNode parent)
+{
+  TheoryId currentTheoryId = Theory::theoryOf(current);
+  if (!TheoryIdSetUtil::setContains(currentTheoryId, visitedTheories))
+  {
+    // current theory not visited, return false
+    return false;
+  }
+
+  if (current == parent)
+  {
+    // top-level and current visited, return true
+    return true;
+  }
+
+  // The current theory has already visited it, so now it depends on the parent
+  // and the type
+  TheoryId parentTheoryId = Theory::theoryOf(parent);
+  if (!TheoryIdSetUtil::setContains(parentTheoryId, visitedTheories))
+  {
+    // parent theory not visited, return false
+    return false;
+  }
+
+  // do we need to consider the type?
+  TypeNode type = current.getType();
+  if (currentTheoryId == parentTheoryId && !type.isInterpretedFinite())
+  {
+    // current and parent are the same theory, and we are infinite, return true
+    return true;
+  }
+  TheoryId typeTheoryId = Theory::theoryOf(type);
+  return TheoryIdSetUtil::setContains(typeTheoryId, visitedTheories);
+}
+
 bool PreRegisterVisitor::alreadyVisited(TNode current, TNode parent) {
 
   Debug("register::internal") << "PreRegisterVisitor::alreadyVisited(" << current << "," << parent << ")" << std::endl;
@@ -47,62 +88,16 @@ bool PreRegisterVisitor::alreadyVisited(TNode current, TNode parent) {
     Debug("register::internal") << "quantifier:true" << std::endl;
     return true;
   }
-
-  TheoryId currentTheoryId = Theory::theoryOf(current);
-  TheoryId parentTheoryId = Theory::theoryOf(parent);
-
-  d_theories = TheoryIdSetUtil::setInsert(currentTheoryId, d_theories);
-  d_theories = TheoryIdSetUtil::setInsert(parentTheoryId, d_theories);
-
-  // Should we use the theory of the type
-  bool useType = false;
-  TheoryId typeTheoryId = THEORY_LAST;
-
-  if (current != parent) {
-    if (currentTheoryId != parentTheoryId) {
-      // If enclosed by different theories it's shared -- in read(a, f(a)) f(a) should be shared with integers
-      TypeNode type = current.getType();
-      useType = true;
-      typeTheoryId = Theory::theoryOf(type);
-    } else {
-      TypeNode type = current.getType();
-      typeTheoryId = Theory::theoryOf(type);
-      if (typeTheoryId != currentTheoryId) {
-        if (type.isInterpretedFinite()) {
-          useType = true;
-        }
-      }
-    }
-  }
   
   // Get the theories that have already visited this node
   TNodeToTheorySetMap::iterator find = d_visited.find(current);
   if (find == d_visited.end()) {
-    if (useType) {
-      d_theories = TheoryIdSetUtil::setInsert(typeTheoryId, d_theories);
-    }
+    // not visited at all, return false
     return false;
   }
 
   TheoryIdSet visitedTheories = (*find).second;
-  if (TheoryIdSetUtil::setContains(currentTheoryId, visitedTheories))
-  {
-    // The current theory has already visited it, so now it depends on the parent and the type
-    if (TheoryIdSetUtil::setContains(parentTheoryId, visitedTheories))
-    {
-      if (useType) {
-        TheoryId typeTheoryId2 = Theory::theoryOf(current.getType());
-        d_theories = TheoryIdSetUtil::setInsert(typeTheoryId2, d_theories);
-        return TheoryIdSetUtil::setContains(typeTheoryId2, visitedTheories);
-      } else {
-        return true;
-      }
-    } else {
-      return false;
-    }
-  } else {
-    return false;
-  }
+  return isAlreadyVisited(visitedTheories, current, parent);
 }
 
 void PreRegisterVisitor::visit(TNode current, TNode parent) {
@@ -112,74 +107,96 @@ void PreRegisterVisitor::visit(TNode current, TNode parent) {
     Debug("register::internal") << toString() << std::endl;
   }
 
-  // Get the theories of the terms
-  TheoryId currentTheoryId = Theory::theoryOf(current);
-  TheoryId parentTheoryId  = Theory::theoryOf(parent);
-
-  // Should we use the theory of the type
-  bool useType = false;
-  TheoryId typeTheoryId = THEORY_LAST;
-
-  if (current != parent) {
-    if (currentTheoryId != parentTheoryId) {
-      // If enclosed by different theories it's shared -- in read(a, f(a)) f(a) should be shared with integers
-      TypeNode type = current.getType();
-      useType = true;
-      typeTheoryId = Theory::theoryOf(type);
-    } else {
-      TypeNode type = current.getType();
-      typeTheoryId = Theory::theoryOf(type);
-      if (typeTheoryId != currentTheoryId) {
-        if (type.isInterpretedFinite()) {
-          useType = true;
-        }
-      }
-    }
-  }
-
+  // get the theories we already preregistered with
   TheoryIdSet visitedTheories = d_visited[current];
+
+  // call the preregistration on current, parent or type theories and update
+  // visitedTheories.
+  preRegister(d_engine, visitedTheories, current, parent);
+
   Debug("register::internal")
       << "PreRegisterVisitor::visit(" << current << "," << parent
-      << "): previously registered with "
+      << "): now registered with "
       << TheoryIdSetUtil::setToString(visitedTheories) << std::endl;
-  if (!TheoryIdSetUtil::setContains(currentTheoryId, visitedTheories))
+  // update the theories set for current
+  d_visited[current] = visitedTheories;
+  Assert(d_visited.find(current) != d_visited.end());
+  Assert(alreadyVisited(current, parent));
+}
+
+void PreRegisterVisitor::preRegister(TheoryEngine* te,
+                                     TheoryIdSet& visitedTheories,
+                                     TNode current,
+                                     TNode parent)
+{
+  // Preregister with the current theory, if necessary
+  TheoryId currentTheoryId = Theory::theoryOf(current);
+  preRegisterWithTheory(te, visitedTheories, currentTheoryId, current, parent);
+
+  if (current != parent)
   {
-    visitedTheories =
-        TheoryIdSetUtil::setInsert(currentTheoryId, visitedTheories);
-    d_visited[current] = visitedTheories;
-    Theory* th = d_engine->theoryOf(currentTheoryId);
-    th->preRegisterTerm(current);
-    Debug("register::internal") << "PreRegisterVisitor::visit(" << current << "," << parent << "): adding " << currentTheoryId << std::endl;
+    // preregister with parent theory, if necessary
+    TheoryId parentTheoryId = Theory::theoryOf(parent);
+    preRegisterWithTheory(te, visitedTheories, parentTheoryId, current, parent);
+
+    // Note that if enclosed by different theories it's shared, for example,
+    // in read(a, f(a)), f(a) should be shared with integers.
+    TypeNode type = current.getType();
+    if (currentTheoryId != parentTheoryId || type.isInterpretedFinite())
+    {
+      // preregister with the type's theory, if necessary
+      TheoryId typeTheoryId = Theory::theoryOf(type);
+      preRegisterWithTheory(te, visitedTheories, typeTheoryId, current, parent);
+    }
   }
-  if (!TheoryIdSetUtil::setContains(parentTheoryId, visitedTheories))
+}
+void PreRegisterVisitor::preRegisterWithTheory(TheoryEngine* te,
+                                               TheoryIdSet& visitedTheories,
+                                               TheoryId id,
+                                               TNode current,
+                                               TNode parent)
+{
+  if (TheoryIdSetUtil::setContains(id, visitedTheories))
   {
-    visitedTheories =
-        TheoryIdSetUtil::setInsert(parentTheoryId, visitedTheories);
-    d_visited[current] = visitedTheories;
-    Theory* th = d_engine->theoryOf(parentTheoryId);
-    th->preRegisterTerm(current);
-    Debug("register::internal") << "PreRegisterVisitor::visit(" << current << "," << parent << "): adding " << parentTheoryId << std::endl;
+    // already registered
+    return;
   }
-  if (useType) {
-    if (!TheoryIdSetUtil::setContains(typeTheoryId, visitedTheories))
+  if (Configuration::isAssertionBuild())
+  {
+    Debug("register::internal")
+        << "PreRegisterVisitor::visit(" << current << "," << parent
+        << "): adding " << id << std::endl;
+    // This should never throw an exception, since theories should be
+    // guaranteed to be initialized.
+    // These checks don't work with finite model finding, because it
+    // uses Rational constants to represent cardinality constraints,
+    // even though arithmetic isn't actually involved.
+    if (!options::finiteModelFind())
     {
-      visitedTheories =
-          TheoryIdSetUtil::setInsert(typeTheoryId, visitedTheories);
-      d_visited[current] = visitedTheories;
-      Theory* th = d_engine->theoryOf(typeTheoryId);
-      th->preRegisterTerm(current);
-      Debug("register::internal") << "PreRegisterVisitor::visit(" << current << "," << parent << "): adding " << parentTheoryId << std::endl;
+      if (!te->isTheoryEnabled(id))
+      {
+        const LogicInfo& l = te->getLogicInfo();
+        LogicInfo newLogicInfo = l.getUnlockedCopy();
+        newLogicInfo.enableTheory(id);
+        newLogicInfo.lock();
+        std::stringstream ss;
+        ss << "The logic was specified as " << l.getLogicString()
+           << ", which doesn't include " << id
+           << ", but found a term in that theory." << std::endl
+           << "You might want to extend your logic to "
+           << newLogicInfo.getLogicString() << std::endl;
+        throw LogicException(ss.str());
+      }
     }
   }
-  Debug("register::internal")
-      << "PreRegisterVisitor::visit(" << current << "," << parent
-      << "): now registered with "
-      << TheoryIdSetUtil::setToString(visitedTheories) << std::endl;
-
-  Assert(d_visited.find(current) != d_visited.end());
-  Assert(alreadyVisited(current, parent));
+  // call the theory's preRegisterTerm method
+  visitedTheories = TheoryIdSetUtil::setInsert(id, visitedTheories);
+  Theory* th = te->theoryOf(id);
+  th->preRegisterTerm(current);
 }
 
+void PreRegisterVisitor::start(TNode node) {}
+
 std::string SharedTermsVisitor::toString() const {
   std::stringstream ss;
   TNodeVisitedMap::const_iterator it = d_visited.begin();
@@ -206,59 +223,14 @@ bool SharedTermsVisitor::alreadyVisited(TNode current, TNode parent) const {
     return true;
   }
   TNodeVisitedMap::const_iterator find = d_visited.find(current);
-
   // If node is not visited at all, just return false
   if (find == d_visited.end()) {
     Debug("register::internal") << "1:false" << std::endl;
     return false;
   }
 
-  TheoryIdSet theories = (*find).second;
-
-  TheoryId currentTheoryId = Theory::theoryOf(current);
-  TheoryId parentTheoryId  = Theory::theoryOf(parent);
-
-  // Should we use the theory of the type
-  bool useType = false;
-  TheoryId typeTheoryId = THEORY_LAST;
-
-  if (current != parent) {
-    if (currentTheoryId != parentTheoryId) {
-      // If enclosed by different theories it's shared -- in read(a, f(a)) f(a) should be shared with integers
-      TypeNode type = current.getType();
-      useType = true;
-      typeTheoryId = Theory::theoryOf(type);
-    } else {
-      TypeNode type = current.getType();
-      typeTheoryId = Theory::theoryOf(type);
-      if (typeTheoryId != currentTheoryId) {
-        if (type.isInterpretedFinite()) {
-          useType = true;
-        }
-      }
-    }
-  }
-
-  if (TheoryIdSetUtil::setContains(currentTheoryId, theories))
-  {
-    if (TheoryIdSetUtil::setContains(parentTheoryId, theories))
-    {
-      if (useType)
-      {
-        return TheoryIdSetUtil::setContains(typeTheoryId, theories);
-      }
-      else
-      {
-        return true;
-      }
-    }
-    else
-    {
-      return false;
-    }
-  } else {
-    return false;
-  }
+  TheoryIdSet visitedTheories = (*find).second;
+  return isAlreadyVisited(visitedTheories, current, parent);
 }
 
 void SharedTermsVisitor::visit(TNode current, TNode parent) {
@@ -267,66 +239,16 @@ void SharedTermsVisitor::visit(TNode current, TNode parent) {
   if (Debug.isOn("register::internal")) {
     Debug("register::internal") << toString() << std::endl;
   }
-
-  // Get the theories of the terms
-  TheoryId currentTheoryId = Theory::theoryOf(current);
-  TheoryId parentTheoryId  = Theory::theoryOf(parent);
-
-  // Should we use the theory of the type
-  bool useType = false;
-  TheoryId typeTheoryId = THEORY_LAST;
-
-  if (current != parent) {
-    if (currentTheoryId != parentTheoryId) {
-      // If enclosed by different theories it's shared -- in read(a, f(a)) f(a) should be shared with integers
-      TypeNode type = current.getType();
-      useType = true;
-      typeTheoryId = Theory::theoryOf(type);
-    } else {
-      TypeNode type = current.getType();
-      typeTheoryId = Theory::theoryOf(type);
-      if (typeTheoryId != currentTheoryId) {
-        if (type.isInterpretedFinite()) {
-          useType = true;
-        }
-      }
-    }
-  }
-
   TheoryIdSet visitedTheories = d_visited[current];
-  Debug("register::internal")
-      << "SharedTermsVisitor::visit(" << current << "," << parent
-      << "): previously registered with "
-      << TheoryIdSetUtil::setToString(visitedTheories) << std::endl;
-  if (!TheoryIdSetUtil::setContains(currentTheoryId, visitedTheories))
-  {
-    visitedTheories =
-        TheoryIdSetUtil::setInsert(currentTheoryId, visitedTheories);
-    Debug("register::internal") << "SharedTermsVisitor::visit(" << current << "," << parent << "): adding " << currentTheoryId << std::endl;
-  }
-  if (!TheoryIdSetUtil::setContains(parentTheoryId, visitedTheories))
-  {
-    visitedTheories =
-        TheoryIdSetUtil::setInsert(parentTheoryId, visitedTheories);
-    Debug("register::internal") << "SharedTermsVisitor::visit(" << current << "," << parent << "): adding " << parentTheoryId << std::endl;
-  }
-  if (useType) {
-    if (!TheoryIdSetUtil::setContains(typeTheoryId, visitedTheories))
-    {
-      visitedTheories =
-          TheoryIdSetUtil::setInsert(typeTheoryId, visitedTheories);
-      Debug("register::internal") << "SharedTermsVisitor::visit(" << current << "," << parent << "): adding " << typeTheoryId << std::endl;
-    }
-  }
-  Debug("register::internal")
-      << "SharedTermsVisitor::visit(" << current << "," << parent
-      << "): now registered with "
-      << TheoryIdSetUtil::setToString(visitedTheories) << std::endl;
+
+  // preregister the term with the current, parent or type theories, as needed
+  PreRegisterVisitor::preRegister(d_engine, visitedTheories, current, parent);
 
   // Record the new theories that we visited
   d_visited[current] = visitedTheories;
 
   // If there is more than two theories and a new one has been added notify the shared terms database
+  TheoryId currentTheoryId = Theory::theoryOf(current);
   if (TheoryIdSetUtil::setDifference(
           visitedTheories, TheoryIdSetUtil::setInsert(currentTheoryId)))
   {
@@ -338,7 +260,7 @@ void SharedTermsVisitor::visit(TNode current, TNode parent) {
 }
 
 void SharedTermsVisitor::start(TNode node) {
-  clear();
+  d_visited.clear();
   d_atom = node;
 }
 
@@ -350,3 +272,5 @@ void SharedTermsVisitor::clear() {
   d_atom = TNode();
   d_visited.clear();
 }
+
+}  // namespace CVC4
diff --git a/src/theory/term_registration_visitor.h b/src/theory/term_registration_visitor.h
index 94fc83b30..6e09d90cc 100644
--- a/src/theory/term_registration_visitor.h
+++ b/src/theory/term_registration_visitor.h
@@ -51,21 +51,16 @@ class PreRegisterVisitor {
   TNodeToTheorySetMap d_visited;
 
   /**
-   * A set of all theories in the term
-   */
-  theory::TheoryIdSet d_theories;
-
-  /**
    * String representation of the visited map, for debugging purposes.
    */
   std::string toString() const;
 
  public:
-  /** Returned set tells us which theories there are */
-  typedef theory::TheoryIdSet return_type;
+  /** required to instantiate template for NodeVisitor */
+  using return_type = void;
 
   PreRegisterVisitor(TheoryEngine* engine, context::Context* context)
-      : d_engine(engine), d_visited(context), d_theories(0)
+      : d_engine(engine), d_visited(context)
   {
   }
 
@@ -82,14 +77,43 @@ class PreRegisterVisitor {
   void visit(TNode current, TNode parent);
 
   /**
-   * Marks the node as the starting literal.
+   * Marks the node as the starting literal, which does nothing. This method
+   * is required to instantiate template for NodeVisitor.
    */
-  void start(TNode node) {}
+  void start(TNode node);
+
+  /** Called when the visitor is finished with a term, do nothing */
+  void done(TNode node) {}
 
   /**
-   * Notifies the engine of all the theories used.
+   * Preregister the term current occuring under term parent.  This calls
+   * Theory::preRegisterTerm for the theories of current and parent, as well
+   * as the theory of current's type, if it is finite.
+   *
+   * This method takes a set of theories visitedTheories that have already
+   * preregistered current and updates this set with the theories that
+   * preregister current during this call
+   *
+   * @param te Pointer to the theory engine containing the theories
+   * @param visitedTheories The theories that have already preregistered current
+   * @param current The term to preregister
+   * @param parent The parent term of current
    */
-  theory::TheoryIdSet done(TNode node) { return d_theories; }
+  static void preRegister(TheoryEngine* te,
+                          theory::TheoryIdSet& visitedTheories,
+                          TNode current,
+                          TNode parent);
+
+ private:
+  /**
+   * Helper for above, called whether we wish to register a term with a theory
+   * given by an identifier id.
+   */
+  static void preRegisterWithTheory(TheoryEngine* te,
+                                    theory::TheoryIdSet& visitedTheories,
+                                    theory::TheoryId id,
+                                    TNode current,
+                                    TNode parent);
 };
 
 
@@ -100,9 +124,6 @@ class PreRegisterVisitor {
  */
 class SharedTermsVisitor {
 
-  /** The shared terms database */
-  SharedTermsDatabase& d_sharedTerms;
-
   /**
    * Cache from preprocessing of atoms.
    */
@@ -118,14 +139,16 @@ class SharedTermsVisitor {
   /** 
    * The initial atom.
    */
-  TNode d_atom; 
-    
-public:
+  TNode d_atom;
 
-  typedef void return_type;
+ public:
+  /** required to instantiate template for NodeVisitor */
+  using return_type = void;
 
-  SharedTermsVisitor(SharedTermsDatabase& sharedTerms)
-  : d_sharedTerms(sharedTerms) {}
+  SharedTermsVisitor(TheoryEngine* te, SharedTermsDatabase& sharedTerms)
+      : d_engine(te), d_sharedTerms(sharedTerms)
+  {
+  }
 
   /**
    * Returns true is current has already been pre-registered with both current and parent theories.
@@ -136,9 +159,9 @@ public:
    * Pre-registeres current with any of the current and parent theories that haven't seen the term yet.
    */
   void visit(TNode current, TNode parent);
-  
+
   /**
-   * Marks the node as the starting literal.
+   * Marks the node as the starting literal, which clears the state.
    */
   void start(TNode node);
 
@@ -151,6 +174,12 @@ public:
    * Clears the internal state.
    */   
   void clear();
+
+ private:
+  /** The engine */
+  TheoryEngine* d_engine;
+  /** The shared terms database */
+  SharedTermsDatabase& d_sharedTerms;
 };
 
 
diff --git a/src/theory/theory_engine.cpp b/src/theory/theory_engine.cpp
index 9ba97fa77..a2e608b9f 100644
--- a/src/theory/theory_engine.cpp
+++ b/src/theory/theory_engine.cpp
@@ -224,7 +224,6 @@ TheoryEngine::TheoryEngine(context::Context* context,
       d_quantEngine(nullptr),
       d_decManager(new DecisionManager(userContext)),
       d_relManager(nullptr),
-      d_preRegistrationVisitor(this, context),
       d_eager_model_building(false),
       d_inConflict(context, false),
       d_inSatMode(false),
@@ -292,47 +291,11 @@ void TheoryEngine::preRegister(TNode preprocessed) {
       // should not have witness
       Assert(!expr::hasSubtermKind(kind::WITNESS, preprocessed));
 
-      // Pre-register the terms in the atom
-      theory::TheoryIdSet theories = NodeVisitor<PreRegisterVisitor>::run(
-          d_preRegistrationVisitor, preprocessed);
-      theories = TheoryIdSetUtil::setRemove(THEORY_BOOL, theories);
-      // Remove the top theory, if any more that means multiple theories were
-      // involved
-      bool multipleTheories =
-          TheoryIdSetUtil::setRemove(Theory::theoryOf(preprocessed), theories);
-      if (Configuration::isAssertionBuild())
-      {
-        TheoryId i;
-        // This should never throw an exception, since theories should be
-        // guaranteed to be initialized.
-        // These checks don't work with finite model finding, because it
-        // uses Rational constants to represent cardinality constraints,
-        // even though arithmetic isn't actually involved.
-        if (!options::finiteModelFind())
-        {
-          while ((i = TheoryIdSetUtil::setPop(theories)) != THEORY_LAST)
-          {
-            if (!d_logicInfo.isTheoryEnabled(i))
-            {
-              LogicInfo newLogicInfo = d_logicInfo.getUnlockedCopy();
-              newLogicInfo.enableTheory(i);
-              newLogicInfo.lock();
-              std::stringstream ss;
-              ss << "The logic was specified as " << d_logicInfo.getLogicString()
-                << ", which doesn't include " << i
-                << ", but found a term in that theory." << std::endl
-                << "You might want to extend your logic to "
-                << newLogicInfo.getLogicString() << std::endl;
-              throw LogicException(ss.str());
-            }
-          }
-        }
-      }
-
-      // pre-register with the shared solver, which also handles
-      // calling prepregister on individual theories.
+      // pre-register with the shared solver, which handles
+      // calling prepregister on individual theories, adding shared terms,
+      // and setting up equalities to propagate in the shared term database.
       Assert(d_sharedSolver != nullptr);
-      d_sharedSolver->preRegisterShared(preprocessed, multipleTheories);
+      d_sharedSolver->preRegister(preprocessed);
     }
 
     // Leaving pre-register
@@ -1307,8 +1270,6 @@ void TheoryEngine::lemma(theory::TrustNode tlemma,
   // get the node
   Node node = tlemma.getNode();
   Node lemma = tlemma.getProven();
-  Trace("te-lemma") << "Lemma, input: " << lemma << ", property = " << p
-                    << std::endl;
 
   Assert(!expr::hasFreeVar(lemma));
 
diff --git a/src/theory/theory_engine.h b/src/theory/theory_engine.h
index e2e7850b5..224d84664 100644
--- a/src/theory/theory_engine.h
+++ b/src/theory/theory_engine.h
@@ -31,7 +31,6 @@
 #include "theory/interrupted.h"
 #include "theory/rewriter.h"
 #include "theory/sort_inference.h"
-#include "theory/term_registration_visitor.h"
 #include "theory/theory.h"
 #include "theory/theory_preprocessor.h"
 #include "theory/trust_node.h"
@@ -159,9 +158,6 @@ class TheoryEngine {
   /** The relevance manager */
   std::unique_ptr<theory::RelevanceManager> d_relManager;
 
-  /** Default visitor for pre-registration */
-  PreRegisterVisitor d_preRegistrationVisitor;
-
   /** are we in eager model building mode? (see setEagerModelBuilding). */
   bool d_eager_model_building;