1 #ifdef _WIN32
2 #include <fcntl.h>
3 #include <io.h>
4 #include <stdio.h>
5 #endif // _WIN32
6
7 #include <cstdlib>
8 #include <fstream>
9 #include <iostream>
10 #include <string>
11
12 #include <marisa.h>
13
14 #include "cmdopt.h"
15
16 namespace {
17
18 int param_num_tries = MARISA_DEFAULT_NUM_TRIES;
19 marisa::TailMode param_tail_mode = MARISA_DEFAULT_TAIL;
20 marisa::NodeOrder param_node_order = MARISA_DEFAULT_ORDER;
21 marisa::CacheLevel param_cache_level = MARISA_DEFAULT_CACHE;
22 const char *output_filename = NULL;
23
print_help(const char * cmd)24 void print_help(const char *cmd) {
25 std::cerr << "Usage: " << cmd << " [OPTION]... [FILE]...\n\n"
26 "Options:\n"
27 " -n, --num-tries=[N] limit the number of tries"
28 " [" << MARISA_MIN_NUM_TRIES << ", " << MARISA_MAX_NUM_TRIES
29 << "] (default: 3)\n"
30 " -t, --text-tail build a dictionary with text TAIL (default)\n"
31 " -b, --binary-tail build a dictionary with binary TAIL\n"
32 " -w, --weight-order arrange siblings in weight order (default)\n"
33 " -l, --label-order arrange siblings in label order\n"
34 " -c, --cache-level=[N] specify the cache size"
35 " [1, 5] (default: 3)\n"
36 " -o, --output=[FILE] write tries to FILE (default: stdout)\n"
37 " -h, --help print this help\n"
38 << std::endl;
39 }
40
read_keys(std::istream & input,marisa::Keyset * keyset)41 void read_keys(std::istream &input, marisa::Keyset *keyset) {
42 std::string line;
43 while (std::getline(input, line)) {
44 const std::string::size_type delim_pos = line.find_last_of('\t');
45 float weight = 1.0F;
46 if (delim_pos != line.npos) {
47 char *end_of_value;
48 weight = (float)std::strtod(&line[delim_pos + 1], &end_of_value);
49 if (*end_of_value == '\0') {
50 line.resize(delim_pos);
51 }
52 }
53 keyset->push_back(line.c_str(), line.length(), weight);
54 }
55 }
56
build(const char * const * args,std::size_t num_args)57 int build(const char * const *args, std::size_t num_args) {
58 marisa::Keyset keyset;
59 if (num_args == 0) try {
60 read_keys(std::cin, &keyset);
61 } catch (const marisa::Exception &ex) {
62 std::cerr << ex.what() << ": failed to read keys" << std::endl;
63 return 10;
64 }
65
66 for (std::size_t i = 0; i < num_args; ++i) try {
67 std::ifstream input_file(args[i], std::ios::binary);
68 if (!input_file) {
69 std::cerr << "error: failed to open: " << args[i] << std::endl;
70 return 11;
71 }
72 read_keys(input_file, &keyset);
73 } catch (const marisa::Exception &ex) {
74 std::cerr << ex.what() << ": failed to read keys" << std::endl;
75 return 12;
76 }
77
78 marisa::Trie trie;
79 try {
80 trie.build(keyset, param_num_tries | param_tail_mode | param_node_order |
81 param_cache_level);
82 } catch (const marisa::Exception &ex) {
83 std::cerr << ex.what() << ": failed to build a dictionary" << std::endl;
84 return 20;
85 }
86
87 std::cerr << "#keys: " << trie.num_keys() << std::endl;
88 std::cerr << "#nodes: " << trie.num_nodes() << std::endl;
89 std::cerr << "size: " << trie.io_size() << std::endl;
90
91 if (output_filename != NULL) {
92 try {
93 trie.save(output_filename);
94 } catch (const marisa::Exception &ex) {
95 std::cerr << ex.what() << ": failed to write a dictionary to file: "
96 << output_filename << std::endl;
97 return 30;
98 }
99 } else {
100 #ifdef _WIN32
101 const int stdout_fileno = ::_fileno(stdout);
102 if (stdout_fileno < 0) {
103 std::cerr << "error: failed to get the file descriptor of "
104 "standard output" << std::endl;
105 return 31;
106 }
107 if (::_setmode(stdout_fileno, _O_BINARY) == -1) {
108 std::cerr << "error: failed to set binary mode" << std::endl;
109 return 32;
110 }
111 #endif // _WIN32
112 try {
113 std::cout << trie;
114 } catch (const marisa::Exception &ex) {
115 std::cerr << ex.what()
116 << ": failed to write a dictionary to standard output" << std::endl;
117 return 33;
118 }
119 }
120 return 0;
121 }
122
123 } // namespace
124
main(int argc,char * argv[])125 int main(int argc, char *argv[]) {
126 std::ios::sync_with_stdio(false);
127
128 ::cmdopt_option long_options[] = {
129 { "max-num-tries", 1, NULL, 'n' },
130 { "text-tail", 0, NULL, 't' },
131 { "binary-tail", 0, NULL, 'b' },
132 { "weight-order", 0, NULL, 'w' },
133 { "label-order", 0, NULL, 'l' },
134 { "cache-level", 1, NULL, 'c' },
135 { "output", 1, NULL, 'o' },
136 { "help", 0, NULL, 'h' },
137 { NULL, 0, NULL, 0 }
138 };
139 ::cmdopt_t cmdopt;
140 ::cmdopt_init(&cmdopt, argc, argv, "n:tbwlc:o:h", long_options);
141 int label;
142 while ((label = ::cmdopt_get(&cmdopt)) != -1) {
143 switch (label) {
144 case 'n': {
145 char *end_of_value;
146 const long value = std::strtol(cmdopt.optarg, &end_of_value, 10);
147 if ((*end_of_value != '\0') || (value <= 0) ||
148 (value > MARISA_MAX_NUM_TRIES)) {
149 std::cerr << "error: option `-n' with an invalid argument: "
150 << cmdopt.optarg << std::endl;
151 return 1;
152 }
153 param_num_tries = (int)value;
154 break;
155 }
156 case 't': {
157 param_tail_mode = MARISA_TEXT_TAIL;
158 break;
159 }
160 case 'b': {
161 param_tail_mode = MARISA_BINARY_TAIL;
162 break;
163 }
164 case 'w': {
165 param_node_order = MARISA_WEIGHT_ORDER;
166 break;
167 }
168 case 'l': {
169 param_node_order = MARISA_LABEL_ORDER;
170 break;
171 }
172 case 'c': {
173 char *end_of_value;
174 const long value = std::strtol(cmdopt.optarg, &end_of_value, 10);
175 if ((*end_of_value != '\0') || (value < 1) || (value > 5)) {
176 std::cerr << "error: option `-c' with an invalid argument: "
177 << cmdopt.optarg << std::endl;
178 return 2;
179 } else if (value == 1) {
180 param_cache_level = MARISA_TINY_CACHE;
181 } else if (value == 2) {
182 param_cache_level = MARISA_SMALL_CACHE;
183 } else if (value == 3) {
184 param_cache_level = MARISA_NORMAL_CACHE;
185 } else if (value == 4) {
186 param_cache_level = MARISA_LARGE_CACHE;
187 } else if (value == 5) {
188 param_cache_level = MARISA_HUGE_CACHE;
189 }
190 break;
191 }
192 case 'o': {
193 output_filename = cmdopt.optarg;
194 break;
195 }
196 case 'h': {
197 print_help(argv[0]);
198 return 0;
199 }
200 default: {
201 return 1;
202 }
203 }
204 }
205 return build(cmdopt.argv + cmdopt.optind,
206 static_cast<std::size_t>(cmdopt.argc - cmdopt.optind));
207 }
208