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 ){ }
   /**