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

   1 #include <stdio.h>
   2 #include <mtbl.h>
   3 #include <memory>
   4 #include <string>
   5 #include <unordered_set>
   6 #include <string.h>
   7 #include <assert.h>
   8 #include "count.pb.h"
   9
  10 using namespace std;
  11
  12 void merge_count(void* userdata,
  13                  const uint8_t *key, size_t len_key,
  14                  const uint8_t *val0, size_t len_val0,
  15                  const uint8_t *val1, size_t len_val1,
  16                  uint8_t **merged_val, size_t *len_merged_val)
  17 {
  18         Count c0, c1;
  19         c0.ParseFromArray(val0, len_val0);
  20         c1.ParseFromArray(val1, len_val1);
  21
  22         Count c;
  23
  24         c.set_white(c0.white() + c1.white());
  25         c.set_draw(c0.draw() + c1.draw());
  26         c.set_black(c0.black() + c1.black());
  27         c.set_sum_white_elo(c0.sum_white_elo() + c1.sum_white_elo());
  28         c.set_sum_black_elo(c0.sum_black_elo() + c1.sum_black_elo());
  29         c.set_num_elo(c0.num_elo() + c1.num_elo());
  30         if (c0.first_timestamp() <= c1.first_timestamp()) {
  31                 c.set_opening_num(c0.opening_num());
  32                 if (c0.has_first_timestamp()) {
  33                         c.set_first_timestamp(c0.first_timestamp());
  34                 }
  35                 c.set_pgn_file_num(c0.pgn_file_num());
  36                 c.set_pgn_start_position(c0.pgn_start_position());
  37         } else {
  38                 c.set_opening_num(c1.opening_num());
  39                 if (c1.has_first_timestamp()) {
  40                         c.set_first_timestamp(c1.first_timestamp());
  41                 }
  42                 c.set_pgn_file_num(c1.pgn_file_num());
  43                 c.set_pgn_start_position(c1.pgn_start_position());
  44         }
  45
  46         // Merge the moves, with deduplication.
  47         unordered_set<string> moves;
  48         for (int i = 0; i < c0.move_size(); ++i) {
  49                 moves.insert(c0.move(i));
  50                 c.add_move(c0.move(i));
  51         }
  52         for (int i = 0; i < c1.move_size(); ++i) {
  53                 if (!moves.count(c1.move(i))) {
  54                         c.add_move(c1.move(i));
  55                 }
  56         }
  57
  58         static string buf;  // Keep allocated.
  59         c.SerializeToString(&buf);
  60
  61         *merged_val = reinterpret_cast<uint8_t *>(malloc(buf.size()));
  62         *len_merged_val = buf.size();
  63         memcpy(*merged_val, buf.data(), buf.size());
  64 }
  65
  66 int main(int argc, char **argv)
  67 {
  68         mtbl_merger_options* mopt = mtbl_merger_options_init();
  69         mtbl_merger_options_set_merge_func(mopt, merge_count, NULL);
  70         mtbl_merger* merger = mtbl_merger_init(mopt);
  71
  72         for (int i = 1; i < argc - 1; ++i) {
  73                 mtbl_reader* mtbl = mtbl_reader_init(argv[i], NULL);
  74                 mtbl_merger_add_source(merger, mtbl_reader_source(mtbl));
  75         }
  76
  77         mtbl_writer_options* wopt = mtbl_writer_options_init();
  78         mtbl_writer_options_set_block_size(wopt, 65536);
  79         mtbl_writer* writer = mtbl_writer_init(argv[argc - 1], wopt);
  80         mtbl_source_write(mtbl_merger_source(merger), writer);
  81         mtbl_writer_destroy(&writer);
  82 }