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

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