Tune compression options a bit for intermediates (binloader) vs. final output (binmer...
[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_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 }