Cleaning up the printing of theory model representative sets. (#1538)

2 files changed, 225 insertions, 201 deletions

diff --git a/src/printer/cvc/cvc_printer.cpp b/src/printer/cvc/cvc_printer.cpp
index f20cb7cce..27105c3b4 100644
--- a/src/printer/cvc/cvc_printer.cpp
+++ b/src/printer/cvc/cvc_printer.cpp
@@ -1010,121 +1010,121 @@ void CvcPrinter::toStream(std::ostream& out, const CommandStatus* s) const
 
 }/* CvcPrinter::toStream(CommandStatus*) */
 
-void CvcPrinter::toStream(std::ostream& out,
-                          const Model& m,
-                          const Command* c) const
+namespace {
+
+void DeclareTypeCommandToStream(std::ostream& out,
+                                const theory::TheoryModel& model,
+                                const DeclareTypeCommand& command)
 {
-  const theory::TheoryModel& tm = (const theory::TheoryModel&) m;
-  if(dynamic_cast<const DeclareTypeCommand*>(c) != NULL) {
-    TypeNode tn = TypeNode::fromType( ((const DeclareTypeCommand*)c)->getType() );
-    if (options::modelUninterpDtEnum() && tn.isSort())
+  TypeNode type_node = TypeNode::fromType(command.getType());
+  const std::vector<Node>* type_reps =
+      model.getRepSet()->getTypeRepsOrNull(type_node);
+  if (options::modelUninterpDtEnum() && type_node.isSort()
+      && type_reps != nullptr)
+  {
+    out << "DATATYPE" << std::endl;
+    out << "  " << command.getSymbol() << " = ";
+    for (size_t i = 0; i < type_reps->size(); i++)
     {
-      const theory::RepSet* rs = tm.getRepSet();
-      if (rs->d_type_reps.find(tn) != rs->d_type_reps.end())
+      if (i > 0)
       {
-        out << "DATATYPE" << std::endl;
-        out << "  " << dynamic_cast<const DeclareTypeCommand*>(c)->getSymbol()
-            << " = ";
-        for (size_t i = 0; i < (*rs->d_type_reps.find(tn)).second.size(); i++)
-        {
-          if (i > 0)
-          {
-            out << "| ";
-          }
-          out << (*rs->d_type_reps.find(tn)).second[i] << " ";
-        }
-        out << std::endl << "END;" << std::endl;
+        out << "| ";
       }
-      else
+      out << (*type_reps)[i] << " ";
+    }
+    out << std::endl << "END;" << std::endl;
+  }
+  else if (type_node.isSort() && type_reps != nullptr)
+  {
+    out << "% cardinality of " << type_node << " is " << type_reps->size()
+        << std::endl;
+    out << command << std::endl;
+    for (Node type_rep : *type_reps)
+    {
+      if (type_rep.isVar())
       {
-        if (tn.isSort())
-        {
-          // print the cardinality
-          if (rs->d_type_reps.find(tn) != rs->d_type_reps.end())
-          {
-            out << "% cardinality of " << tn << " is "
-                << (*rs->d_type_reps.find(tn)).second.size() << std::endl;
-          }
-        }
-        out << c << std::endl;
-        if (tn.isSort())
-        {
-          // print the representatives
-          if (rs->d_type_reps.find(tn) != rs->d_type_reps.end())
-          {
-            for (size_t i = 0; i < (*rs->d_type_reps.find(tn)).second.size();
-                 i++)
-            {
-              if ((*rs->d_type_reps.find(tn)).second[i].isVar())
-              {
-                out << (*rs->d_type_reps.find(tn)).second[i] << " : " << tn
-                    << ";" << std::endl;
-              }
-              else
-              {
-                out << "% rep: " << (*rs->d_type_reps.find(tn)).second[i]
-                    << std::endl;
-              }
-            }
-          }
-        }
+        out << type_rep << " : " << type_node << ";" << std::endl;
       }
-    }
-  } else if(dynamic_cast<const DeclareFunctionCommand*>(c) != NULL) {
-    Node n = Node::fromExpr( ((const DeclareFunctionCommand*)c)->getFunction() );
-    if(n.getKind() == kind::SKOLEM) {
-      // don't print out internal stuff
-      return;
-    }
-    TypeNode tn = n.getType();
-    out << n << " : ";
-    if( tn.isFunction() || tn.isPredicate() ){
-      out << "(";
-      for( size_t i=0; i<tn.getNumChildren()-1; i++ ){
-        if( i>0 ) out << ", ";
-        out << tn[i];
+      else
+      {
+        out << "% rep: " << type_rep << std::endl;
       }
-      out << ") -> " << tn.getRangeType();
-    }else{
-      out << tn;
     }
-    Node val = Node::fromExpr(tm.getSmtEngine()->getValue(n.toExpr()));
-    if( options::modelUninterpDtEnum() && val.getKind() == kind::STORE ) {
-      const theory::RepSet* rs = tm.getRepSet();
-      TypeNode tn = val[1].getType();
-      if (tn.isSort() && rs->d_type_reps.find(tn) != rs->d_type_reps.end())
+  }
+  else
+  {
+    out << command << std::endl;
+  }
+}
+
+void DeclareFunctionCommandToStream(std::ostream& out,
+                                    const theory::TheoryModel& model,
+                                    const DeclareFunctionCommand& command)
+{
+  Node n = Node::fromExpr(command.getFunction());
+  if (n.getKind() == kind::SKOLEM)
+  {
+    // don't print out internal stuff
+    return;
+  }
+  TypeNode tn = n.getType();
+  out << n << " : ";
+  if (tn.isFunction() || tn.isPredicate())
+  {
+    out << "(";
+    for (size_t i = 0; i < tn.getNumChildren() - 1; i++)
+    {
+      if (i > 0)
       {
-        Cardinality indexCard((*rs->d_type_reps.find(tn)).second.size());
-        val = theory::arrays::TheoryArraysRewriter::normalizeConstant( val, indexCard );
+        out << ", ";
       }
+      out << tn[i];
     }
-    out << " = " << val << ";" << std::endl;
-
-/*
-    //for table format (work in progress)
-    bool printedModel = false;
-    if( tn.isFunction() ){
-      if( options::modelFormatMode()==MODEL_FORMAT_MODE_TABLE ){
-        //specialized table format for functions
-        RepSetIterator riter( &d_rep_set );
-        riter.setFunctionDomain( n );
-        while( !riter.isFinished() ){
-          std::vector< Node > children;
-          children.push_back( n );
-          for( int i=0; i<riter.getNumTerms(); i++ ){
-            children.push_back( riter.getTerm( i ) );
-          }
-          Node nn = NodeManager::currentNM()->mkNode( APPLY_UF, children );
-          Node val = getValue( nn );
-          out << val << " ";
-          riter.increment();
-        }
-        printedModel = true;
+    out << ") -> " << tn.getRangeType();
+  }
+  else
+  {
+    out << tn;
+  }
+  Node val = Node::fromExpr(model.getSmtEngine()->getValue(n.toExpr()));
+  if (options::modelUninterpDtEnum() && val.getKind() == kind::STORE)
+  {
+    TypeNode type_node = val[1].getType();
+    if (tn.isSort())
+    {
+      if (const std::vector<Node>* type_reps =
+              model.getRepSet()->getTypeRepsOrNull(type_node))
+      {
+        Cardinality indexCard(type_reps->size());
+        val = theory::arrays::TheoryArraysRewriter::normalizeConstant(
+            val, indexCard);
       }
     }
-*/
-  }else{
-    out << c << std::endl;
+  }
+  out << " = " << val << ";" << std::endl;
+}
+
+}  // namespace
+
+void CvcPrinter::toStream(std::ostream& out,
+                          const Model& model,
+                          const Command* command) const
+{
+  const auto* theory_model = dynamic_cast<const theory::TheoryModel*>(&model);
+  AlwaysAssert(theory_model != nullptr);
+  if (const auto* declare_type_command =
+          dynamic_cast<const DeclareTypeCommand*>(command))
+  {
+    DeclareTypeCommandToStream(out, *theory_model, *declare_type_command);
+  }
+  else if (const auto* dfc =
+               dynamic_cast<const DeclareFunctionCommand*>(command))
+  {
+    DeclareFunctionCommandToStream(out, *theory_model, *dfc);
+  }
+  else
+  {
+    out << command << std::endl;
   }
 }
 
diff --git a/src/printer/smt2/smt2_printer.cpp b/src/printer/smt2/smt2_printer.cpp
index 54fc10719..24fd0924f 100644
--- a/src/printer/smt2/smt2_printer.cpp
+++ b/src/printer/smt2/smt2_printer.cpp
@@ -1277,116 +1277,140 @@ void Smt2Printer::toStream(std::ostream& out, const Model& m) const
   }
 }
 
-void Smt2Printer::toStream(std::ostream& out,
-                           const Model& m,
-                           const Command* c) const
+namespace {
+
+void DeclareTypeCommandToStream(std::ostream& out,
+                                const theory::TheoryModel& model,
+                                const DeclareTypeCommand& command,
+                                Variant variant)
 {
-  const theory::TheoryModel& tm = (const theory::TheoryModel&) m;
-  if(dynamic_cast<const DeclareTypeCommand*>(c) != NULL) {
-    TypeNode tn = TypeNode::fromType( ((const DeclareTypeCommand*)c)->getType() );
-    const theory::RepSet* rs = tm.getRepSet();
-    const std::map<TypeNode, std::vector<Node> >& type_reps = rs->d_type_reps;
-
-    std::map< TypeNode, std::vector< Node > >::const_iterator tn_iterator = type_reps.find( tn );
-    if( options::modelUninterpDtEnum() && tn.isSort() && tn_iterator != type_reps.end() ){
-      if(d_variant == smt2_6_variant) {
-        out << "(declare-datatypes ((" << dynamic_cast<const DeclareTypeCommand*>(c)->getSymbol() << " 0)) (";
-      }else{
-        out << "(declare-datatypes () ((" << dynamic_cast<const DeclareTypeCommand*>(c)->getSymbol() << " ";
-      }
-      for( size_t i=0, N = tn_iterator->second.size(); i < N; i++ ){
-        out << "(" << (*tn_iterator).second[i] << ")";
-      }
-      out << ")))" << endl;
-    } else {
-      if( tn.isSort() ){
-        //print the cardinality
-        if( tn_iterator != type_reps.end() ) {
-          out << "; cardinality of " << tn << " is " << tn_iterator->second.size() << endl;
-        }
+  TypeNode tn = TypeNode::fromType(command.getType());
+  const std::vector<Node>* type_refs = model.getRepSet()->getTypeRepsOrNull(tn);
+  if (options::modelUninterpDtEnum() && tn.isSort() && type_refs != nullptr)
+  {
+    if (variant == smt2_6_variant)
+    {
+      out << "(declare-datatypes ((" << command.getSymbol() << " 0)) (";
+    }
+    else
+    {
+      out << "(declare-datatypes () ((" << command.getSymbol() << " ";
+    }
+    for (Node type_ref : *type_refs)
+    {
+      out << "(" << type_ref << ")";
+    }
+    out << ")))" << endl;
+  }
+  else if (tn.isSort() && type_refs != nullptr)
+  {
+    // print the cardinality
+    out << "; cardinality of " << tn << " is " << type_refs->size() << endl;
+    out << command << endl;
+    // print the representatives
+    for (Node type_ref : *type_refs)
+    {
+      if (type_ref.isVar())
+      {
+        out << "(declare-fun " << quoteSymbol(type_ref) << " () " << tn << ")"
+            << endl;
       }
-      out << c << endl;
-      if( tn.isSort() ){
-        //print the representatives
-        if( tn_iterator != type_reps.end() ){
-          for( size_t i = 0, N = (*tn_iterator).second.size(); i < N; i++ ){
-            TNode current = (*tn_iterator).second[i];
-            if( current.isVar() ){
-              out << "(declare-fun " << quoteSymbol(current) << " () " << tn << ")" << endl;
-            }else{
-              out << "; rep: " << current << endl;
-            }
-          }
-        }
+      else
+      {
+        out << "; rep: " << type_ref << endl;
       }
     }
-  } else if(dynamic_cast<const DeclareFunctionCommand*>(c) != NULL) {
-    const DeclareFunctionCommand* dfc = (const DeclareFunctionCommand*)c;
-    Node n = Node::fromExpr( dfc->getFunction() );
-    if(dfc->getPrintInModelSetByUser()) {
-      if(!dfc->getPrintInModel()) {
-        return;
-      }
-    } else if(n.getKind() == kind::SKOLEM) {
-      // don't print out internal stuff
+  }
+  else
+  {
+    out << command << endl;
+  }
+}
+
+void DeclareFunctionCommandToStream(std::ostream& out,
+                                    const theory::TheoryModel& model,
+                                    const DeclareFunctionCommand& command)
+{
+  Node n = Node::fromExpr(command.getFunction());
+  if (command.getPrintInModelSetByUser())
+  {
+    if (!command.getPrintInModel())
+    {
       return;
     }
-    Node val = Node::fromExpr(tm.getSmtEngine()->getValue(n.toExpr()));
-    if(val.getKind() == kind::LAMBDA) {
-      out << "(define-fun " << n << " " << val[0]
-          << " " << n.getType().getRangeType()
-          << " " << val[1] << ")" << endl;
-    } else {
-      if( options::modelUninterpDtEnum() && val.getKind() == kind::STORE ) {
-        TypeNode tn = val[1].getType();
-        const theory::RepSet* rs = tm.getRepSet();
-        if (tn.isSort() && rs->d_type_reps.find(tn) != rs->d_type_reps.end())
-        {
-          Cardinality indexCard((*rs->d_type_reps.find(tn)).second.size());
-          val = theory::arrays::TheoryArraysRewriter::normalizeConstant( val, indexCard );
-        }
-      }
-      out << "(define-fun " << n << " () "
-          << n.getType() << " ";
-      if(val.getType().isInteger() && n.getType().isReal() && !n.getType().isInteger()) {
-        //toStreamReal(out, val, true);
-        toStreamRational(out, val.getConst<Rational>(), true);
-        //out << val << ".0";
-      } else {
-        out << val;
+  }
+  else if (n.getKind() == kind::SKOLEM)
+  {
+    // don't print out internal stuff
+    return;
+  }
+  Node val = Node::fromExpr(model.getSmtEngine()->getValue(n.toExpr()));
+  if (val.getKind() == kind::LAMBDA)
+  {
+    out << "(define-fun " << n << " " << val[0] << " "
+        << n.getType().getRangeType() << " " << val[1] << ")" << endl;
+  }
+  else
+  {
+    if (options::modelUninterpDtEnum() && val.getKind() == kind::STORE)
+    {
+      TypeNode tn = val[1].getType();
+      const std::vector<Node>* type_refs =
+          model.getRepSet()->getTypeRepsOrNull(tn);
+      if (tn.isSort() && type_refs != nullptr)
+      {
+        Cardinality indexCard(type_refs->size());
+        val = theory::arrays::TheoryArraysRewriter::normalizeConstant(
+            val, indexCard);
       }
-      out << ")" << endl;
     }
-/*
-    //for table format (work in progress)
-    bool printedModel = false;
-    if( tn.isFunction() ){
-      if( options::modelFormatMode()==MODEL_FORMAT_MODE_TABLE ){
-        //specialized table format for functions
-        RepSetIterator riter( &d_rep_set );
-        riter.setFunctionDomain( n );
-        while( !riter.isFinished() ){
-          std::vector< Node > children;
-          children.push_back( n );
-          for( int i=0; i<riter.getNumTerms(); i++ ){
-            children.push_back( riter.getTerm( i ) );
-          }
-          Node nn = NodeManager::currentNM()->mkNode( APPLY_UF, children );
-          Node val = getValue( nn );
-          out << val << " ";
-          riter.increment();
-        }
-        printedModel = true;
-      }
+    out << "(define-fun " << n << " () " << n.getType() << " ";
+    if (val.getType().isInteger() && n.getType().isReal()
+        && !n.getType().isInteger())
+    {
+      // toStreamReal(out, val, true);
+      toStreamRational(out, val.getConst<Rational>(), true);
+      // out << val << ".0";
     }
-*/
-  } else {
-    DatatypeDeclarationCommand* c1 = (DatatypeDeclarationCommand*)c;
-    const vector<DatatypeType>& datatypes = c1->getDatatypes();
-    if (!datatypes[0].isTuple()) {
-      out << c << endl;
+    else
+    {
+      out << val;
+    }
+    out << ")" << endl;
+  }
+}
+
+}  // namespace
+
+void Smt2Printer::toStream(std::ostream& out,
+                           const Model& model,
+                           const Command* command) const
+{
+  const theory::TheoryModel* theory_model =
+      dynamic_cast<const theory::TheoryModel*>(&model);
+  AlwaysAssert(theory_model != nullptr);
+  if (const DeclareTypeCommand* dtc =
+          dynamic_cast<const DeclareTypeCommand*>(command))
+  {
+    DeclareTypeCommandToStream(out, *theory_model, *dtc, d_variant);
+  }
+  else if (const DeclareFunctionCommand* dfc =
+               dynamic_cast<const DeclareFunctionCommand*>(command))
+  {
+    DeclareFunctionCommandToStream(out, *theory_model, *dfc);
+  }
+  else if (const DatatypeDeclarationCommand* datatype_declaration_command =
+               dynamic_cast<const DatatypeDeclarationCommand*>(command))
+  {
+    if (!datatype_declaration_command->getDatatypes()[0].isTuple())
+    {
+      out << command << endl;
     }
   }
+  else
+  {
+    Unreachable();
+  }
 }
 
 void Smt2Printer::toStreamSygus(std::ostream& out, TNode n) const