From 0c3c5f6c1ec2a614a9e6c7213ece994d5597572a Mon Sep 17 00:00:00 2001
From: Ming-Hung Tsai <mtsai@redhat.com>
Date: Tue, 28 Jul 2020 00:32:30 +0800
Subject: [PATCH] [sequence_generator] Fix parameters and span boundary
 checking

---
 base/sequence_generator.cc | 90 +++++++++++++++++++++++++-------------
 1 file changed, 60 insertions(+), 30 deletions(-)

diff --git a/base/sequence_generator.cc b/base/sequence_generator.cc
index a811b48..cd0e3f9 100644
--- a/base/sequence_generator.cc
+++ b/base/sequence_generator.cc
@@ -14,12 +14,12 @@ namespace {
 					   uint64_t size,
 					   uint64_t step)
 			: begin_(begin),
-			  end_(begin + (size / step) * step),
 			  step_(step),
 			  current_(begin),
-			  rounds_(0) {
-			if (size < step)
-				throw std::runtime_error("size must be greater than the step");
+			  rounds_(0)
+		{
+			verify_parameters(size, step);
+			end_ = begin + (size / step) * step;
 		}
 
 		forward_sequence_generator()
@@ -38,6 +38,8 @@ namespace {
 		}
 
 		void reset(uint64_t begin, uint64_t size, uint64_t step = 1) {
+			verify_parameters(size, step);
+
 			begin_ = begin;
 			end_ = begin + (size / step) * step;
 			step_ = step;
@@ -45,6 +47,14 @@ namespace {
 			rounds_ = 0;
 		}
 
+		static void verify_parameters(uint64_t size, uint64_t step) {
+			if (!size || !step)
+				throw std::runtime_error("size or step must be non-zero");
+
+			if (size < step)
+				throw std::runtime_error("size must be greater than the step");
+		}
+
 		uint64_t get_rounds() {
 			return rounds_;
 		}
@@ -68,11 +78,18 @@ namespace {
 					  uint64_t step,
 					  unsigned seq_nr = 1)
 			: begin_(begin),
-			  nr_steps_(size / step),
 			  step_(step),
 			  max_forward_steps_(seq_nr),
 			  nr_generated_(0)
 		{
+			if (!size || !step || !seq_nr)
+				throw std::runtime_error("size, step, or forward steps must be non-zero");
+
+			if (size < step)
+				throw std::runtime_error("size must be greater than the step");
+
+			nr_steps_ = size / step;
+
 			if (!max_forward_steps_ || max_forward_steps_ > nr_steps_)
 				throw std::runtime_error("invalid number of forward steps");
 
@@ -84,7 +101,44 @@ namespace {
 			// FIXME: eliminate if-else
 			uint64_t step_idx = (max_forward_steps_ > 1) ?
 				next_forward_step() : next_random_step();
+
+			return begin_ + step_idx * step_;
+		}
+
+	private:
+		void reset_forward_generator() {
+			uint64_t begin = peek_random_step();
+
+			unsigned seq_nr = (std::rand() % max_forward_steps_) + 1;
+			base::run_set<uint64_t>::const_iterator it = rand_map_.upper_bound(begin);
+			if (it != rand_map_.end())
+				seq_nr = std::min<uint64_t>(seq_nr, *it->begin_ - begin);
+			else
+				seq_nr = std::min<uint64_t>(seq_nr, nr_steps_ - begin);
+
+			forward_gen_.reset(begin, seq_nr);
+		}
+
+		uint64_t next_forward_step() {
+			uint64_t step_idx = forward_gen_.next();
+			consume_random_map(step_idx);
+
+			if (forward_gen_.get_rounds())
+				reset_forward_generator();
+
+			return step_idx;
+		}
+
+		uint64_t next_random_step() {
+			uint64_t step_idx = peek_random_step();
+			consume_random_map(step_idx);
+
+			return step_idx;
+		}
+
+		void consume_random_map(uint64_t step_idx) {
 			rand_map_.add(step_idx);
+
 			++nr_generated_;
 
 			// wrap-around
@@ -92,33 +146,9 @@ namespace {
 				rand_map_.clear();
 				nr_generated_ = 0;
 			}
-
-			return begin_ + step_idx * step_;
 		}
 
-	private:
-		void reset_forward_generator() {
-			uint64_t begin = next_random_step();
-			unsigned seq_nr = (std::rand() % max_forward_steps_) + 1;
-			forward_gen_.reset(begin, seq_nr);
-		}
-
-		uint64_t next_forward_step() {
-			uint64_t step_idx;
-
-			bool found = true;
-			while (found) {
-				step_idx = forward_gen_.next();
-				found = rand_map_.member(step_idx);
-
-				if (found || forward_gen_.get_rounds())
-					reset_forward_generator();
-			}
-
-			return step_idx;
-		}
-
-		uint64_t next_random_step() const {
+		uint64_t peek_random_step() const {
 			uint64_t step_idx;
 
 			bool found = true;