epan/wmem/wmem_tree.c

   1 /* wmem_tree.c
   2  * Wireshark Memory Manager Red-Black Tree
   3  * Based on the red-black tree implementation in epan/emem.*
   4  * Copyright 2013, Evan Huus <eapache@gmail.com>
   5  *
   6  * $Id$
   7  *
   8  * Wireshark - Network traffic analyzer
   9  * By Gerald Combs <gerald@wireshark.org>
  10  * Copyright 1998 Gerald Combs
  11  *
  12  * This program is free software; you can redistribute it and/or modify
  13  * it under the terms of the GNU General Public License as published by
  14  * the Free Software Foundation; either version 2 of the License, or
  15  * (at your option) any later version.
  16  *
  17  * This program is distributed in the hope that it will be useful,
  18  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  19  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  20  * GNU General Public License for more details.
  21  *
  22  * You should have received a copy of the GNU General Public License along
  23  * with this program; if not, write to the Free Software Foundation, Inc.,
  24  * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  25  */
  26
  27 #include "config.h"
  28
  29 #include <ctype.h>
  30 #include <string.h>
  31 #include <stdio.h>
  32 #include <glib.h>
  33
  34 #include "wmem_core.h"
  35 #include "wmem_tree.h"
  36 #include "wmem_user_cb.h"
  37
  38 typedef enum _wmem_node_color_t {
  39     WMEM_NODE_COLOR_RED,
  40     WMEM_NODE_COLOR_BLACK
  41 } wmem_node_color_t;
  42
  43 struct _wmem_tree_node_t {
  44     struct _wmem_tree_node_t *parent;
  45     struct _wmem_tree_node_t *left;
  46     struct _wmem_tree_node_t *right;
  47
  48     void    *data;
  49     guint32  key32;
  50
  51     wmem_node_color_t color;
  52     gboolean          is_subtree;
  53 };
  54
  55 typedef struct _wmem_tree_node_t wmem_tree_node_t;
  56
  57 struct _wmem_tree_t {
  58     wmem_allocator_t *master;
  59     wmem_allocator_t *allocator;
  60     wmem_tree_node_t *root;
  61     guint             master_cb_id;
  62     guint             slave_cb_id;
  63 };
  64
  65 static wmem_tree_node_t *
  66 node_uncle(wmem_tree_node_t *node)
  67 {
  68     wmem_tree_node_t *parent, *grandparent;
  69
  70     parent = node->parent;
  71     if (parent == NULL) {
  72         return NULL;
  73     }
  74
  75     grandparent = parent->parent;
  76     if (grandparent == NULL) {
  77         return NULL;
  78     }
  79
  80     if (parent == grandparent->left) {
  81         return grandparent->right;
  82     }
  83     else {
  84         return grandparent->left;
  85     }
  86 }
  87
  88 static void rb_insert_case1(wmem_tree_t *tree, wmem_tree_node_t *node);
  89 static void rb_insert_case2(wmem_tree_t *tree, wmem_tree_node_t *node);
  90
  91 static void
  92 rotate_left(wmem_tree_t *tree, wmem_tree_node_t *node)
  93 {
  94     if (node->parent) {
  95         if (node->parent->left == node) {
  96             node->parent->left = node->right;
  97         }
  98         else {
  99             node->parent->right = node->right;
 100         }
 101     }
 102     else {
 103         tree->root = node->right;
 104     }
 105
 106     node->right->parent = node->parent;
 107     node->parent        = node->right;
 108     node->right         = node->right->left;
 109     if (node->right) {
 110         node->right->parent = node;
 111     }
 112     node->parent->left = node;
 113 }
 114
 115 static void
 116 rotate_right(wmem_tree_t *tree, wmem_tree_node_t *node)
 117 {
 118     if (node->parent) {
 119         if (node->parent->left == node) {
 120             node->parent->left = node->left;
 121         }
 122         else {
 123             node->parent->right = node->left;
 124         }
 125     }
 126     else {
 127         tree->root = node->left;
 128     }
 129
 130     node->left->parent = node->parent;
 131     node->parent       = node->left;
 132     node->left         = node->left->right;
 133     if (node->left) {
 134         node->left->parent = node;
 135     }
 136     node->parent->right = node;
 137 }
 138
 139 static void
 140 rb_insert_case5(wmem_tree_t *tree, wmem_tree_node_t *node)
 141 {
 142     wmem_tree_node_t *parent, *grandparent;
 143
 144     parent      = node->parent;
 145     grandparent = parent->parent;
 146
 147     parent->color      = WMEM_NODE_COLOR_BLACK;
 148     grandparent->color = WMEM_NODE_COLOR_RED;
 149
 150     if (node == parent->left && parent == grandparent->left) {
 151         rotate_right(tree, grandparent);
 152     }
 153     else {
 154         rotate_left(tree, grandparent);
 155     }
 156 }
 157
 158 static void
 159 rb_insert_case4(wmem_tree_t *tree, wmem_tree_node_t *node)
 160 {
 161     wmem_tree_node_t *parent, *grandparent;
 162
 163     parent      = node->parent;
 164     grandparent = parent->parent;
 165     if (!grandparent) {
 166         return;
 167     }
 168
 169     if (node == parent->right && parent == grandparent->left) {
 170         rotate_left(tree, parent);
 171         node = node->left;
 172     }
 173     else if (node == parent->left && parent == grandparent->right) {
 174         rotate_right(tree, parent);
 175         node = node->right;
 176     }
 177
 178     rb_insert_case5(tree, node);
 179 }
 180
 181 static void
 182 rb_insert_case3(wmem_tree_t *tree, wmem_tree_node_t *node)
 183 {
 184     wmem_tree_node_t *parent, *grandparent, *uncle;
 185
 186     uncle = node_uncle(node);
 187
 188     if (uncle && uncle->color == WMEM_NODE_COLOR_RED) {
 189         parent      = node->parent;
 190         grandparent = parent->parent;
 191
 192         parent->color      = WMEM_NODE_COLOR_BLACK;
 193         uncle->color       = WMEM_NODE_COLOR_BLACK;
 194         grandparent->color = WMEM_NODE_COLOR_RED;
 195
 196         rb_insert_case1(tree, grandparent);
 197     }
 198     else {
 199         rb_insert_case4(tree, node);
 200     }
 201 }
 202
 203 static void
 204 rb_insert_case2(wmem_tree_t *tree, wmem_tree_node_t *node)
 205 {
 206     /* parent is always non-NULL here */
 207     if (node->parent->color == WMEM_NODE_COLOR_RED) {
 208         rb_insert_case3(tree, node);
 209     }
 210 }
 211
 212 static void
 213 rb_insert_case1(wmem_tree_t *tree, wmem_tree_node_t *node)
 214 {
 215     wmem_tree_node_t *parent = node->parent;
 216
 217     if (parent == NULL) {
 218         node->color = WMEM_NODE_COLOR_BLACK;
 219     }
 220     else {
 221         rb_insert_case2(tree, node);
 222     }
 223 }
 224
 225 wmem_tree_t *
 226 wmem_tree_new(wmem_allocator_t *allocator)
 227 {
 228     wmem_tree_t *tree;
 229
 230     tree = wmem_new(allocator, wmem_tree_t);
 231     tree->master    = allocator;
 232     tree->allocator = allocator;
 233     tree->root      = NULL;
 234
 235     return tree;
 236 }
 237
 238 static gboolean
 239 wmem_tree_reset_cb(wmem_allocator_t *allocator _U_, wmem_cb_event_t event,
 240         void *user_data)
 241 {
 242     wmem_tree_t *tree = (wmem_tree_t *)user_data;
 243
 244     tree->root = NULL;
 245
 246     if (event == WMEM_CB_DESTROY_EVENT) {
 247         wmem_unregister_callback(tree->master, tree->master_cb_id);
 248         wmem_free(tree->master, tree);
 249     }
 250
 251     return TRUE;
 252 }
 253
 254 static gboolean
 255 wmem_tree_destroy_cb(wmem_allocator_t *allocator _U_, wmem_cb_event_t event _U_,
 256         void *user_data)
 257 {
 258     wmem_tree_t *tree = (wmem_tree_t *)user_data;
 259
 260     wmem_unregister_callback(tree->allocator, tree->slave_cb_id);
 261
 262     return FALSE;
 263 }
 264
 265 wmem_tree_t *
 266 wmem_tree_new_autoreset(wmem_allocator_t *master, wmem_allocator_t *slave)
 267 {
 268     wmem_tree_t *tree;
 269
 270     tree = wmem_new(master, wmem_tree_t);
 271     tree->master    = master;
 272     tree->allocator = slave;
 273     tree->root      = NULL;
 274
 275     tree->master_cb_id = wmem_register_callback(master, wmem_tree_destroy_cb,
 276             tree);
 277     tree->slave_cb_id  = wmem_register_callback(slave, wmem_tree_reset_cb,
 278             tree);
 279
 280     return tree;
 281 }
 282
 283 gboolean
 284 wmem_tree_is_empty(wmem_tree_t *tree)
 285 {
 286     return tree->root == NULL;
 287 }
 288
 289 static wmem_tree_node_t *
 290 create_node(wmem_allocator_t *allocator, wmem_tree_node_t *parent, guint32 key,
 291         void *data, wmem_node_color_t color, gboolean is_subtree)
 292 {
 293     wmem_tree_node_t *node;
 294
 295     node = wmem_new(allocator, wmem_tree_node_t);
 296
 297     node->left   = NULL;
 298     node->right  = NULL;
 299     node->parent = parent;
 300
 301     node->key32 = key;
 302     node->data  = data;
 303
 304     node->color      = color;
 305     node->is_subtree = is_subtree;
 306
 307     return node;
 308 }
 309
 310 #define CREATE_DATA(TRANSFORM, DATA) ((TRANSFORM) ? (TRANSFORM)(DATA) : (DATA))
 311 static void *
 312 lookup_or_insert32(wmem_tree_t *tree, guint32 key,
 313         void*(*func)(void*), void* data, gboolean is_subtree, gboolean replace)
 314 {
 315     wmem_tree_node_t *node     = tree->root;
 316     wmem_tree_node_t *new_node = NULL;
 317
 318     /* is this the first node ?*/
 319     if (!node) {
 320         new_node = create_node(tree->allocator, NULL, key,
 321                 CREATE_DATA(func, data), WMEM_NODE_COLOR_BLACK, is_subtree);
 322         tree->root = new_node;
 323         return new_node->data;
 324     }
 325
 326     /* it was not the new root so walk the tree until we find where to
 327      * insert this new leaf.
 328      */
 329     while (!new_node) {
 330         /* this node already exists, so just return the data pointer*/
 331         if (key == node->key32) {
 332             if (replace) {
 333                 node->data = CREATE_DATA(func, data);
 334             }
 335             return node->data;
 336         }
 337         else if (key < node->key32) {
 338             if (node->left) {
 339                 node = node->left;
 340             }
 341             else {
 342                 /* new node to the left */
 343                 new_node = create_node(tree->allocator, node, key,
 344                         CREATE_DATA(func, data), WMEM_NODE_COLOR_RED,
 345                         is_subtree);
 346                 node->left = new_node;
 347             }
 348         }
 349         else if (key > node->key32) {
 350             if (node->right) {
 351                 node = node->right;
 352             }
 353             else {
 354                 /* new node to the right */
 355                 new_node = create_node(tree->allocator, node, key,
 356                         CREATE_DATA(func, data), WMEM_NODE_COLOR_RED,
 357                         is_subtree);
 358                 node->right = new_node;
 359             }
 360         }
 361     }
 362
 363     /* node will now point to the newly created node */
 364     rb_insert_case1(tree, new_node);
 365
 366     return new_node->data;
 367 }
 368
 369 void
 370 wmem_tree_insert32(wmem_tree_t *tree, guint32 key, void *data)
 371 {
 372     lookup_or_insert32(tree, key, NULL, data, FALSE, TRUE);
 373 }
 374
 375 void *
 376 wmem_tree_lookup32(wmem_tree_t *tree, guint32 key)
 377 {
 378     wmem_tree_node_t *node = tree->root;
 379
 380     while (node) {
 381         if (key == node->key32) {
 382             return node->data;
 383         }
 384         else if (key < node->key32) {
 385             node = node->left;
 386         }
 387         else if (key > node->key32) {
 388             node = node->right;
 389         }
 390     }
 391
 392     return NULL;
 393 }
 394
 395 void *
 396 wmem_tree_lookup32_le(wmem_tree_t *tree, guint32 key)
 397 {
 398     wmem_tree_node_t *node = tree->root;
 399
 400     while (node) {
 401         if (key == node->key32) {
 402             return node->data;
 403         }
 404         else if (key < node->key32) {
 405             if (node->left == NULL) {
 406                 break;
 407             }
 408             node = node->left;
 409         }
 410         else if (key > node->key32) {
 411             if (node->right == NULL) {
 412                 break;
 413             }
 414             node = node->right;
 415         }
 416     }
 417
 418     if (!node) {
 419         return NULL;
 420     }
 421
 422     /* If we are still at the root of the tree this means that this node
 423      * is either smaller than the search key and then we return this
 424      * node or else there is no smaller key available and then
 425      * we return NULL.
 426      */
 427     if (node->parent == NULL) {
 428         if (key > node->key32) {
 429             return node->data;
 430         } else {
 431             return NULL;
 432         }
 433     }
 434
 435     if (node->key32 <= key) {
 436         /* if our key is <= the search key, we have the right node */
 437         return node->data;
 438     }
 439     else if (node == node->parent->left) {
 440         /* our key is bigger than the search key and we're a left child,
 441          * we have to check if any of our ancestors are smaller. */
 442         while (node) {
 443             if (key > node->key32) {
 444                 return node->data;
 445             }
 446             node=node->parent;
 447         }
 448         return NULL;
 449     }
 450     else {
 451         /* our key is bigger than the search key and we're a right child,
 452          * our parent is the one we want */
 453         return node->parent->data;
 454     }
 455 }
 456
 457 /* Strings are stored as an array of uint32 containing the string characters
 458    with 4 characters in each uint32.
 459    The first byte of the string is stored as the most significant byte.
 460    If the string is not a multiple of 4 characters in length the last
 461    uint32 containing the string bytes are padded with 0 bytes.
 462    After the uint32's containing the string, there is one final terminator
 463    uint32 with the value 0x00000001
 464 */
 465 static guint32 *
 466 pack_string(const gchar *key, guint32 *divx, guint32 flags)
 467 {
 468     guint32 *aligned = NULL;
 469     guint32 len = (guint32) strlen(key);
 470     guint32 i;
 471     guint32 tmp;
 472
 473     *divx = (len+3)/4 + 1;
 474
 475     aligned = (guint32 *)wmem_alloc(NULL, *divx * sizeof (guint32));
 476
 477     /* pack the bytes one one by one into guint32s */
 478     tmp = 0;
 479     for (i = 0;i < len;i++) {
 480         unsigned char ch;
 481
 482         ch = (unsigned char)key[i];
 483         if (flags & WMEM_TREE_STRING_NOCASE) {
 484             if (isupper(ch)) {
 485                 ch = tolower(ch);
 486             }
 487         }
 488         tmp <<= 8;
 489         tmp |= ch;
 490         if (i%4 == 3) {
 491             aligned[i/4] = tmp;
 492             tmp = 0;
 493         }
 494     }
 495     /* add required padding to the last uint32 */
 496     if (i%4 != 0) {
 497         while (i%4 != 0) {
 498             i++;
 499             tmp <<= 8;
 500         }
 501         aligned[i/4-1] = tmp;
 502     }
 503
 504     /* add the terminator */
 505     aligned[*divx-1] = 0x00000001;
 506
 507     return aligned;
 508 }
 509
 510 void
 511 wmem_tree_insert_string(wmem_tree_t* tree, const gchar* k, void* v, guint32 flags)
 512 {
 513     wmem_tree_key_t key[2];
 514     guint32 *aligned;
 515     guint32 divx;
 516
 517     aligned = pack_string(k, &divx, flags);
 518
 519     key[0].length = divx;
 520     key[0].key = aligned;
 521     key[1].length = 0;
 522     key[1].key = NULL;
 523
 524     wmem_tree_insert32_array(tree, key, v);
 525     wmem_free(NULL, aligned);
 526 }
 527
 528 void *
 529 wmem_tree_lookup_string(wmem_tree_t* tree, const gchar* k, guint32 flags)
 530 {
 531     wmem_tree_key_t key[2];
 532     guint32 *aligned;
 533     guint32 divx;
 534     void *ret;
 535
 536     aligned = pack_string(k, &divx, flags);
 537
 538     key[0].length = divx;
 539     key[0].key = aligned;
 540     key[1].length = 0;
 541     key[1].key = NULL;
 542
 543     ret = wmem_tree_lookup32_array(tree, key);
 544     wmem_free(NULL, aligned);
 545     return ret;
 546 }
 547
 548 static void *
 549 create_sub_tree(void* d)
 550 {
 551     return wmem_tree_new(((wmem_tree_t *)d)->allocator);
 552 }
 553
 554 void
 555 wmem_tree_insert32_array(wmem_tree_t *tree, wmem_tree_key_t *key, void *data)
 556 {
 557     wmem_tree_t *insert_tree = NULL;
 558     wmem_tree_key_t *cur_key;
 559     guint32 i, insert_key32 = 0;
 560
 561     for (cur_key = key; cur_key->length > 0; cur_key++) {
 562         for (i = 0; i < cur_key->length; i++) {
 563             /* Insert using the previous key32 */
 564             if (!insert_tree) {
 565                 insert_tree = tree;
 566             } else {
 567                 insert_tree = (wmem_tree_t *)lookup_or_insert32(insert_tree,
 568                         insert_key32, create_sub_tree, tree, TRUE, FALSE);
 569             }
 570             insert_key32 = cur_key->key[i];
 571         }
 572     }
 573
 574     g_assert(insert_tree);
 575
 576     wmem_tree_insert32(insert_tree, insert_key32, data);
 577 }
 578
 579 void *
 580 wmem_tree_lookup32_array_helper(wmem_tree_t *tree, wmem_tree_key_t *key,
 581         void*(*helper)(wmem_tree_t*, guint32))
 582 {
 583     wmem_tree_t *lookup_tree = NULL;
 584     wmem_tree_key_t *cur_key;
 585     guint32 i, lookup_key32 = 0;
 586
 587     if (!tree || !key) {
 588         return NULL;
 589     }
 590
 591     for (cur_key = key; cur_key->length > 0; cur_key++) {
 592         for (i = 0; i < cur_key->length; i++) {
 593             /* Lookup using the previous key32 */
 594             if (!lookup_tree) {
 595                 lookup_tree = tree;
 596             }
 597             else {
 598                 lookup_tree =
 599                     (wmem_tree_t *)(*helper)(lookup_tree, lookup_key32);
 600                 if (!lookup_tree) {
 601                     return NULL;
 602                 }
 603             }
 604             lookup_key32 = cur_key->key[i];
 605         }
 606     }
 607
 608     /* Assert if we didn't get any valid keys */
 609     g_assert(lookup_tree);
 610
 611     return (*helper)(lookup_tree, lookup_key32);
 612 }
 613
 614 void *
 615 wmem_tree_lookup32_array(wmem_tree_t *tree, wmem_tree_key_t *key)
 616 {
 617     return wmem_tree_lookup32_array_helper(tree, key, wmem_tree_lookup32);
 618 }
 619
 620 void *
 621 wmem_tree_lookup32_array_le(wmem_tree_t *tree, wmem_tree_key_t *key)
 622 {
 623     return wmem_tree_lookup32_array_helper(tree, key, wmem_tree_lookup32_le);
 624 }
 625
 626 static gboolean
 627 wmem_tree_foreach_nodes(wmem_tree_node_t* node, wmem_foreach_func callback,
 628         void *user_data)
 629 {
 630     gboolean stop_traverse = FALSE;
 631
 632     if (!node) {
 633         return FALSE;
 634     }
 635
 636     if (node->left) {
 637         if (wmem_tree_foreach_nodes(node->left, callback, user_data)) {
 638             return TRUE;
 639         }
 640     }
 641
 642     if (node->is_subtree == TRUE) {
 643         stop_traverse = wmem_tree_foreach((wmem_tree_t *)node->data,
 644                 callback, user_data);
 645     } else {
 646         stop_traverse = callback(node->data, user_data);
 647     }
 648
 649     if (stop_traverse) {
 650         return TRUE;
 651     }
 652
 653     if(node->right) {
 654         if (wmem_tree_foreach_nodes(node->right, callback, user_data)) {
 655             return TRUE;
 656         }
 657     }
 658
 659     return FALSE;
 660 }
 661
 662 gboolean
 663 wmem_tree_foreach(wmem_tree_t* tree, wmem_foreach_func callback,
 664         void *user_data)
 665 {
 666     if(!tree->root)
 667         return FALSE;
 668
 669     return wmem_tree_foreach_nodes(tree->root, callback, user_data);
 670 }
 671
 672 static void wmem_print_subtree(wmem_tree_t *tree, guint32 level);
 673
 674 static void
 675 wmem_tree_print_nodes(const char *prefix, wmem_tree_node_t *node, guint32 level)
 676 {
 677     guint32 i;
 678
 679     if (!node)
 680         return;
 681
 682     for (i=0; i<level; i++) {
 683         printf("    ");
 684     }
 685
 686     printf("%sNODE:%p parent:%p left:%p right:%p colour:%s key:%u %s:%p\n",
 687             prefix,
 688             (void *)node, (void *)node->parent,
 689             (void *)node->left, (void *)node->right,
 690             node->color?"Black":"Red", node->key32,
 691             node->is_subtree?"tree":"data", node->data);
 692     if (node->left)
 693         wmem_tree_print_nodes("L-", node->left, level+1);
 694     if (node->right)
 695         wmem_tree_print_nodes("R-", node->right, level+1);
 696
 697     if (node->is_subtree)
 698         wmem_print_subtree((wmem_tree_t *)node->data, level+1);
 699 }
 700
 701 static void
 702 wmem_print_subtree(wmem_tree_t *tree, guint32 level)
 703 {
 704     guint32 i;
 705
 706     if (!tree)
 707         return;
 708
 709     for (i=0; i<level; i++) {
 710         printf("    ");
 711     }
 712
 713     printf("WMEM tree:%p root:%p\n", (void *)tree, (void *)tree->root);
 714     if (tree->root) {
 715         wmem_tree_print_nodes("Root-", tree->root, level);
 716     }
 717 }
 718
 719 void
 720 wmem_print_tree(wmem_tree_t *tree)
 721 {
 722     wmem_print_subtree(tree, 0);
 723 }
 724
 725 /*
 726  * Editor modelines  -  http://www.wireshark.org/tools/modelines.html
 727  *
 728  * Local variables:
 729  * c-basic-offset: 4
 730  * tab-width: 8
 731  * indent-tabs-mode: nil
 732  * End:
 733  *
 734  * vi: set shiftwidth=4 tabstop=8 expandtab:
 735  * :indentSize=4:tabSize=8:noTabs=true:
 736  */