Fix term registration and non-theory-preprocessed terms in substitutions (#6080)

This fixes two issues for preprocessing:
(1) The term preregistration visitor was calling preRegister on terms multiple times in a SAT context, which the linear arithmetic solver is sensitive to.
(2) It was possible for non-preprocessed terms to appear in assertions if they were on the RHS of substitutions learned by non-clausal simplification, and substituted into assertions post-theory-preprocessing.

To fix (1), the SharedTermsVisitor is update to track which theories has preregistered each term, as is done in the PreRegisterVisitor. To fix (2), we no longer apply-subst after theory preprocessing.

These two fixes are required to fix #6071.

Note: we should performance test this on SMT-LIB.

6 files changed, 67 insertions, 30 deletions

diff --git a/src/smt/process_assertions.cpp b/src/smt/process_assertions.cpp
index 9a6486ded..3c7d88fe8 100644
--- a/src/smt/process_assertions.cpp
+++ b/src/smt/process_assertions.cpp
@@ -341,10 +341,8 @@ bool ProcessAssertions::apply(Assertions& as)
   d_passes["theory-rewrite-eq"]->apply(&assertions);
   // apply theory preprocess, which includes ITE removal
   d_passes["theory-preprocess"]->apply(&assertions);
-  // This is needed because when solving incrementally, removeITEs may
-  // introduce skolems that were solved for earlier and thus appear in the
-  // substitution map.
-  d_passes["apply-substs"]->apply(&assertions);
+  // notice that we do not apply substitutions as a last step here, since
+  // the range of substitutions is not theory-preprocessed.
 
   if (options::bitblastMode() == options::BitblastMode::EAGER)
   {
diff --git a/src/theory/arith/theory_arith_private.cpp b/src/theory/arith/theory_arith_private.cpp
index c096a4896..1cf4a2993 100644
--- a/src/theory/arith/theory_arith_private.cpp
+++ b/src/theory/arith/theory_arith_private.cpp
@@ -1418,7 +1418,7 @@ void TheoryArithPrivate::setupPolynomial(const Polynomial& poly) {
 }
 
 void TheoryArithPrivate::setupAtom(TNode atom) {
-  Assert(isRelationOperator(atom.getKind()));
+  Assert(isRelationOperator(atom.getKind())) << atom;
   Assert(Comparison::isNormalAtom(atom));
   Assert(!isSetup(atom));
   Assert(!d_constraintDatabase.hasLiteral(atom));
diff --git a/src/theory/shared_solver.cpp b/src/theory/shared_solver.cpp
index 91e008f80..87625291c 100644
--- a/src/theory/shared_solver.cpp
+++ b/src/theory/shared_solver.cpp
@@ -33,7 +33,7 @@ SharedSolver::SharedSolver(TheoryEngine& te, ProofNodeManager* pnm)
       d_logicInfo(te.getLogicInfo()),
       d_sharedTerms(&d_te, d_te.getSatContext(), d_te.getUserContext(), pnm),
       d_preRegistrationVisitor(&te, d_te.getSatContext()),
-      d_sharedTermsVisitor(&te, d_sharedTerms)
+      d_sharedTermsVisitor(&te, d_sharedTerms, d_te.getSatContext())
 {
 }
 
@@ -44,6 +44,7 @@ bool SharedSolver::needsEqualityEngine(theory::EeSetupInfo& esi)
 
 void SharedSolver::preRegister(TNode atom)
 {
+  Trace("theory") << "SharedSolver::preRegister atom " << atom << std::endl;
   // 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
@@ -68,6 +69,7 @@ void SharedSolver::preRegister(TNode atom)
     // Theory::preRegisterTerm possibly multiple times.
     NodeVisitor<PreRegisterVisitor>::run(d_preRegistrationVisitor, atom);
   }
+  Trace("theory") << "SharedSolver::preRegister atom finished" << std::endl;
 }
 
 void SharedSolver::preNotifySharedFact(TNode atom)
diff --git a/src/theory/term_registration_visitor.cpp b/src/theory/term_registration_visitor.cpp
index 3a3003ba9..253f21d98 100644
--- a/src/theory/term_registration_visitor.cpp
+++ b/src/theory/term_registration_visitor.cpp
@@ -113,8 +113,9 @@ void PreRegisterVisitor::visit(TNode current, TNode parent) {
   TheoryIdSet visitedTheories = d_visited[current];
 
   // call the preregistration on current, parent or type theories and update
-  // visitedTheories.
-  preRegister(d_engine, visitedTheories, current, parent);
+  // visitedTheories. The set of preregistering theories coincides with
+  // visitedTheories here.
+  preRegister(d_engine, visitedTheories, current, parent, visitedTheories);
 
   Debug("register::internal")
       << "PreRegisterVisitor::visit(" << current << "," << parent
@@ -129,17 +130,20 @@ void PreRegisterVisitor::visit(TNode current, TNode parent) {
 void PreRegisterVisitor::preRegister(TheoryEngine* te,
                                      TheoryIdSet& visitedTheories,
                                      TNode current,
-                                     TNode parent)
+                                     TNode parent,
+                                     TheoryIdSet preregTheories)
 {
   // Preregister with the current theory, if necessary
   TheoryId currentTheoryId = Theory::theoryOf(current);
-  preRegisterWithTheory(te, visitedTheories, currentTheoryId, current, parent);
+  preRegisterWithTheory(
+      te, visitedTheories, currentTheoryId, current, parent, preregTheories);
 
   if (current != parent)
   {
     // preregister with parent theory, if necessary
     TheoryId parentTheoryId = Theory::theoryOf(parent);
-    preRegisterWithTheory(te, visitedTheories, parentTheoryId, current, parent);
+    preRegisterWithTheory(
+        te, visitedTheories, parentTheoryId, current, parent, preregTheories);
 
     // Note that if enclosed by different theories it's shared, for example,
     // in read(a, f(a)), f(a) should be shared with integers.
@@ -148,7 +152,8 @@ void PreRegisterVisitor::preRegister(TheoryEngine* te,
     {
       // preregister with the type's theory, if necessary
       TheoryId typeTheoryId = Theory::theoryOf(type);
-      preRegisterWithTheory(te, visitedTheories, typeTheoryId, current, parent);
+      preRegisterWithTheory(
+          te, visitedTheories, typeTheoryId, current, parent, preregTheories);
     }
   }
 }
@@ -156,11 +161,18 @@ void PreRegisterVisitor::preRegisterWithTheory(TheoryEngine* te,
                                                TheoryIdSet& visitedTheories,
                                                TheoryId id,
                                                TNode current,
-                                               TNode parent)
+                                               TNode parent,
+                                               TheoryIdSet preregTheories)
 {
   if (TheoryIdSetUtil::setContains(id, visitedTheories))
   {
-    // already registered
+    // already visited
+    return;
+  }
+  visitedTheories = TheoryIdSetUtil::setInsert(id, visitedTheories);
+  if (TheoryIdSetUtil::setContains(id, preregTheories))
+  {
+    // already pregregistered
     return;
   }
   if (Configuration::isAssertionBuild())
@@ -192,7 +204,6 @@ void PreRegisterVisitor::preRegisterWithTheory(TheoryEngine* te,
     }
   }
   // call the theory's preRegisterTerm method
-  visitedTheories = TheoryIdSetUtil::setInsert(id, visitedTheories);
   Theory* th = te->theoryOf(id);
   th->preRegisterTerm(current);
 }
@@ -242,13 +253,19 @@ void SharedTermsVisitor::visit(TNode current, TNode parent) {
     Debug("register::internal") << toString() << std::endl;
   }
   TheoryIdSet visitedTheories = d_visited[current];
+  TheoryIdSet preregTheories = d_preregistered[current];
 
   // preregister the term with the current, parent or type theories, as needed
-  PreRegisterVisitor::preRegister(d_engine, visitedTheories, current, parent);
+  PreRegisterVisitor::preRegister(
+      d_engine, visitedTheories, current, parent, preregTheories);
 
   // Record the new theories that we visited
   d_visited[current] = visitedTheories;
 
+  // add visited theories to those who have preregistered
+  d_preregistered[current] =
+      TheoryIdSetUtil::setUnion(preregTheories, 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(
diff --git a/src/theory/term_registration_visitor.h b/src/theory/term_registration_visitor.h
index e3dc1977a..c99ed6b99 100644
--- a/src/theory/term_registration_visitor.h
+++ b/src/theory/term_registration_visitor.h
@@ -59,8 +59,8 @@ class PreRegisterVisitor {
   /** required to instantiate template for NodeVisitor */
   using return_type = void;
 
-  PreRegisterVisitor(TheoryEngine* engine, context::Context* context)
-      : d_engine(engine), d_visited(context)
+  PreRegisterVisitor(TheoryEngine* engine, context::Context* c)
+      : d_engine(engine), d_visited(c)
   {
   }
 
@@ -98,11 +98,15 @@ class PreRegisterVisitor {
    * @param visitedTheories The theories that have already preregistered current
    * @param current The term to preregister
    * @param parent The parent term of current
+   * @param preregTheories The theories that have already preregistered current.
+   * If there is no theory sharing, this coincides with visitedTheories.
+   * Otherwise, visitedTheories may be a subset of preregTheories.
    */
   static void preRegister(TheoryEngine* te,
                           theory::TheoryIdSet& visitedTheories,
                           TNode current,
-                          TNode parent);
+                          TNode parent,
+                          theory::TheoryIdSet preregTheories);
 
  private:
   /**
@@ -113,7 +117,8 @@ class PreRegisterVisitor {
                                     theory::TheoryIdSet& visitedTheories,
                                     theory::TheoryId id,
                                     TNode current,
-                                    TNode parent);
+                                    TNode parent,
+                                    theory::TheoryIdSet preregTheories);
 };
 
 
@@ -123,14 +128,10 @@ class PreRegisterVisitor {
  * been visited already, we need to visit it again, since we need to associate it with both atoms.
  */
 class SharedTermsVisitor {
-
-  /**
-   * Cache from preprocessing of atoms.
-   */
-  typedef std::unordered_map<TNode, theory::TheoryIdSet, TNodeHashFunction>
-      TNodeVisitedMap;
-  TNodeVisitedMap d_visited;
-
+  using TNodeVisitedMap =
+      std::unordered_map<TNode, theory::TheoryIdSet, TNodeHashFunction>;
+  using TNodeToTheorySetMap =
+      context::CDHashMap<TNode, theory::TheoryIdSet, TNodeHashFunction>;
   /**
    * String representation of the visited map, for debugging purposes.
    */
@@ -145,8 +146,10 @@ class SharedTermsVisitor {
   /** required to instantiate template for NodeVisitor */
   using return_type = void;
 
-  SharedTermsVisitor(TheoryEngine* te, SharedTermsDatabase& sharedTerms)
-      : d_engine(te), d_sharedTerms(sharedTerms)
+  SharedTermsVisitor(TheoryEngine* te,
+                     SharedTermsDatabase& sharedTerms,
+                     context::Context* c)
+      : d_engine(te), d_sharedTerms(sharedTerms), d_preregistered(c)
   {
   }
 
@@ -180,6 +183,10 @@ class SharedTermsVisitor {
   TheoryEngine* d_engine;
   /** The shared terms database */
   SharedTermsDatabase& d_sharedTerms;
+  /** Cache of nodes we have visited in this traversal */
+  TNodeVisitedMap d_visited;
+  /** (Global) cache of nodes we have preregistered in this SAT context */
+  TNodeToTheorySetMap d_preregistered;
 };
 
 
diff --git a/test/regress/regress1/strings/issue6071-arith-prereg-i.smt2 b/test/regress/regress1/strings/issue6071-arith-prereg-i.smt2
new file mode 100644
index 000000000..13fcaf428
--- /dev/null
+++ b/test/regress/regress1/strings/issue6071-arith-prereg-i.smt2
@@ -0,0 +1,13 @@
+; COMMAND-LINE: -i
+; EXPECT: sat
+(set-logic ALL)
+(set-option :strings-lazy-pp false)
+(declare-fun ufbi4 (Bool Bool Bool Bool) Int)
+(declare-fun str0 () String)
+(declare-fun str8 () String)
+(declare-fun i16 () Int)
+(assert (= (str.to_int str0) (ufbi4 false true false (= 0 i16))))
+(push)
+(assert (= str8 (str.from_int (str.to_int str0))))
+(push)
+(check-sat)