Standardize collectModelInfo and theory-specific collectRelevantTerms (#4909)

This is work towards a configurable approach to equality engine management. This PR does not change any behavior, it only reorganizes the code.

This PR introduces the standard template for collectModelInfo, which isolates the usage of equality engine in relation to the model. In the future, theories will be encouraged to use the standard template for collectModelInfo and override collectRelevantTerms/collectModelValues only. This is to allow custom theory-independent modifications to building models (e.g. don't assert equality engine to model if we are using a centralized approach).

This PR standardizes TheoryArrays and TheoryDatatypes custom implementation of collectRelevantTerms as work towards using the standard template for collectModelInfo. Notice this required separating two portions of a loop in TheoryArrays::collectModelInfo which was doing two different things (collecting arrays and relevant terms).

6 files changed, 228 insertions, 130 deletions

diff --git a/src/theory/arrays/theory_arrays.cpp b/src/theory/arrays/theory_arrays.cpp
index 85759b75f..51e1b367c 100644
--- a/src/theory/arrays/theory_arrays.cpp
+++ b/src/theory/arrays/theory_arrays.cpp
@@ -1090,119 +1090,41 @@ void TheoryArrays::computeCareGraph()
 
 bool TheoryArrays::collectModelInfo(TheoryModel* m)
 {
-  set<Node> termSet;
-
-  // Compute terms appearing in assertions and shared terms
+  // Compute terms appearing in assertions and shared terms, and also
+  // include additional reads due to the RIntro1 and RIntro2 rules.
+  std::set<Node> termSet;
   computeRelevantTerms(termSet);
 
-  // Compute arrays that we need to produce representatives for and also make sure RIntro1 reads are included in the relevant set of reads
+  // Send the equality engine information to the model
+  if (!m->assertEqualityEngine(d_equalityEngine, &termSet))
+  {
+    return false;
+  }
+
   NodeManager* nm = NodeManager::currentNM();
+  // Compute arrays that we need to produce representatives for
   std::vector<Node> arrays;
-  bool computeRep, isArray;
   eq::EqClassesIterator eqcs_i = eq::EqClassesIterator(d_equalityEngine);
   for (; !eqcs_i.isFinished(); ++eqcs_i) {
     Node eqc = (*eqcs_i);
-    isArray = eqc.getType().isArray();
-    if (!isArray) {
+    if (!eqc.getType().isArray())
+    {
+      // not an array, skip
       continue;
     }
-    computeRep = false;
     eq::EqClassIterator eqc_i = eq::EqClassIterator(eqc, d_equalityEngine);
     for (; !eqc_i.isFinished(); ++eqc_i) {
       Node n = *eqc_i;
       // If this EC is an array type and it contains something other than STORE nodes, we have to compute a representative explicitly
-      if (isArray && termSet.find(n) != termSet.end()) {
-        if (n.getKind() == kind::STORE) {
-          // Make sure RIntro1 reads are included
-          Node r = nm->mkNode(kind::SELECT, n, n[1]);
-          Trace("arrays::collectModelInfo") << "TheoryArrays::collectModelInfo, adding RIntro1 read: " << r << endl;
-          termSet.insert(r);
-        }
-        else if (!computeRep) {
+      if (termSet.find(n) != termSet.end())
+      {
+        if (n.getKind() != kind::STORE)
+        {
           arrays.push_back(n);
-          computeRep = true;
-        }
-      }
-    }
-  }
-
-  // Now do a fixed-point iteration to get all reads that need to be included because of RIntro2 rule
-  bool changed;
-  do {
-    changed = false;
-    eqcs_i = eq::EqClassesIterator(d_equalityEngine);
-    for (; !eqcs_i.isFinished(); ++eqcs_i) {
-      Node eqc = (*eqcs_i);
-      eq::EqClassIterator eqc_i = eq::EqClassIterator(eqc, d_equalityEngine);
-      for (; !eqc_i.isFinished(); ++eqc_i) {
-        Node n = *eqc_i;
-        if (n.getKind() == kind::SELECT && termSet.find(n) != termSet.end()) {
-
-          // Find all terms equivalent to n[0] and get corresponding read terms
-          Node array_eqc = d_equalityEngine->getRepresentative(n[0]);
-          eq::EqClassIterator array_eqc_i =
-              eq::EqClassIterator(array_eqc, d_equalityEngine);
-          for (; !array_eqc_i.isFinished(); ++array_eqc_i) {
-            Node arr = *array_eqc_i;
-            if (arr.getKind() == kind::STORE
-                && termSet.find(arr) != termSet.end()
-                && !d_equalityEngine->areEqual(arr[1], n[1]))
-            {
-              Node r = nm->mkNode(kind::SELECT, arr, n[1]);
-              if (termSet.find(r) == termSet.end()
-                  && d_equalityEngine->hasTerm(r))
-              {
-                Trace("arrays::collectModelInfo") << "TheoryArrays::collectModelInfo, adding RIntro2(a) read: " << r << endl;
-                termSet.insert(r);
-                changed = true;
-              }
-              r = nm->mkNode(kind::SELECT, arr[0], n[1]);
-              if (termSet.find(r) == termSet.end()
-                  && d_equalityEngine->hasTerm(r))
-              {
-                Trace("arrays::collectModelInfo") << "TheoryArrays::collectModelInfo, adding RIntro2(b) read: " << r << endl;
-                termSet.insert(r);
-                changed = true;
-              }
-            }
-          }
-
-          // Find all stores in which n[0] appears and get corresponding read terms
-          const CTNodeList* instores = d_infoMap.getInStores(array_eqc);
-          size_t it = 0;
-          for(; it < instores->size(); ++it) {
-            TNode instore = (*instores)[it];
-            Assert(instore.getKind() == kind::STORE);
-            if (termSet.find(instore) != termSet.end()
-                && !d_equalityEngine->areEqual(instore[1], n[1]))
-            {
-              Node r = nm->mkNode(kind::SELECT, instore, n[1]);
-              if (termSet.find(r) == termSet.end()
-                  && d_equalityEngine->hasTerm(r))
-              {
-                Trace("arrays::collectModelInfo") << "TheoryArrays::collectModelInfo, adding RIntro2(c) read: " << r << endl;
-                termSet.insert(r);
-                changed = true;
-              }
-              r = nm->mkNode(kind::SELECT, instore[0], n[1]);
-              if (termSet.find(r) == termSet.end()
-                  && d_equalityEngine->hasTerm(r))
-              {
-                Trace("arrays::collectModelInfo") << "TheoryArrays::collectModelInfo, adding RIntro2(d) read: " << r << endl;
-                termSet.insert(r);
-                changed = true;
-              }
-            }
-          }
+          break;
         }
       }
     }
-  } while (changed);
-
-  // Send the equality engine information to the model
-  if (!m->assertEqualityEngine(d_equalityEngine, &termSet))
-  {
-    return false;
   }
 
   // Build a list of all the relevant reads, indexed by the store representative
@@ -2339,6 +2261,136 @@ TrustNode TheoryArrays::expandDefinition(Node node)
   return TrustNode::null();
 }
 
+void TheoryArrays::computeRelevantTerms(std::set<Node>& termSet,
+                                        bool includeShared)
+{
+  // include all standard terms
+  std::set<Kind> irrKinds;
+  computeRelevantTermsInternal(termSet, irrKinds, includeShared);
+
+  NodeManager* nm = NodeManager::currentNM();
+  // make sure RIntro1 reads are included in the relevant set of reads
+  eq::EqClassesIterator eqcs_i = eq::EqClassesIterator(d_equalityEngine);
+  for (; !eqcs_i.isFinished(); ++eqcs_i)
+  {
+    Node eqc = (*eqcs_i);
+    if (!eqc.getType().isArray())
+    {
+      // not an array, skip
+      continue;
+    }
+    eq::EqClassIterator eqc_i = eq::EqClassIterator(eqc, d_equalityEngine);
+    for (; !eqc_i.isFinished(); ++eqc_i)
+    {
+      Node n = *eqc_i;
+      if (termSet.find(n) != termSet.end())
+      {
+        if (n.getKind() == kind::STORE)
+        {
+          // Make sure RIntro1 reads are included
+          Node r = nm->mkNode(kind::SELECT, n, n[1]);
+          Trace("arrays::collectModelInfo")
+              << "TheoryArrays::collectModelInfo, adding RIntro1 read: " << r
+              << endl;
+          termSet.insert(r);
+        }
+      }
+    }
+  }
+
+  // Now do a fixed-point iteration to get all reads that need to be included
+  // because of RIntro2 rule
+  bool changed;
+  do
+  {
+    changed = false;
+    eqcs_i = eq::EqClassesIterator(d_equalityEngine);
+    for (; !eqcs_i.isFinished(); ++eqcs_i)
+    {
+      Node eqc = (*eqcs_i);
+      eq::EqClassIterator eqc_i = eq::EqClassIterator(eqc, d_equalityEngine);
+      for (; !eqc_i.isFinished(); ++eqc_i)
+      {
+        Node n = *eqc_i;
+        if (n.getKind() == kind::SELECT && termSet.find(n) != termSet.end())
+        {
+          // Find all terms equivalent to n[0] and get corresponding read terms
+          Node array_eqc = d_equalityEngine->getRepresentative(n[0]);
+          eq::EqClassIterator array_eqc_i =
+              eq::EqClassIterator(array_eqc, d_equalityEngine);
+          for (; !array_eqc_i.isFinished(); ++array_eqc_i)
+          {
+            Node arr = *array_eqc_i;
+            if (arr.getKind() == kind::STORE
+                && termSet.find(arr) != termSet.end()
+                && !d_equalityEngine->areEqual(arr[1], n[1]))
+            {
+              Node r = nm->mkNode(kind::SELECT, arr, n[1]);
+              if (termSet.find(r) == termSet.end()
+                  && d_equalityEngine->hasTerm(r))
+              {
+                Trace("arrays::collectModelInfo")
+                    << "TheoryArrays::collectModelInfo, adding RIntro2(a) "
+                       "read: "
+                    << r << endl;
+                termSet.insert(r);
+                changed = true;
+              }
+              r = nm->mkNode(kind::SELECT, arr[0], n[1]);
+              if (termSet.find(r) == termSet.end()
+                  && d_equalityEngine->hasTerm(r))
+              {
+                Trace("arrays::collectModelInfo")
+                    << "TheoryArrays::collectModelInfo, adding RIntro2(b) "
+                       "read: "
+                    << r << endl;
+                termSet.insert(r);
+                changed = true;
+              }
+            }
+          }
+
+          // Find all stores in which n[0] appears and get corresponding read
+          // terms
+          const CTNodeList* instores = d_infoMap.getInStores(array_eqc);
+          size_t it = 0;
+          for (; it < instores->size(); ++it)
+          {
+            TNode instore = (*instores)[it];
+            Assert(instore.getKind() == kind::STORE);
+            if (termSet.find(instore) != termSet.end()
+                && !d_equalityEngine->areEqual(instore[1], n[1]))
+            {
+              Node r = nm->mkNode(kind::SELECT, instore, n[1]);
+              if (termSet.find(r) == termSet.end()
+                  && d_equalityEngine->hasTerm(r))
+              {
+                Trace("arrays::collectModelInfo")
+                    << "TheoryArrays::collectModelInfo, adding RIntro2(c) "
+                       "read: "
+                    << r << endl;
+                termSet.insert(r);
+                changed = true;
+              }
+              r = nm->mkNode(kind::SELECT, instore[0], n[1]);
+              if (termSet.find(r) == termSet.end()
+                  && d_equalityEngine->hasTerm(r))
+              {
+                Trace("arrays::collectModelInfo")
+                    << "TheoryArrays::collectModelInfo, adding RIntro2(d) "
+                       "read: "
+                    << r << endl;
+                termSet.insert(r);
+                changed = true;
+              }
+            }
+          }
+        }
+      }
+    }
+  } while (changed);
+}
+
 }/* CVC4::theory::arrays namespace */
 }/* CVC4::theory namespace */
 }/* CVC4 namespace */
diff --git a/src/theory/arrays/theory_arrays.h b/src/theory/arrays/theory_arrays.h
index f1cd2ea14..530f8e0e1 100644
--- a/src/theory/arrays/theory_arrays.h
+++ b/src/theory/arrays/theory_arrays.h
@@ -499,6 +499,12 @@ class TheoryArrays : public Theory {
    */
   Node getNextDecisionRequest();
 
+  /**
+   * Compute relevant terms. This includes additional select nodes for the
+   * RIntro1 and RIntro2 rules.
+   */
+  void computeRelevantTerms(std::set<Node>& termSet,
+                            bool includeShared = true) override;
 };/* class TheoryArrays */
 
 }/* CVC4::theory::arrays namespace */
diff --git a/src/theory/datatypes/theory_datatypes.cpp b/src/theory/datatypes/theory_datatypes.cpp
index 4b38ad6bd..e625f57eb 100644
--- a/src/theory/datatypes/theory_datatypes.cpp
+++ b/src/theory/datatypes/theory_datatypes.cpp
@@ -1522,11 +1522,11 @@ bool TheoryDatatypes::collectModelInfo(TheoryModel* m)
   Trace("dt-model") << std::endl;
   printModelDebug( "dt-model" );
   Trace("dt-model") << std::endl;
-  
-  set<Node> termSet;
-  
+
+  std::set<Node> termSet;
+
   // Compute terms appearing in assertions and shared terms, and in inferred equalities
-  getRelevantTerms(termSet);
+  computeRelevantTerms(termSet);
 
   //combine the equality engine
   if (!m->assertEqualityEngine(d_equalityEngine, &termSet))
@@ -2250,12 +2250,14 @@ Node TheoryDatatypes::mkAnd( std::vector< TNode >& assumptions ) {
   }
 }
 
-void TheoryDatatypes::getRelevantTerms( std::set<Node>& termSet ) {
+void TheoryDatatypes::computeRelevantTerms(std::set<Node>& termSet,
+                                           bool includeShared)
+{
   // Compute terms appearing in assertions and shared terms
-  std::set<Kind> irr_kinds;
+  std::set<Kind> irrKinds;
   // testers are not relevant for model construction
-  irr_kinds.insert(APPLY_TESTER);
-  computeRelevantTerms(termSet, irr_kinds);
+  irrKinds.insert(APPLY_TESTER);
+  computeRelevantTermsInternal(termSet, irrKinds, includeShared);
 
   Trace("dt-cmi") << "Have " << termSet.size() << " relevant terms..."
                   << std::endl;
diff --git a/src/theory/datatypes/theory_datatypes.h b/src/theory/datatypes/theory_datatypes.h
index a68caca94..bdc13b5e5 100644
--- a/src/theory/datatypes/theory_datatypes.h
+++ b/src/theory/datatypes/theory_datatypes.h
@@ -357,8 +357,6 @@ private:
   void instantiate( EqcInfo* eqc, Node n );
   /** must communicate fact */
   bool mustCommunicateFact( Node n, Node exp );
-  /** get relevant terms */
-  void getRelevantTerms( std::set<Node>& termSet );
 private:
   //equality queries
   bool hasTerm( TNode a );
@@ -367,7 +365,13 @@ private:
   bool areCareDisequal( TNode x, TNode y );
   TNode getRepresentative( TNode a );
 
- private:
+  /**
+   * Compute relevant terms. In addition to all terms in assertions and shared
+   * terms, this includes datatypes in non-singleton equivalence classes.
+   */
+  void computeRelevantTerms(std::set<Node>& termSet,
+                            bool includeShared = true) override;
+
   /** sygus symmetry breaking utility */
   std::unique_ptr<SygusExtension> d_sygusExtension;
 
diff --git a/src/theory/theory.cpp b/src/theory/theory.cpp
index 4f0cbdb6a..02f84526f 100644
--- a/src/theory/theory.cpp
+++ b/src/theory/theory.cpp
@@ -343,6 +343,23 @@ std::unordered_set<TNode, TNodeHashFunction> Theory::currentlySharedTerms() cons
   return currentlyShared;
 }
 
+bool Theory::collectModelInfo(TheoryModel* m)
+{
+  std::set<Node> termSet;
+  // Compute terms appearing in assertions and shared terms
+  computeRelevantTerms(termSet);
+  // if we are using an equality engine, assert it to the model
+  if (d_equalityEngine != nullptr)
+  {
+    if (!m->assertEqualityEngine(d_equalityEngine, &termSet))
+    {
+      return false;
+    }
+  }
+  // now, collect theory-specific value assigments
+  return collectModelValues(m, termSet);
+}
+
 void Theory::collectTerms(TNode n,
                           set<Kind>& irrKinds,
                           set<Node>& termSet) const
@@ -365,16 +382,9 @@ void Theory::collectTerms(TNode n,
   }
 }
 
-
-void Theory::computeRelevantTerms(set<Node>& termSet, bool includeShared) const
-{
-  set<Kind> irrKinds;
-  computeRelevantTerms(termSet, irrKinds, includeShared);
-}
-
-void Theory::computeRelevantTerms(set<Node>& termSet,
-                                  set<Kind>& irrKinds,
-                                  bool includeShared) const
+void Theory::computeRelevantTermsInternal(std::set<Node>& termSet,
+                                          std::set<Kind>& irrKinds,
+                                          bool includeShared) const
 {
   // Collect all terms appearing in assertions
   irrKinds.insert(kind::EQUAL);
@@ -394,6 +404,17 @@ void Theory::computeRelevantTerms(set<Node>& termSet,
   }
 }
 
+void Theory::computeRelevantTerms(std::set<Node>& termSet, bool includeShared)
+{
+  std::set<Kind> irrKinds;
+  computeRelevantTermsInternal(termSet, irrKinds, includeShared);
+}
+
+bool Theory::collectModelValues(TheoryModel* m, std::set<Node>& termSet)
+{
+  return true;
+}
+
 Theory::PPAssertStatus Theory::ppAssert(TNode in,
                                         SubstitutionMap& outSubstitutions)
 {
diff --git a/src/theory/theory.h b/src/theory/theory.h
index 4feeac394..78c6e34cb 100644
--- a/src/theory/theory.h
+++ b/src/theory/theory.h
@@ -183,13 +183,7 @@ class Theory {
    */
   context::CDList<TNode> d_sharedTerms;
 
-  /**
-   * Helper function for computeRelevantTerms
-   */
-  void collectTerms(TNode n,
-                    std::set<Kind>& irrKinds,
-                    std::set<Node>& termSet) const;
-
+  //---------------------------------- collect model info
   /**
    * Scans the current set of assertions and shared terms top-down
    * until a theory-leaf is reached, and adds all terms found to
@@ -203,11 +197,30 @@ class Theory {
    * includeShared: Whether to include shared terms in termSet. Notice that
    * shared terms are not influenced by irrKinds.
    */
-  void computeRelevantTerms(std::set<Node>& termSet,
-                            std::set<Kind>& irrKinds,
-                            bool includeShared = true) const;
-  /** same as above, but with empty irrKinds */
-  void computeRelevantTerms(std::set<Node>& termSet, bool includeShared = true) const;
+  void computeRelevantTermsInternal(std::set<Node>& termSet,
+                                    std::set<Kind>& irrKinds,
+                                    bool includeShared = true) const;
+  /**
+   * Helper function for computeRelevantTerms
+   */
+  void collectTerms(TNode n,
+                    std::set<Kind>& irrKinds,
+                    std::set<Node>& termSet) const;
+  /**
+   * Same as above, but with empty irrKinds. This version can be overridden
+   * by the theory, e.g. by restricting or extending the set of terms returned
+   * by computeRelevantTermsInternal, which is called by default with no
+   * irrKinds.
+   */
+  virtual void computeRelevantTerms(std::set<Node>& termSet,
+                                    bool includeShared = true);
+  /**
+   * Collect model values, after equality information is added to the model.
+   * The argument termSet is the set of relevant terms returned by
+   * computeRelevantTerms.
+   */
+  virtual bool collectModelValues(TheoryModel* m, std::set<Node>& termSet);
+  //---------------------------------- end collect model info
 
   /**
    * Construct a Theory.
@@ -619,7 +632,7 @@ class Theory {
    * This method returns true if and only if the equality engine of m is
    * consistent as a result of this call.
    */
-  virtual bool collectModelInfo(TheoryModel* m) { return true; }
+  virtual bool collectModelInfo(TheoryModel* m);
   /** if theories want to do something with model after building, do it here */
   virtual void postProcessModel( TheoryModel* m ){ }
   /**