6 files changed, 44 insertions, 19 deletions
diff --git a/src/util/Makefile.am b/src/util/Makefile.am
index a4462d824..d69f0edf0 100644
--- a/src/util/Makefile.am
+++ b/src/util/Makefile.am
@@ -137,7 +137,8 @@ EXTRA_DIST = \
 	array_store_all.i \
 	ascription_type.i \
 	rational.i \
-	hash.i
+	hash.i \
+	util_model.i
 
 DISTCLEANFILES = \
 	integer.h.tmp \
diff --git a/src/util/ite_removal.cpp b/src/util/ite_removal.cpp
index e8a539615..b6e17ba29 100644
--- a/src/util/ite_removal.cpp
+++ b/src/util/ite_removal.cpp
@@ -27,9 +27,6 @@ using namespace std;
 
 namespace CVC4 {
 
-struct IteRewriteAttrTag {};
-typedef expr::Attribute<IteRewriteAttrTag, Node> IteRewriteAttr;
-
 void RemoveITE::run(std::vector<Node>& output, IteSkolemMap& iteSkolemMap)
 {
   for (unsigned i = 0, i_end = output.size(); i < i_end; ++ i) {
@@ -43,9 +40,10 @@ Node RemoveITE::run(TNode node, std::vector<Node>& output,
   Debug("ite") << "removeITEs(" << node << ")" << endl;
 
   // The result may be cached already
-  Node cachedRewrite;
   NodeManager *nodeManager = NodeManager::currentNM();
-  if(nodeManager->getAttribute(node, IteRewriteAttr(), cachedRewrite)) {
+  ITECache::iterator i = d_iteCache.find(node);
+  if(i != d_iteCache.end()) {
+    Node cachedRewrite = (*i).second;
     Debug("ite") << "removeITEs: in-cache: " << cachedRewrite << endl;
     return cachedRewrite.isNull() ? Node(node) : cachedRewrite;
   }
@@ -64,7 +62,7 @@ Node RemoveITE::run(TNode node, std::vector<Node>& output,
       Debug("ite") << "removeITEs(" << node << ") => " << newAssertion << endl;
 
       // Attach the skolem
-      nodeManager->setAttribute(node, IteRewriteAttr(), skolem);
+      d_iteCache[node] = skolem;
 
       // Remove ITEs from the new assertion, rewrite it and push it to the output
       newAssertion = run(newAssertion, output, iteSkolemMap);
@@ -94,15 +92,15 @@ Node RemoveITE::run(TNode node, std::vector<Node>& output,
 
     // If changes, we rewrite
     if(somethingChanged) {
-      cachedRewrite = nodeManager->mkNode(node.getKind(), newChildren);
-      nodeManager->setAttribute(node, IteRewriteAttr(), cachedRewrite);
+      Node cachedRewrite = nodeManager->mkNode(node.getKind(), newChildren);
+      d_iteCache[node] = cachedRewrite;
       return cachedRewrite;
     } else {
-      nodeManager->setAttribute(node, IteRewriteAttr(), Node::null());
+      d_iteCache[node] = Node::null();
       return node;
     }
   } else {
-    nodeManager->setAttribute(node, IteRewriteAttr(), Node::null());
+    d_iteCache[node] = Node::null();
     return node;
   }
 }
diff --git a/src/util/ite_removal.h b/src/util/ite_removal.h
index 7327d4a64..452b2d8a9 100644
--- a/src/util/ite_removal.h
+++ b/src/util/ite_removal.h
@@ -23,15 +23,23 @@
 #include <vector>
 #include "expr/node.h"
 #include "util/dump.h"
+#include "context/context.h"
+#include "context/cdhashmap.h"
 
 namespace CVC4 {
 
 typedef std::hash_map<Node, unsigned, NodeHashFunction> IteSkolemMap;
 
 class RemoveITE {
+  typedef context::CDHashMap<Node, Node, NodeHashFunction> ITECache;
+  ITECache d_iteCache;
 
 public:
 
+  RemoveITE(context::UserContext* u) :
+    d_iteCache(u) {
+  }
+
   /**
    * Removes the ITE nodes by introducing skolem variables. All
    * additional assertions are pushed into assertions. iteSkolemMap
@@ -39,7 +47,7 @@ public:
    * assertions containing the new Boolean ite created in conjunction
    * with that skolem variable.
    */
-  static void run(std::vector<Node>& assertions, IteSkolemMap& iteSkolemMap);
+  void run(std::vector<Node>& assertions, IteSkolemMap& iteSkolemMap);
 
   /**
    * Removes the ITE from the node by introducing skolem
@@ -48,8 +56,8 @@ public:
    * variables to the index in assertions containing the new Boolean
    * ite created in conjunction with that skolem variable.
    */
-  static Node run(TNode node, std::vector<Node>& additionalAssertions,
-                  IteSkolemMap& iteSkolemMap);
+  Node run(TNode node, std::vector<Node>& additionalAssertions,
+           IteSkolemMap& iteSkolemMap);
 
 };/* class RemoveTTE */
 
diff --git a/src/util/util_model.cpp b/src/util/util_model.cpp
index 6dfe2c566..a5768c723 100644
--- a/src/util/util_model.cpp
+++ b/src/util/util_model.cpp
@@ -18,6 +18,7 @@
 #include "expr/command.h"
 #include "smt/smt_engine_scope.h"
 #include "smt/command_list.h"
+#include "printer/printer.h"
 
 #include <vector>
 
@@ -25,6 +26,13 @@ using namespace std;
 
 namespace CVC4 {
 
+std::ostream& operator<<(std::ostream& out, Model& m) {
+  smt::SmtScope smts(&m.d_smt);
+  Expr::dag::Scope scope(out, false);
+  Printer::getPrinter(options::outputLanguage())->toStream(out, m);
+  return out;
+}
+
 Model::Model() :
   d_smt(*smt::currentSmtEngine()) {
 }
diff --git a/src/util/util_model.h b/src/util/util_model.h
index 97010dd45..07f964d46 100644
--- a/src/util/util_model.h
+++ b/src/util/util_model.h
@@ -29,14 +29,21 @@ namespace CVC4 {
 
 class CVC4_PUBLIC Command;
 class CVC4_PUBLIC SmtEngine;
+class CVC4_PUBLIC Model;
+
+std::ostream& operator<<(std::ostream&, Model&) CVC4_PUBLIC;
 
 class CVC4_PUBLIC Model {
-private:
+  friend std::ostream& operator<<(std::ostream&, Model&);
+
+protected:
   /** The SmtEngine we're associated to */
   const SmtEngine& d_smt;
-public:
-  /** construct the base class */
+
+  /** construct the base class; users cannot do this, only CVC4 internals */
   Model();
+
+public:
   /** virtual destructor */
   virtual ~Model() { }
   /** get number of commands to report */
@@ -48,8 +55,6 @@ public:
   virtual Expr getValue(Expr expr) = 0;
   /** get cardinality for sort */
   virtual Cardinality getCardinality(Type t) = 0;
-  /** write the model to a stream */
-  virtual void toStream(std::ostream& out) = 0;
 };/* class Model */
 
 class ModelBuilder
diff --git a/src/util/util_model.i b/src/util/util_model.i
new file mode 100644
index 000000000..0d1b3bc81
--- /dev/null
+++ b/src/util/util_model.i
@@ -0,0 +1,5 @@
+%{
+#include "util/util_model.h"
+%}
+
+%include "util/util_model.h"