fix(exp): Fix expt_0520 malloc failure

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

1 files changed, 38 insertions, 7 deletions

diff --git a/read_helper.h b/read_helper.h
index 311ebbb..292eaac 100644
--- a/read_helper.h
+++ b/read_helper.h
@@ -15,12 +15,18 @@
 #endif
 
 class ReadHelper {
-  unsigned long long max_value = 0, min_value = 0x7fffffffff;
+  typedef unsigned long long key_type;
+  key_type max_value = 0, min_value = 0x7fffffffff;
 
   char const *filename;
+  static constexpr size_t REVERSED_BLOCK = 256;
 
 public:
-  ReadHelper(char const *filename) : filename(filename) {}
+  size_t sample_length, needed_read_length;
+  ReadHelper(char const *filename, size_t sample_length,
+             size_t insertion_length)
+      : filename(filename), sample_length(sample_length),
+        needed_read_length(insertion_length + sample_length) {}
   unsigned long long maxValue() const { return max_value; }
   unsigned long long minValue() const { return min_value; }
 
@@ -35,6 +41,11 @@ public:
     population_vector.push_back(value);
   }
 
+  void finish_read() {
+    this->population_vector.push_back(this->max_value);
+    this->population_vector.push_back(this->min_value);
+  }
+
   static unsigned long randomRow(unsigned long max_value_) {
     std::random_device randomDevice;
     std::mt19937 mt19937(randomDevice());
@@ -42,20 +53,21 @@ public:
     return dst(mt19937);
   }
 
-  void readFile(long sample_length, unsigned long &total_row,
-                unsigned long max_number) {
+  void readFile(unsigned long &total_row) {
+    this->population_vector.clear();
     auto read_number = 0UL;
     FILE *file = fopen(filename, "r");
     assert(file);
     P_ERR("Reading sample");
     for (long long i;
-         read_number < max_number && fscanf(file, "%lld ", &i) != EOF;
+         read_number < needed_read_length && fscanf(file, "%lld ", &i) != EOF;
          store_into_vector(i))
       read_number++;
 
     if (total_row > 0) {
       P_ERR("\rReading skip");
-      read_number = randomRow(total_row - sample_length - max_number - 256);
+      read_number = randomRow(total_row - sample_length - needed_read_length -
+                              REVERSED_BLOCK);
       total_row = read_number;
       // fprintf(stderr, "Skip %lu\n", read_number);
       for (long long i; read_number > 0 && fscanf(file, "%lld ", &i) != EOF;)
@@ -65,12 +77,31 @@ public:
     P_ERR("\rReading population");
     read_number = 0;
 
-    auto remain = sample_length + 256;
+    auto remain = sample_length + REVERSED_BLOCK;
     for (long long i; read_number < remain && fscanf(file, "%lld ", &i) != EOF;
          store_into_vector(i))
       read_number++;
     fclose(file);
     P_ERR("\r");
+    finish_read();
+  }
+
+  void split_into(std::vector<key_type> &sample, std::vector<key_type> &p) {
+    sample.reserve(sample_length);
+    p.reserve(needed_read_length - sample_length);
+    memcpy(sample.data(), population_vector.data(),
+           sizeof(key_type) * sample_length);
+    memcpy(p.data(), population_vector.data() + sample_length,
+           sizeof(key_type) * (needed_read_length - sample_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];
+    memcpy(sample, population_vector.data(), sizeof(key_type) * sample_length);
+    memcpy(p, population_vector.data() + sample_length,
+           sizeof(key_type) * (needed_read_length - sample_length));
   }
 };
+
 #endif
\ No newline at end of file