4 files changed, 65 insertions, 31 deletions

diff --git a/src/printer/smt2/smt2_printer.cpp b/src/printer/smt2/smt2_printer.cpp
index dbb89d857..e02d308da 100644
--- a/src/printer/smt2/smt2_printer.cpp
+++ b/src/printer/smt2/smt2_printer.cpp
@@ -1368,47 +1368,50 @@ void Smt2Printer::toStream(std::ostream& out,
           dynamic_cast<const DeclareTypeCommand*>(command))
   {
     // print out the DeclareTypeCommand
-    TypeNode tn = TypeNode::fromType((*dtc).getType());
-    const std::vector<Node>* type_refs =
-        theory_model->getRepSet()->getTypeRepsOrNull(tn);
-    if (options::modelUninterpDtEnum() && tn.isSort() && type_refs != nullptr)
+    Type t = (*dtc).getType();
+    if (!t.isSort())
     {
-      if (isVariant_2_6(d_variant))
-      {
-        out << "(declare-datatypes ((" << (*dtc).getSymbol() << " 0)) (";
-      }
-      else
-      {
-        out << "(declare-datatypes () ((" << (*dtc).getSymbol() << " ";
-      }
-      for (Node type_ref : *type_refs)
-      {
-        out << "(" << type_ref << ")";
-      }
-      out << ")))" << endl;
+      out << (*dtc) << endl;
     }
-    else if (tn.isSort() && type_refs != nullptr)
+    else
     {
-      // print the cardinality
-      out << "; cardinality of " << tn << " is " << type_refs->size() << endl;
-      out << (*dtc) << endl;
-      // print the representatives
-      for (Node type_ref : *type_refs)
+      std::vector<Expr> elements = theory_model->getDomainElements(t);
+      if (options::modelUninterpDtEnum())
       {
-        if (type_ref.isVar())
+        if (isVariant_2_6(d_variant))
         {
-          out << "(declare-fun " << quoteSymbol(type_ref) << " () " << tn << ")"
-              << endl;
+          out << "(declare-datatypes ((" << (*dtc).getSymbol() << " 0)) (";
         }
         else
         {
-          out << "; rep: " << type_ref << endl;
+          out << "(declare-datatypes () ((" << (*dtc).getSymbol() << " ";
+        }
+        for (const Expr& type_ref : elements)
+        {
+          out << "(" << type_ref << ")";
+        }
+        out << ")))" << endl;
+      }
+      else
+      {
+        // print the cardinality
+        out << "; cardinality of " << t << " is " << elements.size() << endl;
+        out << (*dtc) << endl;
+        // print the representatives
+        for (const Expr& type_ref : elements)
+        {
+          Node trn = Node::fromExpr(type_ref);
+          if (trn.isVar())
+          {
+            out << "(declare-fun " << quoteSymbol(trn) << " () " << t << ")"
+                << endl;
+          }
+          else
+          {
+            out << "; rep: " << trn << endl;
+          }
         }
       }
-    }
-    else
-    {
-      out << (*dtc) << endl;
     }
   }
   else if (const DeclareFunctionCommand* dfc =
diff --git a/src/smt/model.h b/src/smt/model.h
index b435fb5e2..06b616f9b 100644
--- a/src/smt/model.h
+++ b/src/smt/model.h
@@ -96,6 +96,14 @@ class Model {
    * is a predicate over t that indicates a property that t satisfies.
    */
   virtual std::vector<std::pair<Expr, Expr> > getApproximations() const = 0;
+  /** get the domain elements for uninterpreted sort t
+   *
+   * This method gets the interpretation of an uninterpreted sort t.
+   * All models interpret uninterpreted sorts t as finite sets
+   * of domain elements v_1, ..., v_n. This method returns this list for t in
+   * this model.
+   */
+  virtual std::vector<Expr> getDomainElements(Type t) const = 0;
 };/* class Model */
 
 class ModelBuilder {
diff --git a/src/theory/theory_model.cpp b/src/theory/theory_model.cpp
index 8d6511854..168dc4c62 100644
--- a/src/theory/theory_model.cpp
+++ b/src/theory/theory_model.cpp
@@ -117,6 +117,27 @@ std::vector<std::pair<Expr, Expr> > TheoryModel::getApproximations() const
   return approx;
 }
 
+std::vector<Expr> TheoryModel::getDomainElements(Type t) const
+{
+  // must be an uninterpreted sort
+  Assert(t.isSort());
+  std::vector<Expr> elements;
+  TypeNode tn = TypeNode::fromType(t);
+  const std::vector<Node>* type_refs = d_rep_set.getTypeRepsOrNull(tn);
+  if (type_refs == nullptr || type_refs->empty())
+  {
+    // This is called when t is a sort that does not occur in this model.
+    // Sorts are always interpreted as non-empty, thus we add a single element.
+    elements.push_back(t.mkGroundTerm());
+    return elements;
+  }
+  for (const Node& n : *type_refs)
+  {
+    elements.push_back(n.toExpr());
+  }
+  return elements;
+}
+
 Node TheoryModel::getValue(TNode n) const
 {
   //apply substitutions
diff --git a/src/theory/theory_model.h b/src/theory/theory_model.h
index baf3a401c..b120b4501 100644
--- a/src/theory/theory_model.h
+++ b/src/theory/theory_model.h
@@ -239,6 +239,8 @@ public:
   bool hasApproximations() const override;
   /** get approximations */
   std::vector<std::pair<Expr, Expr> > getApproximations() const override;
+  /** get domain elements for uninterpreted sort t */
+  std::vector<Expr> getDomainElements(Type t) const override;
   /** get the representative set object */
   const RepSet* getRepSet() const { return &d_rep_set; }
   /** get the representative set object (FIXME: remove this, see #1199) */