Home | History | Annotate | Download | only in tools 
      1 #include <cstdlib>
      2 #include <iostream>
      3 #include <limits>
      4 #include <string>
      5 #include <vector>
      6 
      7 #include <marisa.h>
      8 
      9 #include "./cmdopt.h"
     10 
     11 namespace {
     12 
     13 std::size_t max_num_results = 10;
     14 bool mmap_flag = true;
     15 bool depth_first_flag = true;
     16 
     17 void print_help(const char *cmd) {
     18   std::cerr << "Usage: " << cmd << " [OPTION]... DIC\n\n"
     19       "Options:\n"
     20       "  -n, --max-num-results=[N]  limits the number of outputs to N"
     21       " (default: 10)\n"
     22       "                             0: no limit\n"
     23       "  -d, --depth-first      predict keys in depth first order(default)\n"
     24       "  -b, --breadth-first    predict keys in breadth first order\n"
     25       "  -m, --mmap-dictionary  use memory-mapped I/O to load a dictionary"
     26       " (default)\n"
     27       "  -r, --read-dictionary  read an entire dictionary into memory\n"
     28       "  -h, --help             print this help\n"
     29       << std::endl;
     30 }
     31 
     32 int predict(const char * const *args, std::size_t num_args) {
     33   if (num_args == 0) {
     34     std::cerr << "error: a dictionary is not specified" << std::endl;
     35     return 10;
     36   } else if (num_args > 1) {
     37     std::cerr << "error: more than one dictionaries are specified"
     38         << std::endl;
     39     return 11;
     40   }
     41 
     42   marisa::Trie trie;
     43   marisa::Mapper mapper;
     44   if (mmap_flag) {
     45     try {
     46       trie.mmap(&mapper, args[0]);
     47     } catch (const marisa::Exception &ex) {
     48       std::cerr << ex.filename() << ':' << ex.line() << ": " << ex.what()
     49           << ": failed to mmap a dictionary file: " << args[0] << std::endl;
     50       return 20;
     51     }
     52   } else {
     53     try {
     54       trie.load(args[0]);
     55     } catch (const marisa::Exception &ex) {
     56       std::cerr << ex.filename() << ':' << ex.line() << ": " << ex.what()
     57           << ": failed to load a dictionary file: " << args[0] << std::endl;
     58       return 21;
     59     }
     60   }
     61 
     62   std::vector<marisa::UInt32> key_ids;
     63   std::vector<std::string> keys;
     64   std::string str;
     65   while (std::getline(std::cin, str)) {
     66     std::size_t num_keys = trie.predict(str);
     67     if (num_keys != 0) {
     68       std::cout << num_keys << " found" << std::endl;
     69       key_ids.clear();
     70       keys.clear();
     71       try {
     72         if (depth_first_flag) {
     73           num_keys = trie.predict_depth_first(
     74               str, &key_ids, &keys, max_num_results);
     75         } else {
     76           num_keys = trie.predict_breadth_first(
     77               str, &key_ids, &keys, max_num_results);
     78         }
     79       } catch (const marisa::Exception &ex) {
     80         std::cerr << ex.filename() << ':' << ex.line() << ": " << ex.what()
     81             << ": failed to predict keys from: " << str << std::endl;
     82         return 30;
     83       }
     84       for (std::size_t i = 0; i < num_keys; ++i) {
     85         std::cout << key_ids[i] << '\t' << keys[i] << '\t' << str << '\n';
     86       }
     87     } else {
     88       std::cout << "not found" << std::endl;
     89     }
     90     if (!std::cout) {
     91       std::cerr << "error: failed to write results to standard output"
     92           << std::endl;
     93       return 31;
     94     }
     95   }
     96 
     97   return 0;
     98 }
     99 
    100 }  // namespace
    101 
    102 int main(int argc, char *argv[]) {
    103   std::ios::sync_with_stdio(false);
    104 
    105   ::cmdopt_option long_options[] = {
    106     { "max-num-results", 1, NULL, 'n' },
    107     { "depth-first", 0, NULL, 'd' },
    108     { "breadth-first", 0, NULL, 'b' },
    109     { "mmap-dictionary", 0, NULL, 'm' },
    110     { "read-dictionary", 0, NULL, 'r' },
    111     { "help", 0, NULL, 'h' },
    112     { NULL, 0, NULL, 0 }
    113   };
    114   ::cmdopt_t cmdopt;
    115   ::cmdopt_init(&cmdopt, argc, argv, "n:dbmrh", long_options);
    116   int label;
    117   while ((label = ::cmdopt_get(&cmdopt)) != -1) {
    118     switch (label) {
    119       case 'n': {
    120         char *end_of_value;
    121         const long value = std::strtol(cmdopt.optarg, &end_of_value, 10);
    122         if ((*end_of_value != '\0') || (value < 0)) {
    123           std::cerr << "error: option `-n' with an invalid argument: "
    124               << cmdopt.optarg << std::endl;
    125         }
    126         if ((value == 0) || ((unsigned long)value > MARISA_MAX_NUM_KEYS)) {
    127           max_num_results = MARISA_MAX_NUM_KEYS;
    128         } else {
    129           max_num_results = (std::size_t)(value);
    130         }
    131         break;
    132       }
    133       case 'd': {
    134         depth_first_flag = true;
    135         break;
    136       }
    137       case 'b': {
    138         depth_first_flag = false;
    139         break;
    140       }
    141       case 'm': {
    142         mmap_flag = true;
    143         break;
    144       }
    145       case 'r': {
    146         mmap_flag = false;
    147         break;
    148       }
    149       case 'h': {
    150         print_help(argv[0]);
    151         return 0;
    152       }
    153       default: {
    154         return 1;
    155       }
    156     }
    157   }
    158   return predict(cmdopt.argv + cmdopt.optind, cmdopt.argc - cmdopt.optind);
    159 }
    160