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

   1 //#define _GLIBCXX_PARALLEL
   2
   3 // Usage: ./binloader IN1 IN2 IN3 ... OUT
   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.h"
  14
  15 using namespace std;
  16
  17 enum Result { WHITE = 0, DRAW, BLACK };
  18 struct Element {
  19         string bpfen_and_move;
  20         Result result;
  21         int opening_num, white_elo, black_elo;
  22
  23         bool operator< (const Element& other) const {
  24                 return bpfen_and_move < other.bpfen_and_move;
  25         }
  26 };
  27
  28 int main(int argc, char **argv)
  29 {
  30         vector<Element> elems;
  31
  32         for (int i = 1; i < argc - 1; ++i) {
  33                 FILE *fp = fopen(argv[i], "rb");
  34                 if (fp == NULL) {
  35                         perror(argv[i]);
  36                         exit(1);
  37                 }
  38                 for ( ;; ) {
  39                         int l = getc(fp);
  40                         if (l == -1) {
  41                                 break;
  42                         }
  43
  44                         string bpfen_and_move;
  45                         bpfen_and_move.resize(l);
  46                         if (fread(&bpfen_and_move[0], l, 1, fp) != 1) {
  47                                 perror("fread()");
  48                 //              exit(1);
  49                                 break;
  50                         }
  51
  52                         int r = getc(fp);
  53                         if (r == -1) {
  54                                 perror("getc()");
  55                                 //exit(1);
  56                                 break;
  57                         }
  58
  59                         int opening_num, white_elo, black_elo;
  60                         if (fread(&white_elo, sizeof(white_elo), 1, fp) != 1) {
  61                                 perror("fread()");
  62                                 //exit(1);
  63                                 break;
  64                         }
  65                         if (fread(&black_elo, sizeof(black_elo), 1, fp) != 1) {
  66                                 perror("fread()");
  67                                 //exit(1);
  68                                 break;
  69                         }
  70                         if (fread(&opening_num, sizeof(opening_num), 1, fp) != 1) {
  71                                 perror("fread()");
  72                                 //exit(1);
  73                                 break;
  74                         }
  75                         elems.emplace_back(Element {move(bpfen_and_move), Result(r), opening_num, white_elo, black_elo});
  76                 }
  77                 fclose(fp);
  78
  79                 printf("Read %ld elems\n", elems.size());
  80         }
  81
  82         printf("Sorting...\n");
  83         sort(elems.begin(), elems.end());
  84
  85         printf("Writing SSTable...\n");
  86         mtbl_writer_options* wopt = mtbl_writer_options_init();
  87         mtbl_writer_options_set_compression(wopt, MTBL_COMPRESSION_SNAPPY);
  88         mtbl_writer* mtbl = mtbl_writer_init(argv[argc - 1], wopt);
  89         Count c;
  90         for (int i = 0; i < elems.size(); ++i) {
  91                 if (elems[i].result == WHITE) {
  92                         ++c.white;
  93                 } else if (elems[i].result == DRAW) {
  94                         ++c.draw;
  95                 } else if (elems[i].result == BLACK) {
  96                         ++c.black;
  97                 }
  98                 c.opening_num = elems[i].opening_num;
  99                 if (elems[i].white_elo >= 100 && elems[i].black_elo >= 100) {
 100                         c.sum_white_elo += elems[i].white_elo;
 101                         c.sum_black_elo += elems[i].black_elo;
 102                         ++c.num_elo;
 103                 }
 104                 if (i == elems.size() - 1 || elems[i].bpfen_and_move != elems[i + 1].bpfen_and_move) {
 105                         mtbl_writer_add(mtbl,
 106                                 (const uint8_t *)elems[i].bpfen_and_move.data(), elems[i].bpfen_and_move.size(),
 107                                 (const uint8_t *)&c, sizeof(c));
 108                         c = Count();
 109                 }
 110         }
 111         mtbl_writer_destroy(&mtbl);
 112 }