mdw@git.distorted.org.uk Git - sgt/agedu/blob - trie.c

   1 /*
   2  * trie.c: implementation of trie.h.
   3  */
   4
   5 #include <assert.h>
   6 #include <stdio.h>
   7 #include <string.h>
   8 #include <stdlib.h>
   9 #include <errno.h>
  10
  11 #include <sys/types.h>
  12 #include <unistd.h>
  13
  14 #include "agedu.h"
  15 #include "malloc.h"
  16 #include "trie.h"
  17
  18 #define alignof(typ) ( offsetof(struct { char c; typ t; }, t) )
  19
  20 /*
  21  * Compare functions for pathnames. Returns the relative order of
  22  * the names, like strcmp; also passes back the offset of the
  23  * first differing character if desired.
  24  */
  25 static int trieccmp(unsigned char a, unsigned char b)
  26 {
  27     a = (a == '\0' ? '\0' : a == pathsep ? '\1' : a+1);
  28     b = (b == '\0' ? '\0' : b == pathsep ? '\1' : b+1);
  29     return (int)a - (int)b;
  30 }
  31
  32 static int triencmp(const char *a, size_t alen,
  33                     const char *b, size_t blen, int *offset)
  34 {
  35     int off = 0;
  36     while (off < alen && off < blen && a[off] == b[off])
  37         off++;
  38     if (offset)
  39         *offset = off;
  40     if (off == alen || off == blen) return (off == blen) - (off == alen);
  41     return trieccmp(a[off], b[off]);
  42 }
  43
  44 static int triecmp(const char *a, const char *b, int *offset)
  45 {
  46     return triencmp(a, strlen(a), b, strlen(b), offset);
  47 }
  48
  49 /* ----------------------------------------------------------------------
  50  * Trie node structures.
  51  *
  52  * The trie format stored in the file consists of three distinct
  53  * node types, each with a distinguishing type field at the start.
  54  *
  55  * TRIE_LEAF is a leaf node; it contains an actual trie_file
  56  * structure, and indicates that when you're searching down the
  57  * trie with a string, you should now expect to encounter
  58  * end-of-string.
  59  *
  60  * TRIE_SWITCH indicates that the set of strings in the trie
  61  * include strings with more than one distinct character after the
  62  * prefix leading up to this point. Hence, it stores multiple
  63  * subnode pointers and a different single character for each one.
  64  *
  65  * TRIE_STRING indicates that at this point everything in the trie
  66  * has the same next few characters; it stores a single mandatory
  67  * string fragment and exactly one subnode pointer.
  68  */
  69 enum {
  70     TRIE_LEAF = 0x7fffe000,
  71     TRIE_SWITCH,
  72     TRIE_STRING
  73 };
  74
  75 struct trie_common {
  76     int type;
  77 };
  78
  79 struct trie_switchentry {
  80     off_t subnode;
  81     int subcount;
  82 };
  83
  84 struct trie_leaf {
  85     struct trie_common c;
  86     struct trie_file file;
  87 };
  88
  89 struct trie_switch {
  90     struct trie_common c;
  91     /*
  92      * sw[0] to sw[len-1] give the subnode pointers and element
  93      * counts. At &sw[len] is stored len single bytes which are
  94      * the characters corresponding to each subnode.
  95      */
  96     int len;
  97     struct trie_switchentry sw[];
  98 };
  99
 100 struct trie_string {
 101     struct trie_common c;
 102     int stringlen;
 103     off_t subnode;
 104     char string[];
 105 };
 106
 107 struct trie_header {
 108     unsigned long magic;
 109     off_t root, indexroot;
 110     int count;
 111     size_t maxpathlen;
 112     int pathsep;
 113 };
 114
 115 /* Union only used for computing alignment */
 116 union trie_node {
 117     struct trie_leaf leaf;
 118     struct {  /* fake trie_switch with indeterminate array length filled in */
 119         struct trie_common c;
 120         int len;
 121         struct trie_switchentry sw[1];
 122     } sw;
 123     struct {  /* fake trie_string with indeterminate array length filled in */
 124         struct trie_common c;
 125         int stringlen;
 126         off_t subnode;
 127         char string[1];
 128     } str;
 129 };
 130 #define TRIE_MAGIC 0x75646761UL
 131 #define TRIE_ALIGN alignof(union trie_node)
 132
 133 /* ----------------------------------------------------------------------
 134  * Trie-building functions.
 135  */
 136
 137 struct tbswitch {
 138     int len;
 139     char c[256];
 140     off_t off[256];
 141     int count[256];
 142 };
 143
 144 struct triebuild {
 145     int fd;
 146     off_t offset;
 147     char *lastpath;
 148     int lastlen, lastsize;
 149     off_t lastoff;
 150     struct tbswitch *switches;
 151     int switchsize;
 152     size_t maxpathlen;
 153 };
 154
 155 static void tb_seek(triebuild *tb, off_t off)
 156 {
 157     tb->offset = off;
 158     if (lseek(tb->fd, off, SEEK_SET) < 0) {
 159         fprintf(stderr, PNAME ": lseek: %s\n", strerror(errno));
 160         exit(1);
 161     }
 162 }
 163
 164 static void tb_write(triebuild *tb, const void *buf, size_t len)
 165 {
 166     tb->offset += len;
 167     while (len > 0) {
 168         int ret = write(tb->fd, buf, len);
 169         if (ret < 0) {
 170             fprintf(stderr, PNAME ": write: %s\n", strerror(errno));
 171             exit(1);
 172         }
 173         len -= ret;
 174         buf = (const void *)((const char *)buf + ret);
 175     }
 176 }
 177
 178 static char trie_align_zeroes[TRIE_ALIGN];
 179
 180 static void tb_align(triebuild *tb)
 181 {
 182     int off = (TRIE_ALIGN - ((tb)->offset % TRIE_ALIGN)) % TRIE_ALIGN;
 183     tb_write(tb, trie_align_zeroes, off);
 184 }
 185
 186 triebuild *triebuild_new(int fd)
 187 {
 188     triebuild *tb = snew(triebuild);
 189     struct trie_header th;
 190
 191     tb->fd = fd;
 192     tb->lastpath = NULL;
 193     tb->lastlen = tb->lastsize = 0;
 194     tb->lastoff = 0;
 195     tb->switches = NULL;
 196     tb->switchsize = 0;
 197     tb->maxpathlen = 0;
 198
 199     th.magic = TRIE_MAGIC;
 200     th.root = th.count = 0;
 201     th.indexroot = 0;
 202     th.maxpathlen = 0;
 203     th.pathsep = (unsigned char)pathsep;
 204
 205     tb_seek(tb, 0);
 206     tb_write(tb, &th, sizeof(th));
 207
 208     return tb;
 209 }
 210
 211 static off_t triebuild_unwind(triebuild *tb, int targetdepth, int *outcount)
 212 {
 213     off_t offset;
 214     int count, depth;
 215
 216     if (tb->lastoff == 0) {
 217         *outcount = 0;
 218         return 0;
 219     }
 220
 221     offset = tb->lastoff;
 222     count = 1;
 223     depth = tb->lastlen + 1;
 224
 225     assert(depth >= targetdepth);
 226
 227     while (depth > targetdepth) {
 228         int odepth = depth;
 229         while (depth > targetdepth &&
 230                (depth-1 > tb->switchsize || tb->switches[depth-1].len == 0))
 231             depth--;
 232         if (odepth > depth) {
 233             /*
 234              * Write out a string node.
 235              */
 236             size_t nodesize = sizeof(struct trie_string) + odepth - depth;
 237             struct trie_string *st = (struct trie_string *)smalloc(nodesize);
 238             st->c.type = TRIE_STRING;
 239             st->stringlen = odepth - depth;
 240             st->subnode = offset;
 241             memcpy(st->string, tb->lastpath + depth, odepth - depth);
 242             tb_align(tb);
 243             offset = tb->offset;
 244             tb_write(tb, st, nodesize);
 245             sfree(st);
 246         }
 247
 248         assert(depth >= targetdepth);
 249         if (depth <= targetdepth)
 250             break;
 251
 252         /*
 253          * Now we expect to be sitting just below a switch node.
 254          * Add our final entry to it and write it out.
 255          */
 256         depth--;
 257         {
 258             struct trie_switch *sw;
 259             char *chars;
 260             size_t nodesize;
 261             int swlen = tb->switches[depth].len;
 262             int i;
 263
 264             assert(swlen > 0);
 265
 266             tb->switches[depth].c[swlen] = tb->lastpath[depth];
 267             tb->switches[depth].off[swlen] = offset;
 268             tb->switches[depth].count[swlen] = count;
 269             swlen++;
 270
 271             nodesize = sizeof(struct trie_switch) +
 272                 swlen * sizeof(struct trie_switchentry) + swlen;
 273             sw = (struct trie_switch *)smalloc(nodesize);
 274             chars = (char *)&sw->sw[swlen];
 275
 276             sw->c.type = TRIE_SWITCH;
 277             sw->len = swlen;
 278             count = 0;
 279             for (i = 0; i < swlen; i++) {
 280                 sw->sw[i].subnode = tb->switches[depth].off[i];
 281                 sw->sw[i].subcount = tb->switches[depth].count[i];
 282                 chars[i] = tb->switches[depth].c[i];
 283
 284                 count += tb->switches[depth].count[i];
 285             }
 286
 287             tb_align(tb);
 288             offset = tb->offset;
 289             tb_write(tb, sw, nodesize);
 290             sfree(sw);
 291
 292             tb->switches[depth].len = 0;   /* clear this node */
 293         }
 294     }
 295
 296     *outcount = count;
 297     return offset;
 298 }
 299
 300 void triebuild_add(triebuild *tb, const char *pathname,
 301                    const struct trie_file *file)
 302 {
 303     int pathlen = strlen(pathname);
 304     int depth;
 305
 306     if (tb->maxpathlen < pathlen+1)
 307         tb->maxpathlen = pathlen+1;
 308
 309     if (tb->lastpath) {
 310         off_t offset;
 311         int count;
 312
 313         /*
 314          * Find the first differing character between this pathname
 315          * and the previous one.
 316          */
 317         int ret = triecmp(tb->lastpath, pathname, &depth);
 318         assert(ret < 0);
 319
 320         /*
 321          * Finalise all nodes above this depth.
 322          */
 323         offset = triebuild_unwind(tb, depth+1, &count);
 324
 325         /*
 326          * Add the final node we just acquired to the switch node
 327          * at our chosen depth, creating it if it isn't already
 328          * there.
 329          */
 330         if (tb->switchsize <= depth) {
 331             int oldsize = tb->switchsize;
 332             tb->switchsize = depth * 3 / 2 + 64;
 333             tb->switches = sresize(tb->switches, tb->switchsize,
 334                                    struct tbswitch);
 335             while (oldsize < tb->switchsize)
 336                 tb->switches[oldsize++].len = 0;
 337         }
 338
 339         tb->switches[depth].c[tb->switches[depth].len] = tb->lastpath[depth];
 340         tb->switches[depth].off[tb->switches[depth].len] = offset;
 341         tb->switches[depth].count[tb->switches[depth].len] = count;
 342         tb->switches[depth].len++;
 343     }
 344
 345     /*
 346      * Write out a leaf node for the new file, and remember its
 347      * file offset.
 348      */
 349     {
 350         struct trie_leaf leaf;
 351
 352         leaf.c.type = TRIE_LEAF;
 353         leaf.file = *file;             /* structure copy */
 354
 355         tb_align(tb);
 356         tb->lastoff = tb->offset;
 357         tb_write(tb, &leaf, sizeof(leaf));
 358     }
 359
 360     /*
 361      * Store this pathname for comparison with the next one.
 362      */
 363     if (tb->lastsize < pathlen+1) {
 364         tb->lastsize = pathlen * 3 / 2 + 64;
 365         tb->lastpath = sresize(tb->lastpath, tb->lastsize, char);
 366     }
 367     strcpy(tb->lastpath, pathname);
 368     tb->lastlen = pathlen;
 369 }
 370
 371 int triebuild_finish(triebuild *tb)
 372 {
 373     struct trie_header th;
 374
 375     th.magic = TRIE_MAGIC;
 376     th.root = triebuild_unwind(tb, 0, &th.count);
 377     th.indexroot = 0;
 378     th.maxpathlen = tb->maxpathlen;
 379     th.pathsep = (unsigned char)pathsep;
 380
 381     tb_seek(tb, 0);
 382     tb_write(tb, &th, sizeof(th));
 383
 384     return th.count;
 385 }
 386
 387 void triebuild_free(triebuild *tb)
 388 {
 389     sfree(tb->switches);
 390     sfree(tb->lastpath);
 391     sfree(tb);
 392 }
 393
 394 /* ----------------------------------------------------------------------
 395  * Querying functions.
 396  */
 397
 398 #define NODE(t, off, type) \
 399     ((const struct type *)((const char *)(t) + (off)))
 400
 401 size_t trie_maxpathlen(const void *t)
 402 {
 403     const struct trie_header *hdr = NODE(t, 0, trie_header);
 404     return hdr->maxpathlen;
 405 }
 406
 407 unsigned long trie_before(const void *t, const char *pathname)
 408 {
 409     const struct trie_header *hdr = NODE(t, 0, trie_header);
 410     int ret = 0, lastcount = hdr->count;
 411     int len = 1 + strlen(pathname), depth = 0;
 412     off_t off = hdr->root;
 413
 414     while (1) {
 415         const struct trie_common *node = NODE(t, off, trie_common);
 416         if (node->type == TRIE_LEAF) {
 417             if (depth < len)
 418                 ret += lastcount;   /* _shouldn't_ happen, but in principle */
 419             return ret;
 420         } else if (node->type == TRIE_STRING) {
 421             const struct trie_string *st = NODE(t, off, trie_string);
 422
 423             int offset;
 424             int cmp = triencmp(st->string, st->stringlen,
 425                                pathname + depth, len-depth, &offset);
 426
 427             if (offset < st->stringlen) {
 428                 if (cmp < 0)
 429                     ret += lastcount;
 430                 return ret;
 431             }
 432
 433             depth += st->stringlen;
 434             off = st->subnode;
 435         } else if (node->type == TRIE_SWITCH) {
 436             const struct trie_switch *sw = NODE(t, off, trie_switch);
 437             const char *chars = (const char *)&sw->sw[sw->len];
 438             int i;
 439
 440             for (i = 0; i < sw->len; i++) {
 441                 int c = chars[i];
 442                 int cmp = trieccmp(pathname[depth], c);
 443                 if (cmp > 0)
 444                     ret += sw->sw[i].subcount;
 445                 else if (cmp < 0)
 446                     return ret;
 447                 else {
 448                     off = sw->sw[i].subnode;
 449                     lastcount = sw->sw[i].subcount;
 450                     depth++;
 451                     break;
 452                 }
 453             }
 454             if (i == sw->len)
 455                 return ret;
 456         }
 457     }
 458 }
 459
 460 void trie_getpath(const void *t, unsigned long n, char *buf)
 461 {
 462     const struct trie_header *hdr = NODE(t, 0, trie_header);
 463     int depth = 0;
 464     off_t off = hdr->root;
 465
 466     while (1) {
 467         const struct trie_common *node = NODE(t, off, trie_common);
 468         if (node->type == TRIE_LEAF) {
 469             assert(depth > 0 && buf[depth-1] == '\0');
 470             return;
 471         } else if (node->type == TRIE_STRING) {
 472             const struct trie_string *st = NODE(t, off, trie_string);
 473
 474             memcpy(buf + depth, st->string, st->stringlen);
 475             depth += st->stringlen;
 476             off = st->subnode;
 477         } else if (node->type == TRIE_SWITCH) {
 478             const struct trie_switch *sw = NODE(t, off, trie_switch);
 479             const char *chars = (const char *)&sw->sw[sw->len];
 480             int i;
 481
 482             for (i = 0; i < sw->len; i++) {
 483                 if (n < sw->sw[i].subcount) {
 484                     buf[depth++] = chars[i];
 485                     off = sw->sw[i].subnode;
 486                     break;
 487                 } else
 488                     n -= sw->sw[i].subcount;
 489             }
 490             assert(i < sw->len);
 491         }
 492     }
 493 }
 494
 495 unsigned long trie_count(const void *t)
 496 {
 497     const struct trie_header *hdr = NODE(t, 0, trie_header);
 498     return hdr->count;
 499 }
 500
 501 char trie_pathsep(const void *t)
 502 {
 503     const struct trie_header *hdr = NODE(t, 0, trie_header);
 504     return (char)hdr->pathsep;
 505 }
 506
 507 struct triewalk_switch {
 508     const struct trie_switch *sw;
 509     int pos, depth, count;
 510 };
 511 struct triewalk {
 512     const void *t;
 513     struct triewalk_switch *switches;
 514     int nswitches, switchsize;
 515     int count;
 516 };
 517 triewalk *triewalk_new(const void *vt)
 518 {
 519     triewalk *tw = snew(triewalk);
 520
 521     tw->t = (const char *)vt;
 522     tw->switches = NULL;
 523     tw->switchsize = 0;
 524     tw->nswitches = -1;
 525     tw->count = 0;
 526
 527     return tw;
 528 }
 529 const struct trie_file *triewalk_next(triewalk *tw, char *buf)
 530 {
 531     off_t off;
 532     int depth;
 533
 534     if (tw->nswitches < 0) {
 535         const struct trie_header *hdr = NODE(tw->t, 0, trie_header);
 536         off = hdr->root;
 537         depth = 0;
 538         tw->nswitches = 0;
 539     } else {
 540         while (1) {
 541             int swpos;
 542             const struct trie_switch *sw;
 543             const char *chars;
 544
 545             if (tw->nswitches == 0) {
 546                 assert(tw->count == NODE(tw->t, 0, trie_header)->count);
 547                 return NULL;           /* run out of trie */
 548             }
 549
 550             swpos = tw->switches[tw->nswitches-1].pos;
 551             sw = tw->switches[tw->nswitches-1].sw;
 552             chars = (const char *)&sw->sw[sw->len];
 553
 554             if (swpos < sw->len) {
 555                 depth = tw->switches[tw->nswitches-1].depth;
 556                 off = sw->sw[swpos].subnode;
 557                 if (buf)
 558                     buf[depth++] = chars[swpos];
 559                 assert(tw->count == tw->switches[tw->nswitches-1].count);
 560                 tw->switches[tw->nswitches-1].count += sw->sw[swpos].subcount;
 561                 tw->switches[tw->nswitches-1].pos++;
 562                 break;
 563             }
 564
 565             tw->nswitches--;
 566         }
 567     }
 568
 569     while (1) {
 570         const struct trie_common *node = NODE(tw->t, off, trie_common);
 571         if (node->type == TRIE_LEAF) {
 572             const struct trie_leaf *lf = NODE(tw->t, off, trie_leaf);
 573             if (buf)
 574                 assert(depth > 0 && buf[depth-1] == '\0');
 575             tw->count++;
 576             return &lf->file;
 577         } else if (node->type == TRIE_STRING) {
 578             const struct trie_string *st = NODE(tw->t, off, trie_string);
 579
 580             if (buf)
 581                 memcpy(buf + depth, st->string, st->stringlen);
 582             depth += st->stringlen;
 583             off = st->subnode;
 584         } else if (node->type == TRIE_SWITCH) {
 585             const struct trie_switch *sw = NODE(tw->t, off, trie_switch);
 586             const char *chars = (const char *)&sw->sw[sw->len];
 587
 588             if (tw->nswitches >= tw->switchsize) {
 589                 tw->switchsize = tw->nswitches * 3 / 2 + 32;
 590                 tw->switches = sresize(tw->switches, tw->switchsize,
 591                                        struct triewalk_switch);
 592             }
 593
 594             tw->switches[tw->nswitches].sw = sw;
 595             tw->switches[tw->nswitches].pos = 1;
 596             tw->switches[tw->nswitches].depth = depth;
 597             tw->switches[tw->nswitches].count = tw->count + sw->sw[0].subcount;
 598             off = sw->sw[0].subnode;
 599             if (buf)
 600                 buf[depth++] = chars[0];
 601             tw->nswitches++;
 602         }
 603     }
 604 }
 605 void triewalk_free(triewalk *tw)
 606 {
 607     sfree(tw->switches);
 608     sfree(tw);
 609 }
 610
 611 void trie_set_index_offset(void *t, off_t ptr)
 612 {
 613     ((struct trie_header *)t)->indexroot = ptr;
 614 }
 615 off_t trie_get_index_offset(const void *t)
 616 {
 617     return ((const struct trie_header *)t)->indexroot;
 618 }
 619
 620 void make_successor(char *pathbuf)
 621 {
 622     int len = strlen(pathbuf);
 623     if (len > 0 && pathbuf[len-1] == pathsep)
 624         len--;
 625     pathbuf[len] = '\001';
 626     pathbuf[len+1] = '\0';
 627 }