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 "merge_count.h"
   9 #include "count.pb.h"
  10
  11 using namespace std;
  12
  13 void merge_serialized_count(void* userdata,
  14                             const uint8_t *key, size_t len_key,
  15                             const uint8_t *val0, size_t len_val0,
  16                             const uint8_t *val1, size_t len_val1,
  17                             uint8_t **merged_val, size_t *len_merged_val)
  18 {
  19         Count c0, c1;
  20         c0.ParseFromArray(val0, len_val0);
  21         c1.ParseFromArray(val1, len_val1);
  22
  23         Count c = merge_count(c0, c1);
  24
  25         static string buf;  // Keep allocated.
  26         c.SerializeToString(&buf);
  27
  28         *merged_val = reinterpret_cast<uint8_t *>(malloc(buf.size()));
  29         *len_merged_val = buf.size();
  30         memcpy(*merged_val, buf.data(), buf.size());
  31 }
  32
  33 int main(int argc, char **argv)
  34 {
  35         mtbl_merger_options* mopt = mtbl_merger_options_init();
  36         mtbl_merger_options_set_merge_func(mopt, merge_serialized_count, NULL);
  37         mtbl_merger* merger = mtbl_merger_init(mopt);
  38
  39         for (int i = 1; i < argc - 1; ++i) {
  40                 mtbl_reader* mtbl = mtbl_reader_init(argv[i], NULL);
  41                 mtbl_merger_add_source(merger, mtbl_reader_source(mtbl));
  42         }
  43
  44         mtbl_writer_options* wopt = mtbl_writer_options_init();
  45         mtbl_writer_options_set_block_size(wopt, 65536);
  46         mtbl_writer* writer = mtbl_writer_init(argv[argc - 1], wopt);
  47         mtbl_source_write(mtbl_merger_source(merger), writer);
  48         mtbl_writer_destroy(&writer);
  49 }