diff --git a/base/io_generator.cc b/base/io_generator.cc
index 0543d31..4d6ebb0 100644
--- a/base/io_generator.cc
+++ b/base/io_generator.cc
@@ -48,8 +48,10 @@ namespace {
 		base::sector_t next_offset() {
 			sector_t r = current_;
 			current_ += block_size_;
-			if (current_ > end_)
+			if (current_ > end_) {
 				current_ = begin_;
+				return begin_;
+			}
 			return r;
 		}
 
@@ -115,8 +117,7 @@ namespace {
 	class base_io_generator: public io_generator {
 	public:
 		base_io_generator(io_generator_options const &opts);
-		virtual bool has_next();
-		virtual void next(base::io &next_io);
+		virtual bool next(base::io &next_io);
 
 	private:
 		offset_generator::ptr
@@ -140,19 +141,17 @@ namespace {
 		  io_size_total_(opts.io_size_) {
 	}
 
-	bool base_io_generator::has_next() {
-		return io_size_finished_ < io_size_total_;
-	}
-
-	void base_io_generator::next(base::io &next_io) {
+	bool base_io_generator::next(base::io &next_io) {
 		if (io_size_finished_ >= io_size_total_)
-			throw std::runtime_error("");
+			return false;
 
 		next_io.op_ = op_gen_->next_op();
 		next_io.sector_ = offset_gen_->next_offset();
 		next_io.size_ = block_size_;
 
 		io_size_finished_ += block_size_;
+
+		return true;
 	}
 
 	offset_generator::ptr
diff --git a/base/io_generator.h b/base/io_generator.h
index 4a33550..16a78d9 100644
--- a/base/io_generator.h
+++ b/base/io_generator.h
@@ -42,8 +42,7 @@ namespace base {
 	public:
 		typedef std::shared_ptr<io_generator> ptr;
 
-		virtual bool has_next() = 0;
-		virtual void next(base::io &next_io) = 0;
+		virtual bool next(base::io &next_io) = 0;
 	};
 
 	io_generator::ptr
diff --git a/persistent-data/data-structures/btree-remove.tcc b/persistent-data/data-structures/btree-remove.tcc
index d222273..852d267 100644
--- a/persistent-data/data-structures/btree-remove.tcc
+++ b/persistent-data/data-structures/btree-remove.tcc
@@ -93,11 +93,12 @@ namespace persistent_data {
 	{
 		using namespace btree_detail;
 
-		unsigned i = 0;
-		bool r = false;
+		unsigned i = *index;
 
 		for (;;) {
-			r = spine.step(block);
+			bool inc = spine.step(block);
+			if (inc)
+				inc_children<ValueTraits>(spine, leaf_rc);
 
 			// patch up the parent to point to the new shadow
 			if (spine.has_parent()) {
@@ -115,9 +116,9 @@ namespace persistent_data {
 				return true;
 			}
 
-			r = rebalance_children<ValueTraits>(spine, key);
+			bool r = rebalance_children<ValueTraits>(spine, key);
 			if (!r)
-				break;
+				return false;
 
 			n = spine.get_node<block_traits>();
 			if (n.get_type() == btree_detail::LEAF) {
@@ -133,7 +134,7 @@ namespace persistent_data {
 			block = n.value_at(i);
 		}
 
-		return r;
+		return true;
 	}
 
 	template <unsigned Levels, typename _>
diff --git a/thin-provisioning/thin_generate_mappings.cc b/thin-provisioning/thin_generate_mappings.cc
index a366eba..1a124bb 100644
--- a/thin-provisioning/thin_generate_mappings.cc
+++ b/thin-provisioning/thin_generate_mappings.cc
@@ -93,10 +93,8 @@ namespace {
 		io_generator::ptr gen = create_io_generator(opts);
 
 		base::io io;
-		while (gen->has_next()) {
+		while (gen->next(io)) {
 			// TODO: support io.size_
-			gen->next(io);
-
 			switch (io.op_) {
 			case base::REQ_OP_READ:
 				process_read(td, pool, io.sector_);
@@ -131,10 +129,10 @@ thin_generate_mappings_cmd::usage(std::ostream &out) const
 	    << "  {-h|--help}\n"
 	    << "  {-o|--output} <output device or file>\n"
 	    << "  {--dev-id} <dev-id>\n"
-	    << "  {--offset} <offset>\n"
-            << "  {--io-size} <io_size>\n"
+	    << "  {--offset} <offset in sectors>\n"
+	    << "  {--io-size} <io-size in sectors>\n"
 	    << "  {--rw write|trim|randwrite|randtrim|randtw}\n"
-            << "  {--size} <size>\n"
+	    << "  {--size} <size in sectors>\n"
 	    << "  {-V|--version}" << endl;
 }
 
diff --git a/thin-provisioning/thin_pool.cc b/thin-provisioning/thin_pool.cc
index 3bdc9e7..3a34ba8 100644
--- a/thin-provisioning/thin_pool.cc
+++ b/thin-provisioning/thin_pool.cc
@@ -72,13 +72,16 @@ thin::insert(block_address thin_block, block_address data_block)
 {
 	uint64_t key[2] = {dev_, thin_block};
 
-	++details_.mapped_blocks_;
-	changed_ = true;
-
 	mapping_tree_detail::block_time bt;
 	bt.block_ = data_block;
 	bt.time_ = pool_.get_time();
-	return pool_.md_->mappings_->insert(key, bt);
+	bool inserted = pool_.md_->mappings_->insert(key, bt);
+
+	changed_ = true;
+	if (inserted)
+		++details_.mapped_blocks_;
+
+	return inserted;
 }
 
 void