git.sesse.net Git - remoteglot-book/blob - binloader.cpp

   1 //#define _GLIBCXX_PARALLEL
   2
   3 // Usage: ./binloader IN1 IN2 IN3 ... OUT NUM_BUCKETS
   4
   5 #include <stdio.h>
   6 #include <vector>
   7 #include <mtbl.h>
   8 #include <algorithm>
   9 #include <utility>
  10 #include <memory>
  11 #include <string>
  12 #include <string.h>
  13 #include "count.pb.h"
  14 #include "hash.h"
  15
  16 #define DUMMY_TIMESTAMP 32503680000
  17
  18 using namespace std;
  19
  20 enum Result { WHITE = 0, DRAW, BLACK };
  21 struct Element {
  22         string bpfen_and_move;
  23         Result result;
  24         int opening_num, white_elo, black_elo;
  25         time_t timestamp;
  26
  27         bool operator< (const Element& other) const {
  28                 return bpfen_and_move < other.bpfen_and_move;
  29         }
  30 };
  31
  32 int main(int argc, char **argv)
  33 {
  34         int num_buckets = atoi(argv[argc - 1]);
  35
  36         vector<vector<Element>> elems;
  37         elems.resize(num_buckets);
  38
  39         size_t num_elems = 0;
  40         for (int i = 1; i < argc - 2; ++i) {
  41                 FILE *fp = fopen(argv[i], "rb");
  42                 if (fp == NULL) {
  43                         perror(argv[i]);
  44                         exit(1);
  45                 }
  46                 for ( ;; ) {
  47                         int l = getc(fp);
  48                         if (l == -1) {
  49                                 break;
  50                         }
  51
  52                         string bpfen_and_move;
  53                         bpfen_and_move.resize(l);
  54                         if (fread(&bpfen_and_move[0], l, 1, fp) != 1) {
  55                                 perror("fread()");
  56                 //              exit(1);
  57                                 break;
  58                         }
  59
  60                         int r = getc(fp);
  61                         if (r == -1) {
  62                                 perror("getc()");
  63                                 //exit(1);
  64                                 break;
  65                         }
  66
  67                         int opening_num, white_elo, black_elo;
  68                         time_t timestamp;
  69                         if (fread(&white_elo, sizeof(white_elo), 1, fp) != 1) {
  70                                 perror("fread()");
  71                                 //exit(1);
  72                                 break;
  73                         }
  74                         if (fread(&black_elo, sizeof(black_elo), 1, fp) != 1) {
  75                                 perror("fread()");
  76                                 //exit(1);
  77                                 break;
  78                         }
  79                         if (fread(&opening_num, sizeof(opening_num), 1, fp) != 1) {
  80                                 perror("fread()");
  81                                 //exit(1);
  82                                 break;
  83                         }
  84                         if (fread(&timestamp, sizeof(timestamp), 1, fp) != 1) {
  85                                 perror("fread()");
  86                                 //exit(1);
  87                                 break;
  88                         }
  89
  90                         int bucket = hash_key_to_bucket(bpfen_and_move.data(), bpfen_and_move.size(), num_buckets);
  91                         elems[bucket].emplace_back(Element {move(bpfen_and_move), Result(r), opening_num, white_elo, black_elo, timestamp});
  92                         ++num_elems;
  93                 }
  94                 fclose(fp);
  95
  96                 printf("Read %ld elems\n", num_elems);
  97         }
  98
  99         printf("Sorting...\n");
 100         for (int i = 0; i < num_buckets; ++i) {
 101                 sort(elems[i].begin(), elems[i].end());
 102         }
 103
 104         printf("Writing SSTables...\n");
 105         string buf;  // Keep allocated.
 106         for (int i = 0; i < num_buckets; ++i) {
 107                 char filename[256];
 108                 snprintf(filename, sizeof(filename), "%s.part%04d", argv[argc - 2], i);
 109
 110                 mtbl_writer_options* wopt = mtbl_writer_options_init();
 111                 mtbl_writer_options_set_compression(wopt, MTBL_COMPRESSION_SNAPPY);
 112                 mtbl_writer* mtbl = mtbl_writer_init(filename, wopt);
 113                 Count c;
 114                 for (size_t j = 0; j < elems[i].size(); ++j) {
 115                         const Element &e = elems[i][j];
 116                         if (e.result == WHITE) {
 117                                 c.set_white(c.white() + 1);
 118                         } else if (e.result == DRAW) {
 119                                 c.set_draw(c.draw() + 1);
 120                         } else if (e.result == BLACK) {
 121                                 c.set_black(c.black() + 1);
 122                         }
 123                         if (e.white_elo >= 100 && e.black_elo >= 100) {
 124                                 c.set_sum_white_elo(c.sum_white_elo() + e.white_elo);
 125                                 c.set_sum_black_elo(c.sum_black_elo() + e.black_elo);
 126                                 c.set_num_elo(c.num_elo() + 1);
 127                         }
 128                         if (!c.has_first_timestamp() || e.timestamp < c.first_timestamp()) {
 129                                 if (e.timestamp != DUMMY_TIMESTAMP) {
 130                                         c.set_first_timestamp(e.timestamp);
 131                                 }
 132                                 c.set_opening_num(e.opening_num);
 133                         }
 134                         if (j == elems[i].size() - 1 || e.bpfen_and_move != elems[i][j + 1].bpfen_and_move) {
 135                                 c.SerializeToString(&buf);
 136                                 mtbl_writer_add(mtbl,
 137                                         (const uint8_t *)e.bpfen_and_move.data(), e.bpfen_and_move.size(),
 138                                         (const uint8_t *)buf.data(), buf.size());
 139                                 c = Count();
 140                         }
 141                 }
 142                 mtbl_writer_destroy(&mtbl);
 143         }
 144 }