1 #include <iostream>
2 #include <string>
3
4 #include <marisa.h>
5
6 #include "cmdopt.h"
7
8 namespace {
9
10 bool mmap_flag = true;
11
print_help(const char * cmd)12 void print_help(const char *cmd) {
13 std::cerr << "Usage: " << cmd << " [OPTION]... DIC\n\n"
14 "Options:\n"
15 " -m, --mmap-dictionary use memory-mapped I/O to load a dictionary"
16 " (default)\n"
17 " -r, --read-dictionary read an entire dictionary into memory\n"
18 " -h, --help print this help\n"
19 << std::endl;
20 }
21
lookup(const char * const * args,std::size_t num_args)22 int lookup(const char * const *args, std::size_t num_args) {
23 if (num_args == 0) {
24 std::cerr << "error: dictionary is not specified" << std::endl;
25 return 10;
26 } else if (num_args > 1) {
27 std::cerr << "error: more than one dictionaries are specified"
28 << std::endl;
29 return 11;
30 }
31
32 marisa::Trie trie;
33 if (mmap_flag) {
34 try {
35 trie.mmap(args[0]);
36 } catch (const marisa::Exception &ex) {
37 std::cerr << ex.what() << ": failed to mmap a dictionary file: "
38 << args[0] << std::endl;
39 return 20;
40 }
41 } else {
42 try {
43 trie.load(args[0]);
44 } catch (const marisa::Exception &ex) {
45 std::cerr << ex.what() << ": failed to load a dictionary file: "
46 << args[0] << std::endl;
47 return 21;
48 }
49 }
50
51 marisa::Agent agent;
52 std::string str;
53 while (std::getline(std::cin, str)) {
54 try {
55 agent.set_query(str.c_str(), str.length());
56 if (trie.lookup(agent)) {
57 std::cout << agent.key().id() << '\t' << str << '\n';
58 } else {
59 std::cout << "-1\t" << str << '\n';
60 }
61 } catch (const marisa::Exception &ex) {
62 std::cerr << ex.what() << ": lookup() failed: " << str << std::endl;
63 return 30;
64 }
65
66 if (!std::cout) {
67 std::cerr << "error: failed to write results to standard output"
68 << std::endl;
69 return 30;
70 }
71 }
72
73 return 0;
74 }
75
76 } // namespace
77
main(int argc,char * argv[])78 int main(int argc, char *argv[]) {
79 std::ios::sync_with_stdio(false);
80
81 ::cmdopt_option long_options[] = {
82 { "mmap-dictionary", 0, NULL, 'm' },
83 { "read-dictionary", 0, NULL, 'r' },
84 { "help", 0, NULL, 'h' },
85 { NULL, 0, NULL, 0 }
86 };
87 ::cmdopt_t cmdopt;
88 ::cmdopt_init(&cmdopt, argc, argv, "mrh", long_options);
89 int label;
90 while ((label = ::cmdopt_get(&cmdopt)) != -1) {
91 switch (label) {
92 case 'm': {
93 mmap_flag = true;
94 break;
95 }
96 case 'r': {
97 mmap_flag = false;
98 break;
99 }
100 case 'h': {
101 print_help(argv[0]);
102 return 0;
103 }
104 default: {
105 return 1;
106 }
107 }
108 }
109 return lookup(cmdopt.argv + cmdopt.optind,
110 static_cast<std::size_t>(cmdopt.argc - cmdopt.optind));
111 }
112