From d068ec8082928001e4abfbbe5490236df9485a87 Mon Sep 17 00:00:00 2001
From: Ming-Hung Tsai <mingnus@gmail.com>
Date: Sat, 27 Feb 2016 15:22:09 +0800
Subject: [PATCH] Add utility class binary_block_counter

---
 persistent-data/block_counter.h               | 30 +++++++++++++++++--
 .../data-structures/btree_counter.h           | 11 +++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/persistent-data/block_counter.h b/persistent-data/block_counter.h
index e4232b4..ea70f93 100644
--- a/persistent-data/block_counter.h
+++ b/persistent-data/block_counter.h
@@ -20,6 +20,7 @@
 #define BLOCK_COUNTER_H
 
 #include "block.h"
+#include "run_set.h"
 
 //----------------------------------------------------------------
 
@@ -32,7 +33,9 @@ namespace persistent_data {
 	public:
 		typedef std::map<block_address, unsigned> count_map;
 
-		void inc(block_address b) {
+		virtual ~block_counter() {}
+
+		virtual void inc(block_address b) {
 			count_map::iterator it = counts_.find(b);
 			if (it == counts_.end())
 				counts_.insert(make_pair(b, 1));
@@ -40,7 +43,7 @@ namespace persistent_data {
 				it->second++;
 		}
 
-		unsigned get_count(block_address b) const {
+		virtual unsigned get_count(block_address b) const {
 			count_map::const_iterator it = counts_.find(b);
 			return (it == counts_.end()) ? 0 : it->second;
 		}
@@ -52,6 +55,29 @@ namespace persistent_data {
 	private:
 		count_map counts_;
 	};
+
+	//----------------------------------------------------------------
+	// Little helper class that keeps track of which blocks
+	// are referenced.
+	//----------------------------------------------------------------
+	class binary_block_counter : public block_counter {
+	public:
+		virtual ~binary_block_counter() {}
+
+		virtual void inc(block_address b) {
+			visited_.add(b);
+		}
+
+		virtual unsigned get_count(block_address b) const {
+			return visited_.member(b) ? 1 : 0;
+		}
+
+		base::run_set<block_address> const& get_visited() const {
+			return visited_;
+		}
+	private:
+		base::run_set<block_address> visited_;
+	};
 }
 
 //----------------------------------------------------------------
diff --git a/persistent-data/data-structures/btree_counter.h b/persistent-data/data-structures/btree_counter.h
index c215626..6e52c29 100644
--- a/persistent-data/data-structures/btree_counter.h
+++ b/persistent-data/data-structures/btree_counter.h
@@ -100,6 +100,17 @@ namespace persistent_data {
 		btree_count_detail::counting_visitor<NoopValueVisitor, noop_damage_visitor, Levels, ValueTraits, ValueCounter> v(noop_vv, noop_dv, bc, vc);
 		tree.visit_depth_first(v);
 	}
+
+	template <unsigned Levels, typename ValueTraits>
+	void count_btree_blocks(btree<Levels, ValueTraits> const &tree, block_counter &bc) {
+		typedef noop_value_visitor<typename ValueTraits::value_type> NoopValueVisitor;
+		NoopValueVisitor noop_vv;
+		noop_damage_visitor noop_dv;
+		typedef noop_value_counter<typename ValueTraits::value_type> NoopValueCounter;
+		NoopValueCounter vc;
+		btree_count_detail::counting_visitor<NoopValueVisitor, noop_damage_visitor, Levels, ValueTraits, NoopValueCounter> v(noop_vv, noop_dv, bc, vc);
+		tree.visit_depth_first(v);
+	}
 }
 
 //----------------------------------------------------------------