From b22495997a8bf427741225ef86edf267ee746d7c Mon Sep 17 00:00:00 2001
From: Ming-Hung Tsai <mingnus@gmail.com>
Date: Sat, 27 Feb 2016 15:22:00 +0800
Subject: [PATCH] Allow counting_visitor to work with damaged btrees

---
 .../data-structures/btree_counter.h           | 28 +++++++++++++------
 .../data-structures/btree_damage_visitor.h    |  5 ++--
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/persistent-data/data-structures/btree_counter.h b/persistent-data/data-structures/btree_counter.h
index ed7b845..c215626 100644
--- a/persistent-data/data-structures/btree_counter.h
+++ b/persistent-data/data-structures/btree_counter.h
@@ -2,35 +2,44 @@
 #define PERSISTENT_DATA_DATA_STRUCTURES_BTREE_COUNTER_H
 
 #include "persistent-data/data-structures/btree.h"
+#include "persistent-data/data-structures/btree_base_visitor.h"
+#include "persistent-data/data-structures/btree_damage_visitor.h"
 #include "persistent-data/block_counter.h"
 
 //----------------------------------------------------------------
 
 namespace persistent_data {
 	namespace btree_count_detail {
-		template <unsigned Levels, typename ValueTraits, typename ValueCounter>
-		class counting_visitor : public btree<Levels, ValueTraits>::visitor {
+		template <typename ValueVisitor, typename DamageVisitor, unsigned Levels, typename ValueTraits, typename ValueCounter>
+		class counting_visitor : public btree_damage_visitor<ValueVisitor, DamageVisitor, Levels, ValueTraits> {
+			typedef btree_damage_visitor<ValueVisitor, DamageVisitor, Levels, ValueTraits> BtreeDamageVisitor;
 		public:
 			typedef btree<Levels, ValueTraits> tree;
 
-			counting_visitor(block_counter &bc, ValueCounter &vc)
-				: bc_(bc),
+			counting_visitor(ValueVisitor &value_visitor,
+					 DamageVisitor &damage_visitor,
+					 block_counter &bc,
+					 ValueCounter &vc)
+				: BtreeDamageVisitor(value_visitor, damage_visitor, false),
+				  bc_(bc),
 				  vc_(vc) {
 			}
 
 			virtual bool visit_internal(node_location const &l,
 						    typename tree::internal_node const &n) {
-				return visit_node(n);
+				return BtreeDamageVisitor::visit_internal(l, n) ?
+					visit_node(n) : false;
 			}
 
 			virtual bool visit_internal_leaf(node_location const &l,
 							 typename tree::internal_node const &n) {
-				return visit_node(n);
+				return BtreeDamageVisitor::visit_internal_leaf(l, n) ?
+					visit_node(n) : false;
 			}
 
 			virtual bool visit_leaf(node_location const &l,
 						typename tree::leaf_node const &n) {
-				if (visit_node(n)) {
+				if (BtreeDamageVisitor::visit_leaf(l, n) && visit_node(n)) {
 					unsigned nr = n.get_nr_entries();
 
 					for (unsigned i = 0; i < nr; i++) {
@@ -85,7 +94,10 @@ namespace persistent_data {
 	// is not corrupt.
 	template <unsigned Levels, typename ValueTraits, typename ValueCounter>
 	void count_btree_blocks(btree<Levels, ValueTraits> const &tree, block_counter &bc, ValueCounter &vc) {
-		btree_count_detail::counting_visitor<Levels, ValueTraits, ValueCounter> v(bc, vc);
+		typedef noop_value_visitor<typename ValueTraits::value_type> NoopValueVisitor;
+		NoopValueVisitor noop_vv;
+		noop_damage_visitor noop_dv;
+		btree_count_detail::counting_visitor<NoopValueVisitor, noop_damage_visitor, Levels, ValueTraits, ValueCounter> v(noop_vv, noop_dv, bc, vc);
 		tree.visit_depth_first(v);
 	}
 }
diff --git a/persistent-data/data-structures/btree_damage_visitor.h b/persistent-data/data-structures/btree_damage_visitor.h
index 5a96d5d..8be623b 100644
--- a/persistent-data/data-structures/btree_damage_visitor.h
+++ b/persistent-data/data-structures/btree_damage_visitor.h
@@ -158,8 +158,9 @@ namespace persistent_data {
 			typedef boost::optional<run64> maybe_run64;
 
 			btree_damage_visitor(ValueVisitor &value_visitor,
-					     DamageVisitor &damage_visitor)
-				: avoid_repeated_visits_(true),
+					     DamageVisitor &damage_visitor,
+					     bool avoid_repeated_visits = true)
+				: avoid_repeated_visits_(avoid_repeated_visits),
 				  value_visitor_(value_visitor),
 				  damage_visitor_(damage_visitor) {
 			}