1 files changed, 25 insertions, 6 deletions

diff --git a/src/theory/arith/arithvar_node_map.h b/src/theory/arith/arithvar_node_map.h
index 3a2cff236..e0e589e57 100644
--- a/src/theory/arith/arithvar_node_map.h
+++ b/src/theory/arith/arithvar_node_map.h
@@ -40,31 +40,50 @@ private:
   ArithVarToNodeMap d_arithVarToNodeMap;
 
 public:
+
+  typedef ArithVarToNodeMap::const_iterator var_iterator;
+
   ArithVarNodeMap() {}
 
   inline bool hasArithVar(TNode x) const {
     return d_nodeToArithVarMap.find(x) != d_nodeToArithVarMap.end();
   }
 
+  inline bool hasNode(ArithVar a) const {
+    return d_arithVarToNodeMap.isKey(a);
+  }
+
   inline ArithVar asArithVar(TNode x) const{
     Assert(hasArithVar(x));
     Assert((d_nodeToArithVarMap.find(x))->second <= ARITHVAR_SENTINEL);
     return (d_nodeToArithVarMap.find(x))->second;
   }
+
   inline Node asNode(ArithVar a) const{
-    Assert(d_arithVarToNodeMap.find(a) != d_arithVarToNodeMap.end());
-    return (d_arithVarToNodeMap.find(a))->second;
+    Assert(hasNode(a));
+    return d_arithVarToNodeMap[a];
   }
 
   inline void setArithVar(TNode x, ArithVar a){
     Assert(!hasArithVar(x));
-    Assert(d_arithVarToNodeMap.find(a) == d_arithVarToNodeMap.end());
-    d_arithVarToNodeMap[a] = x;
+    Assert(!d_arithVarToNodeMap.isKey(a));
+    d_arithVarToNodeMap.set(a, x);
     d_nodeToArithVarMap[x] = a;
   }
 
-  const ArithVarToNodeMap& getArithVarToNodeMap() const {
-    return d_arithVarToNodeMap;
+  inline void remove(ArithVar x){
+    Assert(hasNode(x));
+    Node node = asNode(x);
+
+    d_nodeToArithVarMap.erase(d_nodeToArithVarMap.find(node));
+    d_arithVarToNodeMap.remove(x);
+  }
+
+  var_iterator var_begin() const {
+    return d_arithVarToNodeMap.begin();
+  }
+  var_iterator var_end() const {
+    return d_arithVarToNodeMap.end();
   }
 
 };/* class ArithVarNodeMap */