tools/trietool.c

   1 /* -*- Mode: C; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */
   2 /*
   3  * trietool.c - Trie manipulation tool
   4  * Created: 2006-08-15
   5  * Author:  Theppitak Karoonboonyanan <theppitak@gmail.com>
   6  */
   7
   8 #include <config.h>
   9 #include <string.h>
  10 #include <stdlib.h>
  11 #include <stdio.h>
  12 #include <ctype.h>
  13 #include <locale.h>
  14 #if defined(HAVE_LOCALE_CHARSET)
  15 # include <localcharset.h>
  16 #elif defined (HAVE_LANGINFO_CODESET)
  17 # include <langinfo.h>
  18 # define locale_charset()  nl_langinfo(CODESET)
  19 #endif
  20 #include <iconv.h>
  21
  22 #include <assert.h>
  23
  24 #include <config.h>
  25 #include <datrie/trie.h>
  26
  27 /* iconv encoding name for AlphaChar string */
  28 #define ALPHA_ENC   "UCS-4LE"
  29
  30 #define N_ELEMENTS(a)   (sizeof(a)/sizeof((a)[0]))
  31
  32 typedef struct {
  33     const char *path;
  34     const char *trie_name;
  35     iconv_t     to_alpha_conv;
  36     iconv_t     from_alpha_conv;
  37     Trie       *trie;
  38 } ProgEnv;
  39
  40 static void init_conv           (ProgEnv *env);
  41 static size_t conv_to_alpha     (ProgEnv           *env,
  42                                  const char        *in,
  43                                  AlphaChar         *out,
  44                                  size_t             out_size);
  45 static size_t conv_from_alpha   (ProgEnv           *env,
  46                                  const AlphaChar   *in,
  47                                  char              *out,
  48                                  size_t             out_size);
  49 static void close_conv          (ProgEnv *env);
  50
  51 static int  prepare_trie        (ProgEnv *env);
  52 static int  close_trie          (ProgEnv *env);
  53
  54 static int  decode_switch       (int argc, char *argv[], ProgEnv *env);
  55 static int  decode_command      (int argc, char *argv[], ProgEnv *env);
  56
  57 static int  command_add         (int argc, char *argv[], ProgEnv *env);
  58 static int  command_add_list    (int argc, char *argv[], ProgEnv *env);
  59 static int  command_delete      (int argc, char *argv[], ProgEnv *env);
  60 static int  command_delete_list (int argc, char *argv[], ProgEnv *env);
  61 static int  command_query       (int argc, char *argv[], ProgEnv *env);
  62 static int  command_list        (int argc, char *argv[], ProgEnv *env);
  63
  64 static void usage               (const char *prog_name, int exit_status);
  65
  66 static char *string_trim        (char *s);
  67
  68 int
  69 main (int argc, char *argv[])
  70 {
  71     int     i;
  72     ProgEnv env;
  73     int     ret;
  74
  75     env.path = ".";
  76
  77     init_conv (&env);
  78
  79     i = decode_switch (argc, argv, &env);
  80     if (i == argc)
  81         usage (argv[0], EXIT_FAILURE);
  82
  83     env.trie_name = argv[i++];
  84
  85     if (prepare_trie (&env) != 0)
  86         exit (EXIT_FAILURE);
  87
  88     ret = decode_command (argc - i, argv + i, &env);
  89
  90     if (close_trie (&env) != 0)
  91         exit (EXIT_FAILURE);
  92
  93     close_conv (&env);
  94
  95     return ret;
  96 }
  97
  98 static void
  99 init_conv (ProgEnv *env)
 100 {
 101     const char *prev_locale;
 102     const char *locale_codeset;
 103
 104     prev_locale = setlocale (LC_CTYPE, "");
 105     locale_codeset = locale_charset();
 106     setlocale (LC_CTYPE, prev_locale);
 107
 108     env->to_alpha_conv = iconv_open (ALPHA_ENC, locale_codeset);
 109     env->from_alpha_conv = iconv_open (locale_codeset, ALPHA_ENC);
 110 }
 111
 112 static size_t
 113 conv_to_alpha (ProgEnv *env, const char *in, AlphaChar *out, size_t out_size)
 114 {
 115     char   *in_p = (char *) in;
 116     char   *out_p = (char *) out;
 117     size_t  in_left = strlen (in);
 118     size_t  out_left = out_size * sizeof (AlphaChar);
 119     size_t  res;
 120     const unsigned char *byte_p;
 121
 122     assert (sizeof (AlphaChar) == 4);
 123
 124     /* convert to UCS-4LE */
 125     res = iconv (env->to_alpha_conv, (char **) &in_p, &in_left,
 126                  &out_p, &out_left);
 127
 128     if (res == (size_t) -1)
 129         return res;
 130
 131     /* convert UCS-4LE to AlphaChar string */
 132     res = 0;
 133     for (byte_p = (const unsigned char *) out;
 134          res < out_size && byte_p + 3 < (unsigned char*) out_p;
 135          byte_p += 4)
 136     {
 137         out[res++] = byte_p[0]
 138                      | (byte_p[1] << 8)
 139                      | (byte_p[2] << 16)
 140                      | (byte_p[3] << 24);
 141     }
 142     if (res < out_size) {
 143         out[res] = 0;
 144     }
 145
 146     return res;
 147 }
 148
 149 static size_t
 150 conv_from_alpha (ProgEnv *env, const AlphaChar *in, char *out, size_t out_size)
 151 {
 152     size_t  in_left = alpha_char_strlen (in) * sizeof (AlphaChar);
 153     size_t  res;
 154
 155     assert (sizeof (AlphaChar) == 4);
 156
 157     /* convert AlphaChar to UCS-4LE */
 158     for (res = 0; in[res]; res++) {
 159         unsigned char  b[4];
 160
 161         b[0] = in[res] & 0xff;
 162         b[1] = (in[res] >> 8) & 0xff;
 163         b[2] = (in[res] >> 16) & 0xff;
 164         b[3] = (in[res] >> 24) & 0xff;
 165
 166         memcpy ((char *) &in[res], b, 4);
 167     }
 168
 169     /* convert UCS-4LE to locale codeset */
 170     res = iconv (env->from_alpha_conv, (char **) &in, &in_left,
 171                  &out, &out_size);
 172     *out = 0;
 173
 174     return res;
 175 }
 176
 177 static void
 178 close_conv (ProgEnv *env)
 179 {
 180     iconv_close (env->to_alpha_conv);
 181     iconv_close (env->from_alpha_conv);
 182 }
 183
 184 static int
 185 prepare_trie (ProgEnv *env)
 186 {
 187     char buff[256];
 188
 189     snprintf (buff, sizeof (buff),
 190               "%s/%s.tri", env->path, env->trie_name);
 191     env->trie = trie_new_from_file (buff);
 192
 193     if (!env->trie) {
 194         FILE       *sbm;
 195         AlphaMap   *alpha_map;
 196
 197         snprintf (buff, sizeof (buff),
 198                   "%s/%s.abm", env->path, env->trie_name);
 199         sbm = fopen (buff, "r");
 200         if (!sbm) {
 201             fprintf (stderr, "Cannot open alphabet map file %s\n", buff);
 202             return -1;
 203         }
 204
 205         alpha_map = alpha_map_new ();
 206
 207         while (fgets (buff, sizeof (buff), sbm)) {
 208             int         b, e;
 209
 210             /* read the range
 211              * format: [b,e]
 212              * where: b = begin char, e = end char; both in hex values
 213              */
 214             if (sscanf (buff, " [ %x , %x ] ", &b, &e) != 2)
 215                 continue;
 216             if (b > e) {
 217                 fprintf (stderr, "Range begin (%x) > range end (%x)\n", b, e);
 218                 continue;
 219             }
 220
 221             alpha_map_add_range (alpha_map, b, e);
 222         }
 223
 224         env->trie = trie_new (alpha_map);
 225
 226         alpha_map_free (alpha_map);
 227         fclose (sbm);
 228     }
 229
 230     return 0;
 231 }
 232
 233 static int
 234 close_trie (ProgEnv *env)
 235 {
 236     if (trie_is_dirty (env->trie)) {
 237         char path[256];
 238
 239         snprintf (path, sizeof (path),
 240                   "%s/%s.tri", env->path, env->trie_name);
 241         if (trie_save (env->trie, path) != 0) {
 242             fprintf (stderr, "Cannot save trie to %s\n", path);
 243             return -1;
 244         }
 245     }
 246
 247     trie_free (env->trie);
 248     return 0;
 249 }
 250
 251 static int
 252 decode_switch (int argc, char *argv[], ProgEnv *env)
 253 {
 254     int opt_idx;
 255
 256     for (opt_idx = 1; opt_idx < argc && *argv[opt_idx] == '-'; opt_idx++) {
 257         if (strcmp (argv[opt_idx], "-h") == 0 ||
 258             strcmp (argv[opt_idx], "--help") == 0)
 259         {
 260             usage (argv[0], EXIT_FAILURE);
 261         } else if (strcmp (argv[opt_idx], "-V") == 0 ||
 262                    strcmp (argv[opt_idx], "--version") == 0)
 263         {
 264             printf ("%s\n", VERSION);
 265             exit (EXIT_FAILURE);
 266         } else if (strcmp (argv[opt_idx], "-p") == 0 ||
 267                    strcmp (argv[opt_idx], "--path") == 0)
 268         {
 269             env->path = argv[++opt_idx];
 270         } else if (strcmp (argv[opt_idx], "--") == 0) {
 271             ++opt_idx;
 272             break;
 273         } else {
 274             fprintf (stderr, "Unknown option: %s\n", argv[opt_idx]);
 275             exit (EXIT_FAILURE);
 276         }
 277     }
 278
 279     return opt_idx;
 280 }
 281
 282 static int
 283 decode_command (int argc, char *argv[], ProgEnv *env)
 284 {
 285     int opt_idx;
 286
 287     for (opt_idx = 0; opt_idx < argc; opt_idx++) {
 288         if (strcmp (argv[opt_idx], "add") == 0) {
 289             ++opt_idx;
 290             opt_idx += command_add (argc - opt_idx, argv + opt_idx, env);
 291         } else if (strcmp (argv[opt_idx], "add-list") == 0) {
 292             ++opt_idx;
 293             opt_idx += command_add_list (argc - opt_idx, argv + opt_idx, env);
 294         } else if (strcmp (argv[opt_idx], "delete") == 0) {
 295             ++opt_idx;
 296             opt_idx += command_delete (argc - opt_idx, argv + opt_idx, env);
 297         } else if (strcmp (argv[opt_idx], "delete-list") == 0) {
 298             ++opt_idx;
 299             opt_idx += command_delete_list (argc - opt_idx, argv + opt_idx, env);
 300         } else if (strcmp (argv[opt_idx], "query") == 0) {
 301             ++opt_idx;
 302             opt_idx += command_query (argc - opt_idx, argv + opt_idx, env);
 303         } else if (strcmp (argv[opt_idx], "list") == 0) {
 304             ++opt_idx;
 305             opt_idx += command_list (argc - opt_idx, argv + opt_idx, env);
 306         } else {
 307             fprintf (stderr, "Unknown command: %s\n", argv[opt_idx]);
 308             return EXIT_FAILURE;
 309         }
 310     }
 311
 312     return EXIT_SUCCESS;
 313 }
 314
 315 static int
 316 command_add (int argc, char *argv[], ProgEnv *env)
 317 {
 318     int opt_idx;
 319
 320     opt_idx = 0;
 321     while (opt_idx < argc) {
 322         const char     *key;
 323         AlphaChar       key_alpha[256];
 324         TrieData        data;
 325
 326         key = argv[opt_idx++];
 327         data = (opt_idx < argc) ? atoi (argv[opt_idx++]) : TRIE_DATA_ERROR;
 328
 329         conv_to_alpha (env, key, key_alpha, N_ELEMENTS (key_alpha));
 330         if (!trie_store (env->trie, key_alpha, data)) {
 331             fprintf (stderr, "Failed to add entry '%s' with data %d\n",
 332                      key, data);
 333         }
 334     }
 335
 336     return opt_idx;
 337 }
 338
 339 static int
 340 command_add_list (int argc, char *argv[], ProgEnv *env)
 341 {
 342     const char *enc_name, *input_name;
 343     int         opt_idx;
 344     iconv_t     saved_conv;
 345     FILE       *input;
 346     char        line[256];
 347
 348     enc_name = 0;
 349     opt_idx = 0;
 350     saved_conv = env->to_alpha_conv;
 351     if (strcmp (argv[0], "-e") == 0 ||
 352         strcmp (argv[0], "--encoding") == 0)
 353     {
 354         if (++opt_idx >= argc) {
 355             fprintf (stderr, "add-list option \"%s\" requires encoding name",
 356                      argv[0]);
 357             return opt_idx;
 358         }
 359         enc_name = argv[opt_idx++];
 360     }
 361     if (opt_idx >= argc) {
 362         fprintf (stderr, "add-list requires input word list file name\n");
 363         return opt_idx;
 364     }
 365     input_name = argv[opt_idx++];
 366
 367     if (enc_name) {
 368         iconv_t conv = iconv_open (ALPHA_ENC, enc_name);
 369         if ((iconv_t) -1 == conv) {
 370             fprintf (stderr,
 371                     "Conversion from \"%s\" to \"%s\" is not supported.\n",
 372                     enc_name, ALPHA_ENC);
 373             return opt_idx;
 374         }
 375
 376         env->to_alpha_conv = conv;
 377     }
 378
 379     input = fopen (input_name, "r");
 380     if (!input) {
 381         fprintf (stderr, "add-list: Cannot open input file \"%s\"\n",
 382                  input_name);
 383         goto exit_iconv_openned;
 384     }
 385
 386     while (fgets (line, sizeof line, input)) {
 387         char       *key, *data;
 388         AlphaChar   key_alpha[256];
 389         TrieData    data_val;
 390
 391         key = string_trim (line);
 392         if ('\0' != *key) {
 393             /* find key boundary */
 394             for (data = key; *data && !strchr ("\t,", *data); ++data)
 395                 ;
 396             /* mark key ending and find data begin */
 397             if ('\0' != *data) {
 398                 *data++ = '\0';
 399                 while (isspace (*data))
 400                     ++data;
 401             }
 402             /* decode data */
 403             data_val = ('\0' != *data) ? atoi (data) : TRIE_DATA_ERROR;
 404
 405             /* store the key */
 406             conv_to_alpha (env, key, key_alpha, N_ELEMENTS (key_alpha));
 407             if (!trie_store (env->trie, key_alpha, data_val))
 408                 fprintf (stderr, "Failed to add key '%s' with data %d.\n",
 409                          key, data_val);
 410         }
 411     }
 412
 413     fclose (input);
 414
 415 exit_iconv_openned:
 416     if (enc_name) {
 417         iconv_close (env->to_alpha_conv);
 418         env->to_alpha_conv = saved_conv;
 419     }
 420
 421     return opt_idx;
 422 }
 423
 424 static int
 425 command_delete (int argc, char *argv[], ProgEnv *env)
 426 {
 427     int opt_idx;
 428
 429     for (opt_idx = 0; opt_idx < argc; opt_idx++) {
 430         AlphaChar   key_alpha[256];
 431
 432         conv_to_alpha (env, argv[opt_idx], key_alpha, N_ELEMENTS (key_alpha));
 433         if (!trie_delete (env->trie, key_alpha)) {
 434             fprintf (stderr, "No entry '%s'. Not deleted.\n", argv[opt_idx]);
 435         }
 436     }
 437
 438     return opt_idx;
 439 }
 440
 441 static int
 442 command_delete_list (int argc, char *argv[], ProgEnv *env)
 443 {
 444     const char *enc_name, *input_name;
 445     int         opt_idx;
 446     iconv_t     saved_conv;
 447     FILE   *input;
 448     char    line[256];
 449
 450     enc_name = 0;
 451     opt_idx = 0;
 452     saved_conv = env->to_alpha_conv;
 453     if (strcmp (argv[0], "-e") == 0 ||
 454         strcmp (argv[0], "--encoding") == 0)
 455     {
 456         if (++opt_idx >= argc) {
 457             fprintf (stderr, "delete-list option \"%s\" requires encoding name",
 458                      argv[0]);
 459             return opt_idx;
 460         }
 461         enc_name = argv[opt_idx++];
 462     }
 463     if (opt_idx >= argc) {
 464         fprintf (stderr, "delete-list requires input word list file name\n");
 465         return opt_idx;
 466     }
 467     input_name = argv[opt_idx++];
 468
 469     if (enc_name) {
 470         iconv_t conv = iconv_open (ALPHA_ENC, enc_name);
 471         if ((iconv_t) -1 == conv) {
 472             fprintf (stderr,
 473                     "Conversion from \"%s\" to \"%s\" is not supported.\n",
 474                     enc_name, ALPHA_ENC);
 475             return opt_idx;
 476         }
 477
 478         env->to_alpha_conv = conv;
 479     }
 480
 481     input = fopen (input_name, "r");
 482     if (!input) {
 483         fprintf (stderr, "delete-list: Cannot open input file \"%s\"\n",
 484                  input_name);
 485         goto exit_iconv_openned;
 486     }
 487
 488     while (fgets (line, sizeof line, input)) {
 489         char   *p;
 490
 491         p = string_trim (line);
 492         if ('\0' != *p) {
 493             AlphaChar   key_alpha[256];
 494
 495             conv_to_alpha (env, p, key_alpha, N_ELEMENTS (key_alpha));
 496             if (!trie_delete (env->trie, key_alpha)) {
 497                 fprintf (stderr, "No entry '%s'. Not deleted.\n", p);
 498             }
 499         }
 500     }
 501
 502     fclose (input);
 503
 504 exit_iconv_openned:
 505     if (enc_name) {
 506         iconv_close (env->to_alpha_conv);
 507         env->to_alpha_conv = saved_conv;
 508     }
 509
 510     return opt_idx;
 511 }
 512
 513 static int
 514 command_query (int argc, char *argv[], ProgEnv *env)
 515 {
 516     AlphaChar   key_alpha[256];
 517     TrieData    data;
 518
 519     if (argc == 0) {
 520         fprintf (stderr, "query: No key specified.\n");
 521         return 0;
 522     }
 523
 524     conv_to_alpha (env, argv[0], key_alpha, N_ELEMENTS (key_alpha));
 525     if (trie_retrieve (env->trie, key_alpha, &data)) {
 526         printf ("%d\n", data);
 527     } else {
 528         fprintf (stderr, "query: Key '%s' not found.\n", argv[0]);
 529     }
 530
 531     return 1;
 532 }
 533
 534 static Bool
 535 list_enum_func (const AlphaChar *key, TrieData key_data, void *user_data)
 536 {
 537     ProgEnv    *env = (ProgEnv *) user_data;
 538     char        key_locale[1024];
 539
 540     conv_from_alpha (env, key, key_locale, N_ELEMENTS (key_locale));
 541     printf ("%s\t%d\n", key_locale, key_data);
 542     return TRUE;
 543 }
 544
 545 static int
 546 command_list (int argc, char *argv[], ProgEnv *env)
 547 {
 548     trie_enumerate (env->trie, list_enum_func, (void *) env);
 549     return 0;
 550 }
 551
 552
 553 static void
 554 usage (const char *prog_name, int exit_status)
 555 {
 556     printf ("%s - double-array trie manipulator\n", prog_name);
 557     printf ("Usage: %s [OPTION]... TRIE CMD ARG ...\n", prog_name);
 558     printf (
 559         "Options:\n"
 560         "  -p, --path DIR           set trie directory to DIR [default=.]\n"
 561         "  -h, --help               display this help and exit\n"
 562         "  -V, --version            output version information and exit\n"
 563         "\n"
 564         "Commands:\n"
 565         "  add  WORD DATA ...\n"
 566         "      Add WORD with DATA to trie\n"
 567         "  add-list [OPTION] LISTFILE\n"
 568         "      Add words and data listed in LISTFILE to trie\n"
 569         "      Options:\n"
 570         "          -e, --encoding ENC    specify character encoding of LISTFILE\n"
 571         "  delete WORD ...\n"
 572         "      Delete WORD from trie\n"
 573         "  delete-list [OPTION] LISTFILE\n"
 574         "      Delete words listed in LISTFILE from trie\n"
 575         "      Options:\n"
 576         "          -e, --encoding ENC    specify character encoding of LISTFILE\n"
 577         "  query WORD\n"
 578         "      Query WORD data from trie\n"
 579         "  list\n"
 580         "      List all words in trie\n"
 581     );
 582
 583     exit (exit_status);
 584 }
 585
 586 static char *
 587 string_trim (char *s)
 588 {
 589     char   *p;
 590
 591     /* skip leading white spaces */
 592     while (*s && isspace (*s))
 593         ++s;
 594
 595     /* trim trailing white spaces */
 596     p = s + strlen (s) - 1;
 597     while (isspace (*p))
 598         --p;
 599     *++p = '\0';
 600
 601     return s;
 602 }
 603
 604 /*
 605 vi:ts=4:ai:expandtab
 606 */