feat: Optimize read_helper library

Signed-off-by: KunoiSayami <[email protected]>

1 files changed, 23 insertions, 13 deletions

diff --git a/read_helper.h b/read_helper.h
index d19834a..f8396f2 100644
--- a/read_helper.h
+++ b/read_helper.h
@@ -2,6 +2,7 @@
 #ifndef LOCKFREE_READ_HELPER_H
 #define LOCKFREE_READ_HELPER_H
 
+#include <cassert>
 #include <cstdio>
 #include <random>
 #include <vector>
@@ -34,12 +35,22 @@ class ReadHelper {
   char const *filename;
   static constexpr size_t REVERSED_BLOCK = 256;
 
+  auto genRandomRow(size_t total_row) const {
+    if (total_row != 0) {
+      assert(total_row > (population_length + sample_length + REVERSED_BLOCK));
+      return randomRow(total_row - population_length - sample_length -
+                       REVERSED_BLOCK);
+    }
+    return 0UL;
+  }
+
 public:
-  const size_t sample_length, needed_read_length;
+  const size_t sample_length, population_length, random_number;
   ReadHelper(char const *filename, size_t sample_length,
-             size_t insertion_length)
+             size_t population_length, size_t total_row = 0)
       : filename(filename), sample_length(sample_length),
-        needed_read_length(insertion_length + sample_length) {
+        population_length(population_length),
+        random_number(genRandomRow(total_row)) {
     // printf("%zu %zu\n", sample_length, needed_read_length);
   }
   key_type maxValue() const { return max_value; }
@@ -54,10 +65,11 @@ public:
     return dst(mt19937);
   }
 
-  void readFile(unsigned long *total_row) {
+  void readFile() {
     this->population_vector.clear();
     auto read_number = 0UL;
     FILE *file = fopen(filename, "r");
+
     if (file == nullptr) {
       fprintf(stderr, "Unable to open file %s\n", filename);
       exit(1);
@@ -69,12 +81,10 @@ public:
          store_into_vector(i))
       read_number++;
 
-    if (total_row != nullptr) {
+    if (random_number > 0) {
       P_ERR("\rReading skip");
-      read_number = randomRow(*total_row - sample_length - needed_read_length -
-                              REVERSED_BLOCK);
-      *total_row = read_number;
       // fprintf(stderr, "Skip %lu\n", read_number);
+      read_number = random_number;
       for (key_type i; read_number > 0 && fscanf(file, "%lld ", &i) != EOF;)
         read_number--;
     }
@@ -82,7 +92,7 @@ public:
     P_ERR("\rReading population");
     read_number = 0;
 
-    auto remain = needed_read_length - sample_length + REVERSED_BLOCK;
+    auto remain = population_length + REVERSED_BLOCK;
     for (key_type i; read_number < remain && fscanf(file, "%lld ", &i) != EOF;
          store_into_vector(i))
       read_number++;
@@ -92,25 +102,25 @@ public:
 
   void split_into(std::vector<key_type> &sample, std::vector<key_type> &p) {
     sample.resize(sample_length - 2);
-    p.resize(needed_read_length - sample_length);
+    p.resize(population_length);
     // printf("%zu\n", needed_read_length - sample_length);
     memcpy(sample.data(), population_vector.data(),
            sizeof(key_type) * (sample_length - 2));
     sample.push_back(this->max_value);
     sample.push_back(this->min_value);
     memcpy(p.data(), population_vector.data() + sample_length,
-           sizeof(key_type) * (needed_read_length - sample_length));
+           sizeof(key_type) * population_length);
   }
 
   void split_into(key_type *&sample, key_type *&p) {
     sample = new key_type[sample_length];
-    p = new key_type[needed_read_length - sample_length];
+    p = new key_type[population_length];
     memcpy(sample, population_vector.data(),
            sizeof(key_type) * (sample_length - 2));
     sample[sample_length - 2] = this->max_value;
     sample[sample_length - 1] = this->min_value;
     memcpy(p, population_vector.data() + sample_length,
-           sizeof(key_type) * (needed_read_length - sample_length));
+           sizeof(key_type) * (population_length));
   }
 
   size_t size() const { return this->population_vector.size(); }