Add some utilities to be able to build books in parts, so we do not need as much...
[remoteglot] / 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* mtbl = mtbl_writer_init(argv[argc - 1], NULL);
87         Count c;
88         for (int i = 0; i < elems.size(); ++i) {
89                 if (elems[i].result == WHITE) {
90                         ++c.white;
91                 } else if (elems[i].result == DRAW) {
92                         ++c.draw;
93                 } else if (elems[i].result == BLACK) {
94                         ++c.black;
95                 }
96                 c.opening_num = elems[i].opening_num;
97                 if (elems[i].white_elo >= 100 && elems[i].black_elo >= 100) {
98                         c.sum_white_elo += elems[i].white_elo;
99                         c.sum_black_elo += elems[i].black_elo;
100                         ++c.num_elo;
101                 }
102                 if (i == elems.size() - 1 || elems[i].bpfen_and_move != elems[i + 1].bpfen_and_move) {
103                         mtbl_writer_add(mtbl,
104                                 (const uint8_t *)elems[i].bpfen_and_move.data(), elems[i].bpfen_and_move.size(),
105                                 (const uint8_t *)&c, sizeof(c));
106                         c = Count();
107                 }
108         }
109         mtbl_writer_destroy(&mtbl);
110 }