From 240e782a07ba03cd9da13fb8a40a307847d8fc42 Mon Sep 17 00:00:00 2001
From: Joe Thornber <ejt@redhat.com>
Date: Thu, 25 Aug 2011 16:27:58 +0100
Subject: [PATCH] include ref_count tree in ref counts

---
 btree.h          |  4 ++--
 btree.tcc        |  4 ++--
 space_map_disk.h | 26 +++++++++++++++++++++++---
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/btree.h b/btree.h
index 07e4cb9..0235471 100644
--- a/btree.h
+++ b/btree.h
@@ -331,7 +331,7 @@ namespace persistent_data {
 		};
 
 		// Walks the tree in depth first order
-		void visit(typename visitor::ptr visitor);
+		void visit(typename visitor::ptr visitor) const;
 
 	private:
 		template <typename ValueTraits2>
@@ -357,7 +357,7 @@ namespace persistent_data {
 
 		void walk_tree(typename visitor::ptr visitor,
 			       unsigned level, bool is_root,
-			       block_address b);
+			       block_address b) const;
 
 		typename persistent_data::transaction_manager<BlockSize>::ptr tm_;
 		bool destroy_;
diff --git a/btree.tcc b/btree.tcc
index ac8ee80..f821153 100644
--- a/btree.tcc
+++ b/btree.tcc
@@ -613,7 +613,7 @@ insert_location(btree_detail::shadow_spine<BlockSize> &spine,
 
 template <unsigned Levels, typename ValueTraits, uint32_t BlockSize>
 void
-btree<Levels, ValueTraits, BlockSize>::visit(typename visitor::ptr visitor)
+btree<Levels, ValueTraits, BlockSize>::visit(typename visitor::ptr visitor) const
 {
 	walk_tree(visitor, 0, true, root_);
 }
@@ -623,7 +623,7 @@ void
 btree<Levels, ValueTraits, BlockSize>::
 walk_tree(typename visitor::ptr visitor,
 	  unsigned level, bool is_root,
-	  block_address b)
+	  block_address b) const
 {
 	using namespace btree_detail;
 
diff --git a/space_map_disk.h b/space_map_disk.h
index c314e80..ff78020 100644
--- a/space_map_disk.h
+++ b/space_map_disk.h
@@ -112,6 +112,16 @@ namespace persistent_data {
 			}
 		};
 
+		template <uint32_t BlockSize>
+		class ref_count_validator : public btree_validator<1, ref_count_traits, BlockSize> {
+		public:
+			typedef boost::shared_ptr<ref_count_validator> ptr;
+
+			ref_count_validator(block_counter &counter)
+				: btree_validator<1, ref_count_traits, BlockSize>(counter) {
+			}
+		};
+
 		template <uint32_t BlockSize>
 		class sm_disk_base : public persistent_space_map {
 		public:
@@ -233,6 +243,12 @@ namespace persistent_data {
 				nr_blocks_ = nr_blocks;
 			}
 
+			virtual void check(block_counter &counter) const {
+				typename ref_count_validator<BlockSize>::ptr v(new ref_count_validator<BlockSize>(counter));
+				counter.inc(ref_counts_.get_root());
+				ref_counts_.visit(v);
+			}
+
 		protected:
 			typename transaction_manager<BlockSize>::ptr get_tm() const {
 				return tm_;
@@ -342,8 +358,11 @@ namespace persistent_data {
 				::memcpy(dest, &d, sizeof(d));
 			}
 
-			void check(block_counter &counter) {
+			void check(block_counter &counter) const {
+				sm_disk_base<BlockSize>::check(counter);
+
 				typename bitmap_tree_validator<BlockSize>::ptr v(new bitmap_tree_validator<BlockSize>(counter));
+				counter.inc(bitmaps_.get_root());
 				bitmaps_.visit(v);
 			}
 
@@ -407,9 +426,10 @@ namespace persistent_data {
 				::memcpy(dest, &d, sizeof(d));
 			}
 
-			void check(block_counter &counter) {
-				counter.inc(bitmap_root_);
+			void check(block_counter &counter) const {
+				sm_disk_base<BlockSize>::check(counter);
 
+				counter.inc(bitmap_root_);
 				for (unsigned i = 0; i < entries_.size(); i++)
 					counter.inc(entries_[i].blocknr_);
 			}