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