From dfcd935acb928ff27c4b24a89de37338411d2543 Mon Sep 17 00:00:00 2001
From: Aina Niemetz <aina.niemetz@gmail.com>
Date: Wed, 28 Feb 2018 12:19:25 -0800
Subject: [PATCH] SmtEngine::getAssignment now returns a vector of assignments.
 (#1628)

---
 src/smt/command.cpp    | 18 ++++++++++-
 src/smt/smt_engine.cpp | 70 ++++++++++++++++++++----------------------
 src/smt/smt_engine.h   |  2 +-
 3 files changed, 51 insertions(+), 39 deletions(-)

diff --git a/src/smt/command.cpp b/src/smt/command.cpp
index 25479a20b..040d87250 100644
--- a/src/smt/command.cpp
+++ b/src/smt/command.cpp
@@ -1321,7 +1321,23 @@ void GetAssignmentCommand::invoke(SmtEngine* smtEngine)
 {
   try
   {
-    d_result = smtEngine->getAssignment();
+    std::vector<std::pair<Expr,Expr>> assignments  = smtEngine->getAssignment();
+    vector<SExpr> sexprs;
+    for (const auto& p : assignments)
+    {
+      vector<SExpr> v;
+      if (p.first.getKind() == kind::APPLY)
+      {
+        v.emplace_back(SExpr::Keyword(p.first.getOperator().toString()));
+      }
+      else
+      {
+        v.emplace_back(SExpr::Keyword(p.first.toString()));
+      }
+      v.emplace_back(SExpr::Keyword(p.second.toString()));
+      sexprs.emplace_back(v);
+    }
+    d_result = SExpr(sexprs);
     d_commandStatus = CommandSuccess::instance();
   }
   catch (RecoverableModalException& e)
diff --git a/src/smt/smt_engine.cpp b/src/smt/smt_engine.cpp
index 8176ba3e9..19236f881 100644
--- a/src/smt/smt_engine.cpp
+++ b/src/smt/smt_engine.cpp
@@ -5148,7 +5148,8 @@ bool SmtEngine::addToAssignment(const Expr& ex) {
 }
 
 // TODO(#1108): Simplify the error reporting of this method.
-CVC4::SExpr SmtEngine::getAssignment() {
+vector<pair<Expr, Expr>> SmtEngine::getAssignment()
+{
   Trace("smt") << "SMT getAssignment()" << endl;
   SmtScope smts(this);
   finalOptionsAreSet();
@@ -5170,50 +5171,45 @@ CVC4::SExpr SmtEngine::getAssignment() {
     throw RecoverableModalException(msg);
   }
 
-  if(d_assignments == NULL) {
-    return SExpr(vector<SExpr>());
-  }
-
-  vector<SExpr> sexprs;
-  TypeNode boolType = d_nodeManager->booleanType();
-  TheoryModel* m = d_theoryEngine->getModel();
-  for(AssignmentSet::key_iterator i = d_assignments->key_begin(),
-        iend = d_assignments->key_end();
-      i != iend;
-      ++i) {
-    Assert((*i).getType() == boolType);
+  vector<pair<Expr,Expr>> res;
+  if (d_assignments != nullptr)
+  {
+    TypeNode boolType = d_nodeManager->booleanType();
+    TheoryModel* m = d_theoryEngine->getModel();
+    for (AssignmentSet::key_iterator i = d_assignments->key_begin(),
+                                     iend = d_assignments->key_end();
+         i != iend;
+         ++i)
+    {
+      Node as = *i;
+      Assert(as.getType() == boolType);
 
-    Trace("smt") << "--- getting value of " << *i << endl;
+      Trace("smt") << "--- getting value of " << as << endl;
 
-    // Expand, then normalize
-    unordered_map<Node, Node, NodeHashFunction> cache;
-    Node n = d_private->expandDefinitions(*i, cache);
-    n = Rewriter::rewrite(n);
+      // Expand, then normalize
+      unordered_map<Node, Node, NodeHashFunction> cache;
+      Node n = d_private->expandDefinitions(as, cache);
+      n = Rewriter::rewrite(n);
 
-    Trace("smt") << "--- getting value of " << n << endl;
-    Node resultNode;
-    if(m != NULL) {
-      resultNode = m->getValue(n);
-    }
+      Trace("smt") << "--- getting value of " << n << endl;
+      Node resultNode;
+      if (m != nullptr)
+      {
+        resultNode = m->getValue(n);
+      }
 
-    // type-check the result we got
-    Assert(resultNode.isNull() || resultNode.getType() == boolType);
+      // type-check the result we got
+      Assert(resultNode.isNull() || resultNode.getType() == boolType);
 
-    // ensure it's a constant
-    Assert(resultNode.isConst());
+      // ensure it's a constant
+      Assert(resultNode.isConst());
 
-    vector<SExpr> v;
-    if((*i).getKind() == kind::APPLY) {
-      Assert((*i).getNumChildren() == 0);
-      v.push_back(SExpr(SExpr::Keyword((*i).getOperator().toString())));
-    } else {
-      Assert((*i).isVar());
-      v.push_back(SExpr(SExpr::Keyword((*i).toString())));
+      Assert(as.getKind() == kind::APPLY || as.isVar());
+      Assert(as.getKind() != kind::APPLY || as.getNumChildren() == 0);
+      res.emplace_back(as.toExpr(), resultNode.toExpr());
     }
-    v.push_back(SExpr(SExpr::Keyword(resultNode.toString())));
-    sexprs.push_back(SExpr(v));
   }
-  return SExpr(sexprs);
+  return res;
 }
 
 void SmtEngine::addToModelCommandAndDump(const Command& c, uint32_t flags, bool userVisible, const char* dumpTag) {
diff --git a/src/smt/smt_engine.h b/src/smt/smt_engine.h
index e768bf826..c8601a23d 100644
--- a/src/smt/smt_engine.h
+++ b/src/smt/smt_engine.h
@@ -598,7 +598,7 @@ class CVC4_PUBLIC SmtEngine {
    * INVALID query).  Only permitted if the SmtEngine is set to
    * operate interactively and produce-assignments is on.
    */
-  CVC4::SExpr getAssignment();
+  std::vector<std::pair<Expr, Expr> > getAssignment();
 
   /**
    * Get the last proof (only if immediately preceded by an UNSAT
-- 
2.30.2