1 files changed, 46 insertions, 3 deletions
diff --git a/src/util/stats.cpp b/src/util/stats.cpp
index 5be9525a9..428f051e0 100644
--- a/src/util/stats.cpp
+++ b/src/util/stats.cpp
@@ -18,13 +18,33 @@
  **/
 
 #include "util/stats.h"
+#include "expr/node_manager.h"
+#include "expr/expr_manager_scope.h"
 
 using namespace CVC4;
 
-StatisticsRegistry::StatSet StatisticsRegistry::d_registeredStats;
-
 std::string Stat::s_delim(",");
 
+StatisticsRegistry* StatisticsRegistry::current() {
+  return NodeManager::currentNM()->getStatisticsRegistry();
+}
+
+void StatisticsRegistry::registerStat(Stat* s) throw(AssertionException) {
+#ifdef CVC4_STATISTICS_ON
+  StatSet& registeredStats = NodeManager::currentNM()->getStatisticsRegistry()->d_registeredStats;
+  AlwaysAssert(registeredStats.find(s) == registeredStats.end());
+  registeredStats.insert(s);
+#endif /* CVC4_STATISTICS_ON */
+}/* StatisticsRegistry::registerStat() */
+
+void StatisticsRegistry::unregisterStat(Stat* s) throw(AssertionException) {
+#ifdef CVC4_STATISTICS_ON
+  StatSet& registeredStats = NodeManager::currentNM()->getStatisticsRegistry()->d_registeredStats;
+  AlwaysAssert(registeredStats.find(s) != registeredStats.end());
+  registeredStats.erase(s);
+#endif /* CVC4_STATISTICS_ON */
+}/* StatisticsRegistry::unregisterStat() */
+
 void StatisticsRegistry::flushStatistics(std::ostream& out) {
 #ifdef CVC4_STATISTICS_ON
   for(StatSet::iterator i = d_registeredStats.begin();
@@ -34,5 +54,28 @@ void StatisticsRegistry::flushStatistics(std::ostream& out) {
     s->flushStat(out);
     out << std::endl;
   }
-#endif
+#endif /* CVC4_STATISTICS_ON */
 }/* StatisticsRegistry::flushStatistics() */
+
+StatisticsRegistry::const_iterator StatisticsRegistry::begin() {
+  return NodeManager::currentNM()->getStatisticsRegistry()->d_registeredStats.begin();
+}/* StatisticsRegistry::begin() */
+
+StatisticsRegistry::const_iterator StatisticsRegistry::end() {
+  return NodeManager::currentNM()->getStatisticsRegistry()->d_registeredStats.end();
+}/* StatisticsRegistry::end() */
+
+RegisterStatistic::RegisterStatistic(ExprManager& em, Stat* stat) :
+    d_em(&em), d_stat(stat) {
+  ExprManagerScope ems(*d_em);
+  StatisticsRegistry::registerStat(d_stat);
+}/* RegisterStatistic::RegisterStatistic(ExprManager&, Stat*) */
+
+RegisterStatistic::~RegisterStatistic() {
+  if(d_em != NULL) {
+    ExprManagerScope ems(*d_em);
+    StatisticsRegistry::unregisterStat(d_stat);
+  } else {
+    StatisticsRegistry::unregisterStat(d_stat);
+  }
+}/* RegisterStatistic::~RegisterStatistic() */