From fce7e99d9a818cf00be1ebfd1de7a6852c51307e Mon Sep 17 00:00:00 2001
From: Joe Thornber <ejt@redhat.com>
Date: Thu, 25 Aug 2011 15:51:40 +0100
Subject: [PATCH] improve reference count tracking

---
 btree_validator.h |  8 ++++++--
 metadata.cc       | 18 ++++++++----------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/btree_validator.h b/btree_validator.h
index 316e0b5..4e6fde6 100644
--- a/btree_validator.h
+++ b/btree_validator.h
@@ -76,20 +76,24 @@ namespace persistent_data {
 
 		void visit_internal(unsigned level, bool is_root,
 				    btree_detail::node_ref<uint64_traits, BlockSize> const &n) {
-			counter_.inc(n.get_location());
 			check_duplicate_block(n.get_location());
 			check_block_nr(n);
 			check_max_entries(n);
 			check_nr_entries(n, is_root);
+
+			for (unsigned i = 0; i < n.get_nr_entries(); i++)
+				counter_.inc(n.value_at(i));
 		}
 
 		void visit_internal_leaf(unsigned level, bool is_root,
 					 btree_detail::node_ref<uint64_traits, BlockSize> const &n) {
-			counter_.inc(n.get_location());
 			check_duplicate_block(n.get_location());
 			check_block_nr(n);
 			check_max_entries(n);
 			check_nr_entries(n, is_root);
+
+			for (unsigned i = 0; i < n.get_nr_entries(); i++)
+				counter_.inc(n.value_at(i));
 		}
 
 		void visit_leaf(unsigned level, bool is_root,
diff --git a/metadata.cc b/metadata.cc
index ec7fb43..c27502a 100644
--- a/metadata.cc
+++ b/metadata.cc
@@ -328,22 +328,20 @@ metadata::device_exists(thin_dev_t dev) const
 }
 
 namespace {
-	// FIXME: this doesn't check for non-zero counts in the sm that are
-	// actually zero.
 	optional<error_set::ptr>
-	check_ref_counts(string const &desc, block_counter const &actual,
+	check_ref_counts(string const &desc, block_counter const &counts,
 			 space_map::ptr sm) {
 		error_set::ptr errors(new error_set(desc));
 
 		bool bad = false;
-		block_counter::count_map const &counts = actual.get_counts();
-		block_counter::count_map::const_iterator it, end = counts.end();
-		for (it = counts.begin(); it != end; ++it) {
-			uint32_t ref_count = sm->get_count(it->first);
-			if (ref_count != it->second) {
+		for (block_address b = 0; b < sm->get_nr_blocks(); b++) {
+			uint32_t actual = sm->get_count(b);
+			uint32_t expected = counts.get_count(b);
+
+			if (actual != expected) {
 				ostringstream out;
-				out << it->first << ": was " << ref_count
-				    << ", expected " << it->second;
+				out << b << ": was " << actual
+				    << ", expected " << expected;
 				errors->add_child(out.str());
 				bad = true;
 			}