wsutil/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  * Wireshark - Network traffic analyzer
   7  * By Gerald Combs <gerald@wireshark.org>
   8  * Copyright 1998 Gerald Combs
   9  *
  10  * SPDX-License-Identifier: GPL-2.0-or-later
  11  */
  12
  13 #include "config.h"
  14
  15 #include <string.h>
  16 #include <stdio.h>
  17 #include <glib.h>
  18
  19 #include "wmem-int.h"
  20 #include "wmem_core.h"
  21 #include "wmem_strutl.h"
  22 #include "wmem_tree.h"
  23 #include "wmem_tree-int.h"
  24 #include "wmem_user_cb.h"
  25
  26 static wmem_tree_node_t *
  27 node_uncle(wmem_tree_node_t *node)
  28 {
  29     wmem_tree_node_t *parent, *grandparent;
  30
  31     parent = node->parent;
  32     if (parent == NULL) {
  33         return NULL;
  34     }
  35
  36     grandparent = parent->parent;
  37     if (grandparent == NULL) {
  38         return NULL;
  39     }
  40
  41     if (parent == grandparent->left) {
  42         return grandparent->right;
  43     }
  44     else {
  45         return grandparent->left;
  46     }
  47 }
  48
  49 static void rb_insert_case1(wmem_tree_t *tree, wmem_tree_node_t *node);
  50 static void rb_insert_case2(wmem_tree_t *tree, wmem_tree_node_t *node);
  51
  52 static void
  53 rotate_left(wmem_tree_t *tree, wmem_tree_node_t *node)
  54 {
  55     if (node->parent) {
  56         if (node->parent->left == node) {
  57             node->parent->left = node->right;
  58         }
  59         else {
  60             node->parent->right = node->right;
  61         }
  62     }
  63     else {
  64         tree->root = node->right;
  65     }
  66
  67     node->right->parent = node->parent;
  68     node->parent        = node->right;
  69     node->right         = node->right->left;
  70     if (node->right) {
  71         node->right->parent = node;
  72     }
  73     node->parent->left = node;
  74
  75     if (tree->post_rotation_cb) {
  76         tree->post_rotation_cb (node);
  77     }
  78 }
  79
  80 static void
  81 rotate_right(wmem_tree_t *tree, wmem_tree_node_t *node)
  82 {
  83     if (node->parent) {
  84         if (node->parent->left == node) {
  85             node->parent->left = node->left;
  86         }
  87         else {
  88             node->parent->right = node->left;
  89         }
  90     }
  91     else {
  92         tree->root = node->left;
  93     }
  94
  95     node->left->parent = node->parent;
  96     node->parent       = node->left;
  97     node->left         = node->left->right;
  98     if (node->left) {
  99         node->left->parent = node;
 100     }
 101     node->parent->right = node;
 102
 103
 104     if (tree->post_rotation_cb) {
 105         tree->post_rotation_cb (node);
 106     }
 107 }
 108
 109 static void
 110 rb_insert_case5(wmem_tree_t *tree, wmem_tree_node_t *node)
 111 {
 112     wmem_tree_node_t *parent, *grandparent;
 113
 114     parent      = node->parent;
 115     grandparent = parent->parent;
 116
 117     parent->color      = WMEM_NODE_COLOR_BLACK;
 118     grandparent->color = WMEM_NODE_COLOR_RED;
 119
 120     if (node == parent->left && parent == grandparent->left) {
 121         rotate_right(tree, grandparent);
 122     }
 123     else {
 124         rotate_left(tree, grandparent);
 125     }
 126 }
 127
 128 static void
 129 rb_insert_case4(wmem_tree_t *tree, wmem_tree_node_t *node)
 130 {
 131     wmem_tree_node_t *parent, *grandparent;
 132
 133     parent      = node->parent;
 134     grandparent = parent->parent;
 135     if (!grandparent) {
 136         return;
 137     }
 138
 139     if (node == parent->right && parent == grandparent->left) {
 140         rotate_left(tree, parent);
 141         node = node->left;
 142     }
 143     else if (node == parent->left && parent == grandparent->right) {
 144         rotate_right(tree, parent);
 145         node = node->right;
 146     }
 147
 148     rb_insert_case5(tree, node);
 149 }
 150
 151 static void
 152 rb_insert_case3(wmem_tree_t *tree, wmem_tree_node_t *node)
 153 {
 154     wmem_tree_node_t *parent, *grandparent, *uncle;
 155
 156     uncle = node_uncle(node);
 157
 158     if (uncle && uncle->color == WMEM_NODE_COLOR_RED) {
 159         parent      = node->parent;
 160         grandparent = parent->parent;
 161
 162         parent->color      = WMEM_NODE_COLOR_BLACK;
 163         uncle->color       = WMEM_NODE_COLOR_BLACK;
 164         grandparent->color = WMEM_NODE_COLOR_RED;
 165
 166         rb_insert_case1(tree, grandparent);
 167     }
 168     else {
 169         rb_insert_case4(tree, node);
 170     }
 171 }
 172
 173 static void
 174 rb_insert_case2(wmem_tree_t *tree, wmem_tree_node_t *node)
 175 {
 176     /* parent is always non-NULL here */
 177     if (node->parent->color == WMEM_NODE_COLOR_RED) {
 178         rb_insert_case3(tree, node);
 179     }
 180 }
 181
 182 static void
 183 rb_insert_case1(wmem_tree_t *tree, wmem_tree_node_t *node)
 184 {
 185     wmem_tree_node_t *parent = node->parent;
 186
 187     if (parent == NULL) {
 188         node->color = WMEM_NODE_COLOR_BLACK;
 189     }
 190     else {
 191         rb_insert_case2(tree, node);
 192     }
 193 }
 194
 195 static void
 196 rb_remove_doubleblack(wmem_tree_t *tree, wmem_tree_node_t *node)
 197 {
 198     wmem_tree_node_t *parent = node->parent;
 199     wmem_tree_node_t *sib, *c_nephew, *d_nephew;
 200     bool left;
 201
 202     ws_assert(parent);
 203     if (node == parent->left) {
 204         left = true;
 205         parent->left = NULL;
 206     } else {
 207         left = false;
 208         parent->right = NULL;
 209     }
 210
 211     for (node = NULL; parent; parent = node->parent) {
 212         if (node) {
 213             left = (node == parent->left);
 214         }
 215         if (left) {
 216             sib = parent->right;
 217             c_nephew = sib->left;
 218             d_nephew = sib->right;
 219         } else {
 220             sib = parent->left;
 221             c_nephew = sib->right;
 222             d_nephew = sib->left;
 223         }
 224         if (sib && sib->color == WMEM_NODE_COLOR_RED) {
 225             // Sib red
 226             goto case_d3;
 227         } else if (d_nephew && d_nephew->color == WMEM_NODE_COLOR_RED) {
 228             // D red, Sib black
 229             goto case_d6;
 230         } else if (c_nephew && c_nephew->color == WMEM_NODE_COLOR_RED) {
 231             // C red, Sib, D black
 232             goto case_d5;
 233         } else if (parent->color == WMEM_NODE_COLOR_RED) {
 234             // Parent red, Sib, D, C black
 235             goto case_d4;
 236         } else {
 237             // All black
 238             sib->color = WMEM_NODE_COLOR_RED;
 239             node = parent;
 240         }
 241     }
 242
 243     return;
 244
 245 case_d3:
 246     if (left) {
 247         rotate_left(tree, parent);
 248     } else {
 249         rotate_right(tree, parent);
 250     }
 251     parent->color = WMEM_NODE_COLOR_RED;
 252     sib->color = WMEM_NODE_COLOR_BLACK;
 253     sib = c_nephew;
 254     d_nephew = left ? sib->right : sib->left;
 255     if (d_nephew && d_nephew->color == WMEM_NODE_COLOR_RED)
 256         goto case_d6;
 257     c_nephew = left ? sib->left : sib->right;
 258     if (c_nephew && c_nephew->color == WMEM_NODE_COLOR_RED)
 259         goto case_d5;
 260
 261 case_d4:
 262     sib->color = WMEM_NODE_COLOR_RED;
 263     parent->color = WMEM_NODE_COLOR_BLACK;
 264     return;
 265
 266 case_d5:
 267     if (left) {
 268         rotate_right(tree, sib);
 269     } else {
 270         rotate_left(tree, sib);
 271     }
 272     sib->color = WMEM_NODE_COLOR_RED;
 273     c_nephew->color = WMEM_NODE_COLOR_BLACK;
 274     d_nephew = sib;
 275     sib = c_nephew;
 276     // D red and sib black;
 277
 278 case_d6:
 279     if (left) {
 280         rotate_left(tree, parent);
 281     } else {
 282         rotate_right(tree, parent);
 283     }
 284     sib->color = parent->color;
 285     parent->color = WMEM_NODE_COLOR_BLACK;
 286     d_nephew->color = WMEM_NODE_COLOR_BLACK;
 287     return;
 288 }
 289
 290 static void
 291 rb_remove_node(wmem_tree_t *tree, wmem_tree_node_t *node, bool free_key)
 292 {
 293     wmem_tree_node_t *temp_node;
 294     /* First, if the node has both children, swap the key and data
 295      * with the in-order successor and delete that node instead.
 296      */
 297     if (node->left && node->right) {
 298         temp_node = node->right;
 299         while (temp_node->left) {
 300             temp_node = temp_node->left;
 301         }
 302         node->key = temp_node->key;
 303         node->data = temp_node->data;
 304         node = temp_node;
 305     }
 306
 307     wmem_node_color_t child_color = WMEM_NODE_COLOR_BLACK;
 308     /* node now has one child at most. */
 309     if (node->left) {
 310         temp_node = node->left;
 311         ws_assert(node->right == NULL);
 312     } else {
 313         temp_node = node->right;
 314     }
 315     /* If there is a child, then the child must be red and original
 316      * node black, or else the R-B tree assumptions are wrong and
 317      * there's a problem elsewhere in the code. */
 318     if (temp_node) {
 319         child_color = temp_node->color;
 320         ws_assert(child_color == WMEM_NODE_COLOR_RED);
 321         ws_assert(node->color == WMEM_NODE_COLOR_BLACK);
 322         temp_node->parent = node->parent;
 323         temp_node->color = WMEM_NODE_COLOR_BLACK;
 324     }
 325
 326     if (temp_node == NULL &&
 327         node->color == WMEM_NODE_COLOR_BLACK && node->parent) {
 328
 329         /* Removing will create a "double black" imbalance in the tree and
 330          * we will need to rectify it to keep this a R-B tree.
 331          * This function removes and does any necessary rotations.
 332          */
 333         rb_remove_doubleblack(tree, node);
 334     } else {
 335         // Now remove the node from the tree.
 336         if (node->parent) {
 337             if (node == node->parent->left) {
 338                 node->parent->left = temp_node;
 339             } else {
 340                 node->parent->right = temp_node;
 341             }
 342         } else {
 343             tree->root = temp_node;
 344         }
 345     }
 346
 347     /* Freeing memory is only strictly necessary for a NULL allocator.
 348      * The key is copied in a string tree, a GUINT_TO_POINTER in a
 349      * 32 bit integer tree, and used uncopied in a generic tree, so
 350      * it should be freed in the first case but not the others.
 351      * The value is returned, so not freed under any circumstance.
 352      */
 353     if (free_key) {
 354         wmem_free(tree->data_allocator, (void *)node->key);
 355     }
 356     wmem_free(tree->data_allocator, node);
 357 }
 358
 359 wmem_tree_t *
 360 wmem_tree_new(wmem_allocator_t *allocator)
 361 {
 362     wmem_tree_t *tree;
 363
 364     tree = wmem_new0(allocator, wmem_tree_t);
 365     tree->metadata_allocator    = allocator;
 366     tree->data_allocator = allocator;
 367
 368     return tree;
 369 }
 370
 371 static bool
 372 wmem_tree_reset_cb(wmem_allocator_t *allocator _U_, wmem_cb_event_t event,
 373         void *user_data)
 374 {
 375     wmem_tree_t *tree = (wmem_tree_t *)user_data;
 376
 377     tree->root = NULL;
 378
 379     if (event == WMEM_CB_DESTROY_EVENT) {
 380         wmem_unregister_callback(tree->metadata_allocator, tree->metadata_scope_cb_id);
 381         wmem_free(tree->metadata_allocator, tree);
 382     }
 383
 384     return true;
 385 }
 386
 387 static bool
 388 wmem_tree_destroy_cb(wmem_allocator_t *allocator _U_, wmem_cb_event_t event _U_,
 389         void *user_data)
 390 {
 391     wmem_tree_t *tree = (wmem_tree_t *)user_data;
 392
 393     wmem_unregister_callback(tree->data_allocator, tree->data_scope_cb_id);
 394
 395     return false;
 396 }
 397
 398 wmem_tree_t *
 399 wmem_tree_new_autoreset(wmem_allocator_t *metadata_scope, wmem_allocator_t *data_scope)
 400 {
 401     wmem_tree_t *tree;
 402
 403     tree = wmem_new0(metadata_scope, wmem_tree_t);
 404     tree->metadata_allocator = metadata_scope;
 405     tree->data_allocator = data_scope;
 406
 407     tree->metadata_scope_cb_id = wmem_register_callback(metadata_scope, wmem_tree_destroy_cb,
 408             tree);
 409     tree->data_scope_cb_id  = wmem_register_callback(data_scope, wmem_tree_reset_cb,
 410             tree);
 411
 412     return tree;
 413 }
 414
 415 static void
 416 free_tree_node(wmem_allocator_t *allocator, wmem_tree_node_t* node, bool free_keys, bool free_values)
 417 {
 418     if (node == NULL) {
 419         return;
 420     }
 421
 422     if (node->left) {
 423         free_tree_node(allocator, node->left, free_keys, free_values);
 424     }
 425
 426     if (node->is_subtree) {
 427         wmem_tree_destroy((wmem_tree_t *)node->data, free_keys, free_values);
 428         node->data = NULL;
 429     }
 430
 431     if (node->right) {
 432         free_tree_node(allocator, node->right, free_keys, free_values);
 433     }
 434
 435     if (free_keys) {
 436         wmem_free(allocator, (void*)node->key);
 437     }
 438
 439     if (free_values) {
 440         wmem_free(allocator, node->data);
 441     }
 442     wmem_free(allocator, node);
 443 }
 444
 445 void
 446 wmem_tree_destroy(wmem_tree_t *tree, bool free_keys, bool free_values)
 447 {
 448     free_tree_node(tree->data_allocator, tree->root, free_keys, free_values);
 449     if (tree->metadata_allocator) {
 450         wmem_unregister_callback(tree->metadata_allocator, tree->metadata_scope_cb_id);
 451     }
 452     if (tree->data_allocator) {
 453         wmem_unregister_callback(tree->data_allocator, tree->data_scope_cb_id);
 454     }
 455     wmem_free(tree->metadata_allocator, tree);
 456 }
 457
 458 bool
 459 wmem_tree_is_empty(wmem_tree_t *tree)
 460 {
 461     return tree->root == NULL;
 462 }
 463
 464 static bool
 465 count_nodes(const void *key _U_, void *value _U_, void *userdata)
 466 {
 467     unsigned* count = (unsigned*)userdata;
 468     (*count)++;
 469     return false;
 470 }
 471
 472 unsigned
 473 wmem_tree_count(wmem_tree_t* tree)
 474 {
 475     unsigned count = 0;
 476
 477     /* Recursing through the tree counting each node is the simplest approach.
 478        We don't keep track of the count within the tree because it can get
 479        complicated with subtrees within the tree */
 480     wmem_tree_foreach(tree, count_nodes, &count);
 481
 482     return count;
 483 }
 484
 485 static wmem_tree_node_t *
 486 create_node(wmem_allocator_t *allocator, wmem_tree_node_t *parent, const void *key,
 487         void *data, wmem_node_color_t color, bool is_subtree)
 488 {
 489     wmem_tree_node_t *node;
 490
 491     node = wmem_new(allocator, wmem_tree_node_t);
 492
 493     node->left   = NULL;
 494     node->right  = NULL;
 495     node->parent = parent;
 496
 497     node->key  = key;
 498     node->data = data;
 499
 500     node->color      = color;
 501     node->is_subtree = is_subtree;
 502     node->is_removed = false;
 503
 504     return node;
 505 }
 506
 507 #define CREATE_DATA(TRANSFORM, DATA) ((TRANSFORM) ? (TRANSFORM)(DATA) : (DATA))
 508
 509
 510 /**
 511  * return inserted node
 512  */
 513 static wmem_tree_node_t *
 514 lookup_or_insert32_node(wmem_tree_t *tree, uint32_t key,
 515         void*(*func)(void*), void* data, bool is_subtree, bool replace)
 516 {
 517     wmem_tree_node_t *node     = tree->root;
 518     wmem_tree_node_t *new_node = NULL;
 519
 520     /* is this the first node ?*/
 521     if (!node) {
 522         new_node = create_node(tree->data_allocator, NULL, GUINT_TO_POINTER(key),
 523                 CREATE_DATA(func, data), WMEM_NODE_COLOR_BLACK, is_subtree);
 524         tree->root = new_node;
 525         return new_node;
 526     }
 527
 528     /* it was not the new root so walk the tree until we find where to
 529      * insert this new leaf.
 530      */
 531     while (!new_node) {
 532         /* this node already exists, so just return the data pointer*/
 533         if (key == GPOINTER_TO_UINT(node->key)) {
 534             if (replace) {
 535                 node->data = CREATE_DATA(func, data);
 536             }
 537             return node;
 538         }
 539         else if (key < GPOINTER_TO_UINT(node->key)) {
 540             if (node->left) {
 541                 node = node->left;
 542             }
 543             else {
 544                 /* new node to the left */
 545                 new_node = create_node(tree->data_allocator, node, GUINT_TO_POINTER(key),
 546                         CREATE_DATA(func, data), WMEM_NODE_COLOR_RED,
 547                         is_subtree);
 548                 node->left = new_node;
 549             }
 550         }
 551         else if (key > GPOINTER_TO_UINT(node->key)) {
 552             if (node->right) {
 553                 node = node->right;
 554             }
 555             else {
 556                 /* new node to the right */
 557                 new_node = create_node(tree->data_allocator, node, GUINT_TO_POINTER(key),
 558                         CREATE_DATA(func, data), WMEM_NODE_COLOR_RED,
 559                         is_subtree);
 560                 node->right = new_node;
 561             }
 562         }
 563     }
 564
 565     /* node will now point to the newly created node */
 566     rb_insert_case1(tree, new_node);
 567
 568     return new_node;
 569 }
 570
 571
 572 static void *
 573 lookup_or_insert32(wmem_tree_t *tree, uint32_t key,
 574         void*(*func)(void*), void* data, bool is_subtree, bool replace)
 575 {
 576     wmem_tree_node_t *node = lookup_or_insert32_node(tree, key, func, data, is_subtree, replace);
 577     return node->data;
 578 }
 579
 580 static void *
 581 wmem_tree_lookup(wmem_tree_t *tree, const void *key, compare_func cmp)
 582 {
 583     wmem_tree_node_t *node;
 584
 585     if (tree == NULL || key == NULL) {
 586         return NULL;
 587     }
 588
 589     node = tree->root;
 590
 591     while (node) {
 592         int result = cmp(key, node->key);
 593         if (result == 0) {
 594             return node->data;
 595         }
 596         else if (result < 0) {
 597             node = node->left;
 598         }
 599         else if (result > 0) {
 600             node = node->right;
 601         }
 602     }
 603
 604     return NULL;
 605 }
 606
 607 wmem_tree_node_t *
 608 wmem_tree_insert_node(wmem_tree_t *tree, const void *key, void *data, compare_func cmp)
 609 {
 610     wmem_tree_node_t *node = tree->root;
 611     wmem_tree_node_t *new_node = NULL;
 612
 613     /* is this the first node ?*/
 614     if (!node) {
 615         tree->root = create_node(tree->data_allocator, node, key,
 616                 data, WMEM_NODE_COLOR_BLACK, false);
 617         return tree->root;
 618     }
 619
 620     /* it was not the new root so walk the tree until we find where to
 621      * insert this new leaf.
 622      */
 623     while (!new_node) {
 624         int result = cmp(key, node->key);
 625         if (result == 0) {
 626             node->data = data;
 627             node->is_removed = data ? false : true;
 628             return node;
 629         }
 630         else if (result < 0) {
 631             if (node->left) {
 632                 node = node->left;
 633             }
 634             else {
 635                 new_node = create_node(tree->data_allocator, node, key,
 636                         data, WMEM_NODE_COLOR_RED, false);
 637                 node->left = new_node;
 638             }
 639         }
 640         else if (result > 0) {
 641             if (node->right) {
 642                 node = node->right;
 643             }
 644             else {
 645                 /* new node to the right */
 646                 new_node = create_node(tree->data_allocator, node, key,
 647                         data, WMEM_NODE_COLOR_RED, false);
 648                 node->right = new_node;
 649             }
 650         }
 651     }
 652
 653     /* node will now point to the newly created node */
 654     rb_insert_case1(tree, new_node);
 655
 656     return new_node;
 657 }
 658
 659 void
 660 wmem_tree_insert32(wmem_tree_t *tree, uint32_t key, void *data)
 661 {
 662     lookup_or_insert32(tree, key, NULL, data, false, true);
 663 }
 664
 665 bool wmem_tree_contains32(wmem_tree_t *tree, uint32_t key)
 666 {
 667     if (!tree) {
 668         return false;
 669     }
 670
 671     wmem_tree_node_t *node = tree->root;
 672
 673     while (node) {
 674         if (key == GPOINTER_TO_UINT(node->key)) {
 675             return true;
 676         }
 677         else if (key < GPOINTER_TO_UINT(node->key)) {
 678             node = node->left;
 679         }
 680         else if (key > GPOINTER_TO_UINT(node->key)) {
 681             node = node->right;
 682         }
 683     }
 684
 685     return false;
 686 }
 687
 688 static wmem_tree_node_t*
 689 wmem_tree_lookup32_node(wmem_tree_t *tree, uint32_t key)
 690 {
 691     if (!tree) {
 692         return NULL;
 693     }
 694
 695     wmem_tree_node_t *node = tree->root;
 696
 697     while (node) {
 698         if (key == GPOINTER_TO_UINT(node->key)) {
 699             return node;
 700         }
 701         else if (key < GPOINTER_TO_UINT(node->key)) {
 702             node = node->left;
 703         }
 704         else if (key > GPOINTER_TO_UINT(node->key)) {
 705             node = node->right;
 706         }
 707     }
 708
 709     return NULL;
 710 }
 711
 712 void *
 713 wmem_tree_lookup32(wmem_tree_t *tree, uint32_t key)
 714 {
 715     wmem_tree_node_t *node = wmem_tree_lookup32_node(tree, key);
 716     if (node == NULL) {
 717         return NULL;
 718     }
 719     return node->data;
 720 }
 721
 722 static wmem_tree_node_t*
 723 wmem_tree_lookup32_le_node(wmem_tree_t *tree, uint32_t key)
 724 {
 725     if (!tree) {
 726         return NULL;
 727     }
 728
 729     wmem_tree_node_t *node = tree->root;
 730
 731     while (node) {
 732         if (key == GPOINTER_TO_UINT(node->key)) {
 733             return node;
 734         }
 735         else if (key < GPOINTER_TO_UINT(node->key)) {
 736             if (node->left == NULL) {
 737                 break;
 738             }
 739             node = node->left;
 740         }
 741         else if (key > GPOINTER_TO_UINT(node->key)) {
 742             if (node->right == NULL) {
 743                 break;
 744             }
 745             node = node->right;
 746         }
 747     }
 748
 749     if (!node) {
 750         return NULL;
 751     }
 752
 753     /* If we are still at the root of the tree this means that this node
 754      * is either smaller than the search key and then we return this
 755      * node or else there is no smaller key available and then
 756      * we return NULL.
 757      */
 758     if (node->parent == NULL) {
 759         if (key > GPOINTER_TO_UINT(node->key)) {
 760             return node;
 761         } else {
 762             return NULL;
 763         }
 764     }
 765
 766     if (GPOINTER_TO_UINT(node->key) <= key) {
 767         /* if our key is <= the search key, we have the right node */
 768         return node;
 769     }
 770     else if (node == node->parent->left) {
 771         /* our key is bigger than the search key and we're a left child,
 772          * we have to check if any of our ancestors are smaller. */
 773         while (node) {
 774             if (key > GPOINTER_TO_UINT(node->key)) {
 775                 return node;
 776             }
 777             node=node->parent;
 778         }
 779         return NULL;
 780     }
 781     else {
 782         /* our key is bigger than the search key and we're a right child,
 783          * our parent is the one we want */
 784         return node->parent;
 785     }
 786 }
 787
 788 void *
 789 wmem_tree_lookup32_le(wmem_tree_t *tree, uint32_t key)
 790 {
 791     wmem_tree_node_t *node = wmem_tree_lookup32_le_node(tree, key);
 792     if (node == NULL) {
 793         return NULL;
 794     }
 795
 796     return node->data;
 797 }
 798
 799 void *
 800 wmem_tree_lookup32_le_full(wmem_tree_t *tree, uint32_t key, uint32_t *orig_key)
 801 {
 802     wmem_tree_node_t *node = wmem_tree_lookup32_le_node(tree, key);
 803     if (node == NULL) {
 804         return NULL;
 805     }
 806
 807     *orig_key = GPOINTER_TO_UINT(node->key);
 808     return node->data;
 809 }
 810
 811 static wmem_tree_node_t*
 812 wmem_tree_lookup32_ge_node(wmem_tree_t *tree, uint32_t key)
 813 {
 814     if (!tree) {
 815         return NULL;
 816     }
 817
 818     wmem_tree_node_t *node = tree->root;
 819
 820     while (node) {
 821         if (key == GPOINTER_TO_UINT(node->key)) {
 822             return node;
 823         }
 824         else if (key < GPOINTER_TO_UINT(node->key)) {
 825             if (node->left == NULL) {
 826                 break;
 827             }
 828             node = node->left;
 829         }
 830         else if (key > GPOINTER_TO_UINT(node->key)) {
 831             if (node->right == NULL) {
 832                 break;
 833             }
 834             node = node->right;
 835         }
 836     }
 837
 838     if (!node) {
 839         return NULL;
 840     }
 841
 842     /* If we are still at the root of the tree this means that this node
 843      * is either greater than the search key and then we return this
 844      * node or else there is no greater key available and then
 845      * we return NULL.
 846      */
 847     if (node->parent == NULL) {
 848         if (key < GPOINTER_TO_UINT(node->key)) {
 849             return node;
 850         } else {
 851             return NULL;
 852         }
 853     }
 854
 855     if (GPOINTER_TO_UINT(node->key) >= key) {
 856         /* if our key is >= the search key, we have the right node */
 857         return node;
 858     }
 859     else if (node == node->parent->right) {
 860         /* our key is smaller than the search key and we're a right child,
 861          * we have to check if any of our ancestors are bigger. */
 862         while (node) {
 863             if (key < GPOINTER_TO_UINT(node->key)) {
 864                 return node;
 865             }
 866             node=node->parent;
 867         }
 868         return NULL;
 869     }
 870     else {
 871         /* our key is smaller than the search key and we're a left child,
 872          * our parent is the one we want */
 873         return node->parent;
 874     }
 875 }
 876
 877 void *
 878 wmem_tree_lookup32_ge(wmem_tree_t *tree, uint32_t key)
 879 {
 880     wmem_tree_node_t *node = wmem_tree_lookup32_ge_node(tree, key);
 881     if (node == NULL) {
 882         return NULL;
 883     }
 884
 885     return node->data;
 886 }
 887
 888 void *
 889 wmem_tree_lookup32_ge_full(wmem_tree_t *tree, uint32_t key, uint32_t *orig_key)
 890 {
 891     wmem_tree_node_t *node = wmem_tree_lookup32_ge_node(tree, key);
 892     if (node == NULL) {
 893         return NULL;
 894     }
 895
 896     *orig_key = GPOINTER_TO_UINT(node->key);
 897     return node->data;
 898 }
 899
 900 void *
 901 wmem_tree_remove32(wmem_tree_t *tree, uint32_t key)
 902 {
 903     wmem_tree_node_t *node = wmem_tree_lookup32_node(tree, key);
 904     if (node == NULL) {
 905         return NULL;
 906     }
 907
 908     void *ret = node->data;
 909
 910     /* Remove the node. Do not free the key, because it is a
 911      * GPOINTER_TO_UINT. The value we return.
 912      */
 913     rb_remove_node(tree, node, false);
 914
 915     return ret;
 916 }
 917
 918 void
 919 wmem_tree_insert_string(wmem_tree_t* tree, const char* k, void* v, uint32_t flags)
 920 {
 921     char *key;
 922     compare_func cmp;
 923
 924     key = wmem_strdup(tree->data_allocator, k);
 925
 926     if (flags & WMEM_TREE_STRING_NOCASE) {
 927         cmp = (compare_func)g_ascii_strcasecmp;
 928     } else {
 929         cmp = (compare_func)strcmp;
 930     }
 931
 932     wmem_tree_insert_node(tree, key, v, cmp);
 933 }
 934
 935 void *
 936 wmem_tree_lookup_string(wmem_tree_t* tree, const char* k, uint32_t flags)
 937 {
 938     compare_func cmp;
 939
 940     if (flags & WMEM_TREE_STRING_NOCASE) {
 941         cmp = (compare_func)g_ascii_strcasecmp;
 942     } else {
 943         cmp = (compare_func)strcmp;
 944     }
 945
 946     return wmem_tree_lookup(tree, k, cmp);
 947 }
 948
 949 void *
 950 wmem_tree_remove_string(wmem_tree_t* tree, const char* k, uint32_t flags)
 951 {
 952     void *ret = wmem_tree_lookup_string(tree, k, flags);
 953     if (ret) {
 954         /* Not really a remove, but set data to NULL to mark node with is_removed */
 955         wmem_tree_insert_string(tree, k, NULL, flags);
 956     }
 957     return ret;
 958 }
 959
 960 static void *
 961 create_sub_tree(void* d)
 962 {
 963     return wmem_tree_new(((wmem_tree_t *)d)->data_allocator);
 964 }
 965
 966 void
 967 wmem_tree_insert32_array(wmem_tree_t *tree, wmem_tree_key_t *key, void *data)
 968 {
 969     wmem_tree_t *insert_tree = NULL;
 970     wmem_tree_key_t *cur_key;
 971     uint32_t i, insert_key32 = 0;
 972
 973     for (cur_key = key; cur_key->length > 0; cur_key++) {
 974         for (i = 0; i < cur_key->length; i++) {
 975             /* Insert using the previous key32 */
 976             if (!insert_tree) {
 977                 insert_tree = tree;
 978             } else {
 979                 insert_tree = (wmem_tree_t *)lookup_or_insert32(insert_tree,
 980                         insert_key32, create_sub_tree, tree, true, false);
 981             }
 982             insert_key32 = cur_key->key[i];
 983         }
 984     }
 985
 986     ws_assert(insert_tree);
 987
 988     wmem_tree_insert32(insert_tree, insert_key32, data);
 989 }
 990
 991 static void *
 992 wmem_tree_lookup32_array_helper(wmem_tree_t *tree, wmem_tree_key_t *key,
 993         void*(*helper)(wmem_tree_t*, uint32_t))
 994 {
 995     wmem_tree_t *lookup_tree = NULL;
 996     wmem_tree_key_t *cur_key;
 997     uint32_t i, lookup_key32 = 0;
 998
 999     if (!tree || !key) {
1000         return NULL;
1001     }
1002
1003     for (cur_key = key; cur_key->length > 0; cur_key++) {
1004         for (i = 0; i < cur_key->length; i++) {
1005             /* Lookup using the previous key32 */
1006             if (!lookup_tree) {
1007                 lookup_tree = tree;
1008             }
1009             else {
1010                 lookup_tree =
1011                     (wmem_tree_t *)(*helper)(lookup_tree, lookup_key32);
1012                 if (!lookup_tree) {
1013                     return NULL;
1014                 }
1015             }
1016             lookup_key32 = cur_key->key[i];
1017         }
1018     }
1019
1020     /* Assert if we didn't get any valid keys */
1021     ws_assert(lookup_tree);
1022
1023     return (*helper)(lookup_tree, lookup_key32);
1024 }
1025
1026 void *
1027 wmem_tree_lookup32_array(wmem_tree_t *tree, wmem_tree_key_t *key)
1028 {
1029     return wmem_tree_lookup32_array_helper(tree, key, wmem_tree_lookup32);
1030 }
1031
1032 void *
1033 wmem_tree_lookup32_array_le(wmem_tree_t *tree, wmem_tree_key_t *key)
1034 {
1035     return wmem_tree_lookup32_array_helper(tree, key, wmem_tree_lookup32_le);
1036 }
1037
1038 static bool
1039 wmem_tree_foreach_nodes(wmem_tree_node_t* node, wmem_foreach_func callback,
1040         void *user_data)
1041 {
1042     bool stop_traverse = false;
1043
1044     if (!node) {
1045         return false;
1046     }
1047
1048     if (node->left) {
1049         if (wmem_tree_foreach_nodes(node->left, callback, user_data)) {
1050             return true;
1051         }
1052     }
1053
1054     if (node->is_subtree) {
1055         stop_traverse = wmem_tree_foreach((wmem_tree_t *)node->data,
1056                 callback, user_data);
1057     } else if (!node->is_removed) {
1058         /* No callback for "removed" nodes */
1059         stop_traverse = callback(node->key, node->data, user_data);
1060     }
1061
1062     if (stop_traverse) {
1063         return true;
1064     }
1065
1066     if(node->right) {
1067         if (wmem_tree_foreach_nodes(node->right, callback, user_data)) {
1068             return true;
1069         }
1070     }
1071
1072     return false;
1073 }
1074
1075 bool
1076 wmem_tree_foreach(wmem_tree_t* tree, wmem_foreach_func callback,
1077         void *user_data)
1078 {
1079     if(!tree->root)
1080         return false;
1081
1082     return wmem_tree_foreach_nodes(tree->root, callback, user_data);
1083 }
1084
1085 static void wmem_print_subtree(wmem_tree_t *tree, uint32_t level, wmem_printer_func key_printer, wmem_printer_func data_printer);
1086
1087 static void
1088 wmem_print_indent(uint32_t level) {
1089     uint32_t i;
1090     for (i=0; i<level; i++) {
1091         printf("    ");
1092     }
1093 }
1094
1095 static void
1096 wmem_tree_print_nodes(const char *prefix, wmem_tree_node_t *node, uint32_t level,
1097     wmem_printer_func key_printer, wmem_printer_func data_printer)
1098 {
1099     if (!node)
1100         return;
1101
1102     wmem_print_indent(level);
1103
1104     printf("%sNODE:%p parent:%p left:%p right:%p colour:%s key:%p %s:%p\n",
1105             prefix,
1106             (void *)node, (void *)node->parent,
1107             (void *)node->left, (void *)node->right,
1108             node->color?"Black":"Red", node->key,
1109             node->is_subtree?"tree":"data", node->data);
1110     if (key_printer) {
1111         wmem_print_indent(level);
1112         key_printer(node->key);
1113         printf("\n");
1114     }
1115     if (data_printer && !node->is_subtree) {
1116         wmem_print_indent(level);
1117         data_printer(node->data);
1118         printf("\n");
1119     }
1120
1121     if (node->left)
1122         wmem_tree_print_nodes("L-", node->left, level+1, key_printer, data_printer);
1123     if (node->right)
1124         wmem_tree_print_nodes("R-", node->right, level+1, key_printer, data_printer);
1125
1126     if (node->is_subtree)
1127         wmem_print_subtree((wmem_tree_t *)node->data, level+1, key_printer, data_printer);
1128 }
1129
1130
1131 static void
1132 wmem_print_subtree(wmem_tree_t *tree, uint32_t level, wmem_printer_func key_printer, wmem_printer_func data_printer)
1133 {
1134     if (!tree)
1135         return;
1136
1137     wmem_print_indent(level);
1138
1139     printf("WMEM tree:%p root:%p\n", (void *)tree, (void *)tree->root);
1140     if (tree->root) {
1141         wmem_tree_print_nodes("Root-", tree->root, level, key_printer, data_printer);
1142     }
1143 }
1144
1145 void
1146 wmem_print_tree(wmem_tree_t *tree, wmem_printer_func key_printer, wmem_printer_func data_printer)
1147 {
1148     wmem_print_subtree(tree, 0, key_printer, data_printer);
1149 }
1150 /*
1151  * Editor modelines  -  https://www.wireshark.org/tools/modelines.html
1152  *
1153  * Local variables:
1154  * c-basic-offset: 4
1155  * tab-width: 8
1156  * indent-tabs-mode: nil
1157  * End:
1158  *
1159  * vi: set shiftwidth=4 tabstop=8 expandtab:
1160  * :indentSize=4:tabSize=8:noTabs=true:
1161  */