memo.c

   1 #include <stdlib.h>
   2 #include <string.h>
   3 #include "memo.h"
   4
   5 #define NDEBUG
   6 #include <assert.h>
   7
   8 enum {
   9   put_stack_size = 48,
  10 };
  11
  12 enum memo_type {
  13   MEMO_NODE,
  14   MEMO_LEAF,
  15 };
  16
  17 void memo_node_free(memo_node_ptr t) {
  18   if (!t) return;
  19   if (MEMO_LEAF == t->type) {
  20     free(((memo_leaf_ptr) t)->key);
  21   } else {
  22     memo_node_free(t->left);
  23     memo_node_free(t->right);
  24   }
  25   free(t);
  26 }
  27
  28 void *memo_alloc() {
  29   return malloc(sizeof(struct memo_s));
  30 }
  31
  32 void memo_clear(memo_ptr memo) {
  33   memo_node_free(memo->root);
  34 }
  35
  36 static inline int chartestbit(char c, int bit) {
  37   // Note we label the most significant bit 0, and the least 7.
  38   return (1 << (7 - bit)) & c;
  39 }
  40
  41 static int firstcritbit(const char *key0, const char *key1) {
  42   const char *cp0 = key0;
  43   const char *cp1 = key1;
  44   int bit;
  45
  46   while(*cp0 == *cp1) {
  47     if (!*cp0) {
  48       return 0;
  49     }
  50     cp0++;
  51     cp1++;
  52   }
  53
  54   char c = *cp0 ^ *cp1;
  55   for (bit = 7; !(c >> bit); bit--);
  56   // Subtract bit from 7 because we number them the other way.
  57   // Add 1 because we want to use the sign as an extra bit of information.
  58   // We'll subtract 1 from it later.
  59   int critbit = ((cp0 - key0) << 3) + 7 - bit + 1;
  60   if ((*cp0 >> bit) & 1) return critbit;
  61   return -critbit;
  62 }
  63
  64 static inline int testbit(const char *key, int bit) {
  65   return chartestbit(key[bit >> 3], (bit & 7));
  66 }
  67
  68 static memo_leaf_ptr memo_leaf_at(memo_node_ptr n, const char *key) {
  69   memo_node_ptr p = n;
  70   int len = (strlen(key) << 3) - 1;
  71   for (;;) {
  72     if (MEMO_LEAF == p->type) {
  73       return (memo_leaf_ptr) p;
  74     }
  75     if (len < p->critbit) {
  76       do {
  77         p = p->left;
  78       } while (MEMO_NODE == p->type);
  79       return (memo_leaf_ptr) p;
  80     }
  81     if (testbit(key, p->critbit)) {
  82       p = p->right;
  83     } else {
  84       p = p->left;
  85     }
  86   }
  87 }
  88
  89 int memo_has(memo_ptr memo, const char *key) {
  90   if (!memo->root) return 0;
  91   memo_leaf_ptr p = memo_leaf_at(memo->root, key);
  92   return (!strcmp(p->key, key));
  93 }
  94
  95 memo_it memo_it_at(memo_ptr memo, const char *key) {
  96   if (!memo->root) return NULL;
  97   memo_leaf_ptr p = memo_leaf_at(memo->root, key);
  98   if (!strcmp(p->key, key)) {
  99       return p;
 100   }
 101   return NULL;
 102 }
 103
 104 void *memo_at(memo_ptr memo, const char *key) {
 105   memo_leaf_ptr p = memo_it_at(memo, key);
 106   if (!p) return NULL;
 107   return p->data;
 108 }
 109
 110 static void memo_leaf_put(memo_leaf_ptr n, void *data, const char *key) {
 111   n->type = MEMO_LEAF;
 112   n->data = data;
 113   n->key = strdup(key);
 114 }
 115
 116 memo_it memo_put_with(memo_ptr memo, void *(*fn)(void *), const char *key) {
 117   if (!memo->root) {
 118     memo_leaf_ptr leaf = malloc(sizeof(memo_leaf_t));
 119     memo_leaf_put(leaf, fn(NULL), key);
 120     memo->root = (memo_node_ptr) leaf;
 121     return leaf;
 122   }
 123
 124   memo_node_ptr stack_space[put_stack_size];
 125   memo_node_ptr *stack = stack_space;
 126   int stack_max = put_stack_size;
 127   int i = 0;
 128   memo_node_ptr t = memo->root;
 129   int keylen = (strlen(key) << 3) - 1;
 130   stack[i] = t;
 131   while (MEMO_NODE == t->type) {
 132     // If the key is shorter than the remaining keys on this subtree,
 133     // we can compare it against any of them (and are guaranteed the new
 134     // node must be inserted above this node). However, this is uncommon thus
 135     // not worth optimizing. We let it loop through each time down the
 136     // left path.
 137     if (keylen < t->critbit || !testbit(key, t->critbit)) {
 138       t = t->left;
 139     } else {
 140       t = t->right;
 141     }
 142     i++;
 143     if (i == stack_max) {
 144       // Badly balanced tree? This is going to cost us.
 145       stack_max *= 2;
 146       if (stack == stack_space) {
 147         stack = malloc(stack_max * sizeof(memo_node_ptr));
 148         memcpy(stack, stack_space, put_stack_size * sizeof(memo_node_ptr));
 149       } else stack = realloc(stack, stack_max * sizeof(memo_node_ptr));
 150     }
 151     stack[i] = t;
 152   }
 153
 154   memo_leaf_ptr leaf = (memo_leaf_ptr) t;
 155   int res = firstcritbit(key, leaf->key);
 156   if (!res) {
 157     leaf->data = fn(leaf->data);
 158     return leaf;
 159   }
 160
 161   memo_leaf_ptr pleaf = malloc(sizeof(memo_leaf_t));
 162   memo_node_ptr pnode = malloc(sizeof(memo_node_t));
 163   memo_leaf_put(pleaf, fn(NULL), key);
 164   pnode->type = MEMO_NODE;
 165   pnode->critbit = abs(res) - 1;
 166
 167   // Last pointer on stack is a leaf, hence decrement before test.
 168   do {
 169     i--;
 170   } while(i >= 0 && pnode->critbit < stack[i]->critbit);
 171
 172   if (res > 0) {
 173     // Key is bigger, therefore it goes on the right.
 174     pnode->left = stack[i + 1];
 175     pnode->right = (memo_node_ptr) pleaf;
 176   } else {
 177     // Key is smaller, therefore it goes on the left.
 178     pnode->left = (memo_node_ptr) pleaf;
 179     pnode->right = stack[i + 1];
 180   }
 181   if (i < 0) {
 182     memo->root = pnode;
 183   } else {
 184     if (stack[i]->left == stack[i + 1]) {
 185       stack[i]->left = pnode;
 186     } else {
 187       stack[i]->right = pnode;
 188     }
 189   }
 190   if (stack != stack_space) free(stack);
 191   return pleaf;
 192 }
 193
 194 memo_it memo_put(memo_ptr memo, void *data, const char *key) {
 195   void *returndata(void *p) { return data; }
 196   return memo_put_with(memo, returndata, key);
 197 }
 198
 199 void *memo_remove(memo_ptr memo, const char *key) {
 200   assert(memo->root);
 201   // assert(memo_has(memo, key));  // TODO: write memo_has
 202   memo_node_ptr stack_space[put_stack_size];
 203   memo_node_ptr *stack = stack_space;
 204   int stack_max = put_stack_size;
 205   int i = 0;
 206   memo_node_ptr t = memo->root;
 207   stack[i] = t;
 208   while (MEMO_NODE == t->type) {
 209     assert((strlen(key) << 3) - 1 >= t->critbit);
 210     if (!testbit(key, t->critbit)) {
 211       t = t->left;
 212     } else {
 213       t = t->right;
 214     }
 215     i++;
 216     if (i == stack_max) {
 217       // Badly balanced tree? This is going to cost us.
 218       stack_max *= 2;
 219       if (stack == stack_space) {
 220         stack = malloc(stack_max * sizeof(memo_node_ptr));
 221         memcpy(stack, stack_space, put_stack_size * sizeof(memo_node_ptr));
 222       } else stack = realloc(stack, stack_max * sizeof(memo_node_ptr));
 223     }
 224     stack[i] = t;
 225   }
 226   memo_leaf_ptr p = (memo_leaf_ptr) t;
 227   if (0 == i) {  // We found it at the root.
 228     memo->root = NULL;
 229   } else {
 230     memo_node_ptr sibling;
 231     if (stack[i - 1]->left == stack[i]) {
 232       sibling = stack[i - 1]->right;
 233     } else {
 234       sibling = stack[i - 1]->left;
 235     }
 236     if (1 == i) {  // One-level down: reassign root.
 237       memo->root = sibling;
 238     } else {  // Reassign grandparent.
 239       if (stack[i - 2]->left == stack[i - 1]) {
 240         stack[i - 2]->left = sibling;
 241       } else {
 242         stack[i - 2]->right = sibling;
 243       }
 244     }
 245     free(stack[i - 1]);
 246   }
 247   free(p->key);
 248   void *data = p->data;
 249   free(p);
 250   return data;
 251 }
 252
 253 static int firstcritbit_u(const unsigned char *key0,
 254                           const unsigned char *key1, int len) {
 255   const unsigned char *cp0 = key0, *limit = key0 + len;
 256   const unsigned char *cp1 = key1;
 257
 258   for(;;) {
 259     if (cp0 == limit) return 0;
 260     if (*cp0 != *cp1) break;
 261     cp0++;
 262     cp1++;
 263   }
 264
 265   int bit;
 266
 267   char c = *cp0 ^ *cp1;
 268   for (bit = 7; !(c >> bit); bit--);
 269   // Subtract bit from 7 because we number them the other way.
 270   // Add 1 because we want to use the sign as an extra bit of information.
 271   // We'll subtract 1 from it later.
 272   int critbit = ((cp0 - key0) << 3) + 7 - bit + 1;
 273   if ((*cp0 >> bit) & 1) return critbit;
 274   return -critbit;
 275 }
 276
 277 static memo_leaf_ptr memo_leaf_at_u(memo_node_ptr n,
 278     const unsigned char *key, int len) {
 279   memo_node_ptr p = n;
 280   int bitlen = (len << 3) - 1;
 281   for (;;) {
 282     if (MEMO_LEAF == p->type) {
 283       return (memo_leaf_ptr) p;
 284     }
 285     // TODO: See comment in put_u about fixed length keys.
 286     if (bitlen < p->critbit) {
 287       do {
 288           p = p->left;
 289       } while (MEMO_NODE == p->type);
 290       return (memo_leaf_ptr) p;
 291     }
 292     if (testbit((char *) key, p->critbit)) {
 293       p = p->right;
 294     } else {
 295       p = p->left;
 296     }
 297   }
 298 }
 299
 300 void *memo_at_u(memo_ptr memo, const unsigned char *key, int len) {
 301   if (!memo->root) return NULL;
 302   memo_leaf_ptr p = memo_leaf_at_u(memo->root, key, len);
 303   if (!memcmp(p->key, key, len)) {
 304       return p->data;
 305   }
 306   return NULL;
 307 }
 308
 309 void *memo_sure_at_u(memo_ptr memo, const unsigned char *key, int len) {
 310   assert(memo_has_u(memo, key, len));
 311   return memo_leaf_at_u(memo->root, key, len)->data;
 312 }
 313
 314 static void memo_leaf_put_u(memo_leaf_ptr n, void *data,
 315     const unsigned char *key, int len) {
 316   n->type = MEMO_LEAF;
 317   n->data = data;
 318   n->key = malloc(len);
 319   memcpy(n->key, key, len);
 320 }
 321
 322 memo_leaf_ptr memo_put_u(memo_ptr memo,
 323     void *data, const unsigned char *key, int len) {
 324   if (!memo->root) {
 325     memo_leaf_ptr leaf = malloc(sizeof(memo_leaf_t));
 326     memo_leaf_put_u(leaf, data, key, len);
 327     memo->root = (memo_node_ptr) leaf;
 328     return leaf;
 329   }
 330
 331   memo_node_ptr stack_space[put_stack_size];
 332   memo_node_ptr *stack = stack_space;
 333   int stack_max = put_stack_size;
 334   int i = 0;
 335   memo_node_ptr t = memo->root;
 336   stack[i] = t;
 337   i++;
 338   while (MEMO_NODE == t->type) {
 339     // No length check: all keys should be the same length.
 340     if (!testbit((char *) key, t->critbit)) {
 341       t = t->left;
 342     } else {
 343       t = t->right;
 344     }
 345     if (i == stack_max) {
 346       // Badly balanced tree? This is going to cost us.
 347       stack_max *= 2;
 348       if (stack == stack_space) {
 349         stack = malloc(stack_max * sizeof(memo_node_ptr));
 350         memcpy(stack, stack_space, put_stack_size * sizeof(memo_node_ptr));
 351       } else stack = realloc(stack, stack_max * sizeof(memo_node_ptr));
 352     }
 353     stack[i] = t;
 354     i++;
 355   }
 356
 357   memo_leaf_ptr leaf = (memo_leaf_ptr) t;
 358   int res = firstcritbit_u(key, (unsigned char *) leaf->key, len);
 359   if (!res) {
 360     leaf->data = data;
 361     return leaf;
 362   }
 363
 364   memo_leaf_ptr pleaf = malloc(sizeof(memo_leaf_t));
 365   memo_node_ptr pnode = malloc(sizeof(memo_node_t));
 366   memo_leaf_put_u(pleaf, data, (unsigned char *) key, len);
 367   pnode->type = MEMO_NODE;
 368   pnode->critbit = abs(res) - 1;
 369
 370   // Decrement i so it points to top of stack.
 371   i--;
 372   // Last pointer on stack is a leaf, so decrement before test.
 373   do {
 374     i--;
 375   } while(i >= 0 && pnode->critbit < stack[i]->critbit);
 376
 377   if (res > 0) {
 378     // Key is bigger, therefore it goes on the right.
 379     pnode->left = stack[i + 1];
 380     pnode->right = (memo_node_ptr) pleaf;
 381   } else {
 382     // Key is smaller, therefore it goes on the left.
 383     pnode->left = (memo_node_ptr) pleaf;
 384     pnode->right = stack[i + 1];
 385   }
 386   if (i < 0) {
 387     memo->root = pnode;
 388   } else {
 389     if (stack[i]->left == stack[i + 1]) {
 390       stack[i]->left = pnode;
 391     } else {
 392       stack[i]->right = pnode;
 393     }
 394   }
 395   if (stack != stack_space) free(stack);
 396   return pleaf;
 397 }
 398
 399 static void remove_all_recurse(memo_node_ptr t,
 400     void (*fn)(void *, const char *)) {
 401   if (MEMO_LEAF == t->type) {
 402     memo_leaf_ptr p = (memo_leaf_ptr) t;
 403     if (fn) fn(p->data, p->key);
 404     free(p->key);
 405     free(p);
 406     return;
 407   }
 408   remove_all_recurse(t->left, fn);
 409   remove_all_recurse(t->right, fn);
 410   free(t);
 411 }
 412
 413 void memo_remove_all_with(memo_ptr memo, void (*fn)(void *data, const char *key)) {
 414   if (memo->root) {
 415     remove_all_recurse(memo->root, fn);
 416     memo->root = NULL;
 417   }
 418 }
 419
 420 void memo_remove_all(memo_ptr memo) {
 421   memo_remove_all_with(memo, NULL);
 422 }
 423
 424 void memo_forall(memo_t memo, void (*fn)(void *data, const char *key)) {
 425   void recurse(memo_node_ptr n) {
 426     if (MEMO_LEAF == n->type) {
 427       memo_leaf_ptr p = (memo_leaf_ptr) n;
 428       fn(p->data, p->key);
 429     } else {
 430       recurse(n->left);
 431       recurse(n->right);
 432     }
 433   }
 434   if (memo->root) {
 435     recurse(memo->root);
 436   }
 437 }
 438
 439 int memo_has_u(memo_ptr memo, const unsigned char *key, int len) {
 440   if (!memo->root) return 0;
 441   memo_leaf_ptr p = memo_leaf_at_u(memo->root, key, len);
 442   return (!memcmp(p->key, key, len));
 443 }
 444
 445 memo_it memo_it_at_u(memo_ptr memo, const unsigned char *key, int len) {
 446   if (!memo->root) return NULL;
 447   memo_leaf_ptr p = memo_leaf_at_u(memo->root, key, len);
 448   return memcmp(p->key, key, len) ? NULL : p;
 449 }
 450
 451 int memo_it_insert_u(memo_it *it,
 452     memo_ptr memo, const unsigned char *key, int len) {
 453   if (!memo->root) {
 454     memo_leaf_ptr leaf = malloc(sizeof(memo_leaf_t));
 455     memo_leaf_put_u(leaf, NULL, key, len);
 456     memo->root = (memo_node_ptr) leaf;
 457     *it = leaf;
 458     return 1;
 459   }
 460
 461   // The stack may seem cumbersome, but it seems to beat simpler solutions:
 462   // it appears we usually backtrack only a little, so rather than start from
 463   // the root again, we walk back up the tree.
 464   memo_node_ptr stack_space[put_stack_size];
 465   memo_node_ptr *stack = stack_space;
 466   int stack_max = put_stack_size;
 467   int i = 0;
 468   memo_node_ptr t = memo->root;
 469   stack[i] = t;
 470   while (MEMO_NODE == t->type) {
 471     // No length check: all keys should be the same length.
 472     if (!testbit((char *) key, t->critbit)) {
 473       t = t->left;
 474     } else {
 475       t = t->right;
 476     }
 477     i++;
 478     if (i == stack_max) {
 479       // Badly balanced tree? This is going to cost us.
 480       stack_max *= 2;
 481       if (stack == stack_space) {
 482         stack = malloc(stack_max * sizeof(memo_node_ptr));
 483         memcpy(stack, stack_space, put_stack_size * sizeof(memo_node_ptr));
 484       } else stack = realloc(stack, stack_max * sizeof(memo_node_ptr));
 485     }
 486     stack[i] = t;
 487   }
 488
 489   memo_leaf_ptr leaf = (memo_leaf_ptr) t;
 490   int res = firstcritbit_u(key, (unsigned char *) leaf->key, len);
 491   if (!res) {
 492     *it = leaf;
 493     return 0;
 494   }
 495
 496   memo_leaf_ptr pleaf = malloc(sizeof(memo_leaf_t));
 497   memo_node_ptr pnode = malloc(sizeof(memo_node_t));
 498   memo_leaf_put_u(pleaf, NULL, (unsigned char *) key, len);
 499   pnode->type = MEMO_NODE;
 500   pnode->critbit = abs(res) - 1;
 501
 502   // Walk back up tree to find place to insert new node.
 503   while(i > 0 && pnode->critbit < stack[i - 1]->critbit) i--;
 504
 505   if (res > 0) {
 506     // Key is bigger, therefore it goes on the right.
 507     pnode->left = stack[i];
 508     pnode->right = (memo_node_ptr) pleaf;
 509   } else {
 510     // Key is smaller, therefore it goes on the left.
 511     pnode->left = (memo_node_ptr) pleaf;
 512     pnode->right = stack[i];
 513   }
 514   if (!i) {
 515     memo->root = pnode;
 516   } else {
 517     if (stack[i - 1]->left == stack[i]) {
 518       stack[i - 1]->left = pnode;
 519     } else {
 520       stack[i - 1]->right = pnode;
 521     }
 522   }
 523   if (stack != stack_space) free(stack);
 524   *it = pleaf;
 525   return 1;
 526 }
 527
 528 int memo_it_insert(memo_it *it, memo_ptr memo, const char *key) {
 529   if (!memo->root) {
 530     memo_leaf_ptr leaf = malloc(sizeof(memo_leaf_t));
 531     memo_leaf_put(leaf, NULL, key);
 532     memo->root = (memo_node_ptr) leaf;
 533     *it = leaf;
 534     return 1;
 535   }
 536
 537   // The stack may seem cumbersome, but it seems to beat simpler solutions:
 538   // it appears we usually backtrack only a little, so rather than start from
 539   // the root again, we walk back up the tree.
 540   memo_node_ptr stack_space[put_stack_size];
 541   memo_node_ptr *stack = stack_space;
 542   int stack_max = put_stack_size;
 543   int i = 0;
 544   int keylen = (strlen(key) << 3) - 1;
 545   memo_node_ptr t = memo->root;
 546   stack[i] = t;
 547   while (MEMO_NODE == t->type) {
 548     // If the key is shorter than the remaining keys on this subtree,
 549     // we can compare it against any of them (and are guaranteed the new
 550     // node must be inserted above this node). However, this is uncommon thus
 551     // not worth optimizing. We let it loop through each time down the
 552     // left path.
 553     if (keylen < t->critbit || !testbit(key, t->critbit)) {
 554       t = t->left;
 555     } else {
 556       t = t->right;
 557     }
 558     i++;
 559     if (i == stack_max) {
 560       // Badly balanced tree? This is going to cost us.
 561       stack_max *= 2;
 562       if (stack == stack_space) {
 563         stack = malloc(stack_max * sizeof(memo_node_ptr));
 564         memcpy(stack, stack_space, put_stack_size * sizeof(memo_node_ptr));
 565       } else stack = realloc(stack, stack_max * sizeof(memo_node_ptr));
 566     }
 567     stack[i] = t;
 568   }
 569
 570   memo_leaf_ptr leaf = (memo_leaf_ptr) t;
 571   int res = firstcritbit(key, leaf->key);
 572   if (!res) {
 573     *it = leaf;
 574     return 0;
 575   }
 576
 577   memo_leaf_ptr pleaf = malloc(sizeof(memo_leaf_t));
 578   memo_node_ptr pnode = malloc(sizeof(memo_node_t));
 579   memo_leaf_put(pleaf, NULL, key);
 580   pnode->type = MEMO_NODE;
 581   pnode->critbit = abs(res) - 1;
 582
 583   // Walk back up tree to find place to insert new node.
 584   while(i > 0 && pnode->critbit < stack[i - 1]->critbit) i--;
 585
 586   if (res > 0) {
 587     // Key is bigger, therefore it goes on the right.
 588     pnode->left = stack[i];
 589     pnode->right = (memo_node_ptr) pleaf;
 590   } else {
 591     // Key is smaller, therefore it goes on the left.
 592     pnode->left = (memo_node_ptr) pleaf;
 593     pnode->right = stack[i];
 594   }
 595   if (!i) {
 596     memo->root = pnode;
 597   } else {
 598     if (stack[i - 1]->left == stack[i]) {
 599       stack[i - 1]->left = pnode;
 600     } else {
 601       stack[i - 1]->right = pnode;
 602     }
 603   }
 604   if (stack != stack_space) free(stack);
 605   *it = pleaf;
 606   return 1;
 607 }