mcs/class/System/System.Collections.Generic/RBTree.cs

   1 //
   2 // System.Collections.Generic.RBTree
   3 //
   4 // Authors:
   5 //   Raja R Harinath <rharinath@novell.com>
   6 //
   7
   8 //
   9 // Copyright (C) 2007, Novell, Inc.
  10 //
  11 // Permission is hereby granted, free of charge, to any person obtaining
  12 // a copy of this software and associated documentation files (the
  13 // "Software"), to deal in the Software without restriction, including
  14 // without limitation the rights to use, copy, modify, merge, publish,
  15 // distribute, sublicense, and/or sell copies of the Software, and to
  16 // permit persons to whom the Software is furnished to do so, subject to
  17 // the following conditions:
  18 //
  19 // The above copyright notice and this permission notice shall be
  20 // included in all copies or substantial portions of the Software.
  21 //
  22 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  23 // EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  24 // MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  25 // NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
  26 // LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
  27 // OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
  28 // WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  29 //
  30
  31 #define ONE_MEMBER_CACHE
  32
  33 #if NET_2_0
  34 using System;
  35 using System.Collections;
  36
  37 namespace System.Collections.Generic
  38 {
  39         [Serializable]
  40         internal class RBTree : IEnumerable, IEnumerable<RBTree.Node> {
  41                 public interface INodeHelper<T> {
  42                         int Compare (T key, Node node);
  43                         Node CreateNode (T key);
  44                 }
  45
  46                 public abstract class Node {
  47                         public Node left, right;
  48                         uint size_black;
  49
  50                         const uint black_mask = 1;
  51                         const int black_shift = 1;
  52                         public bool IsBlack {
  53                                 get { return (size_black & black_mask) == black_mask; }
  54                                 set { size_black = value ? (size_black | black_mask) : (size_black & ~black_mask); }
  55                         }
  56
  57                         public uint Size {
  58                                 get { return size_black >> black_shift; }
  59                                 set { size_black = (value << black_shift) | (size_black & black_mask); }
  60                         }
  61
  62                         public uint FixSize ()
  63                         {
  64                                 Size = 1;
  65                                 if (left != null)
  66                                         Size += left.Size;
  67                                 if (right != null)
  68                                         Size += right.Size;
  69                                 return Size;
  70                         }
  71
  72                         public Node ()
  73                         {
  74                                 size_black = 2; // Size == 1, IsBlack = false
  75                         }
  76
  77                         public abstract void SwapValue (Node other);
  78
  79 #if TEST
  80                         public int VerifyInvariants ()
  81                         {
  82                                 int black_depth_l = 0;
  83                                 int black_depth_r = 0;
  84                                 uint size = 1;
  85                                 bool child_is_red = false;
  86                                 if (left != null) {
  87                                         black_depth_l = left.VerifyInvariants ();
  88                                         size += left.Size;
  89                                         child_is_red |= !left.IsBlack;
  90                                 }
  91
  92                                 if (right != null) {
  93                                         black_depth_r = right.VerifyInvariants ();
  94                                         size += right.Size;
  95                                         child_is_red |= !right.IsBlack;
  96                                 }
  97
  98                                 if (black_depth_l != black_depth_r)
  99                                         throw new SystemException ("Internal error: black depth mismatch");
 100
 101                                 if (!IsBlack && child_is_red)
 102                                         throw new SystemException ("Internal error: red-red conflict");
 103                                 if (Size != size)
 104                                         throw new SystemException ("Internal error: metadata error");
 105
 106                                 return black_depth_l + (IsBlack ? 1 : 0);
 107                         }
 108
 109                         public abstract void Dump (string indent);
 110 #endif
 111                 }
 112
 113                 Node root;
 114                 object hlp;
 115                 uint version;
 116
 117 #if ONE_MEMBER_CACHE
 118 #if TARGET_JVM
 119                 static readonly LocalDataStoreSlot _cachedPathStore = System.Threading.Thread.AllocateDataSlot ();
 120
 121                 static List<Node> cached_path {
 122                         get { return (List<Node>) System.Threading.Thread.GetData (_cachedPathStore); }
 123                         set { System.Threading.Thread.SetData (_cachedPathStore, value); }
 124                 }
 125 #else
 126                 [ThreadStatic]
 127                 static List<Node> cached_path;
 128 #endif
 129
 130                 static List<Node> alloc_path ()
 131                 {
 132                         if (cached_path == null)
 133                                 return new List<Node> ();
 134
 135                         List<Node> path = cached_path;
 136                         cached_path = null;
 137                         return path;
 138                 }
 139
 140                 static void release_path (List<Node> path)
 141                 {
 142                         if (cached_path == null || cached_path.Capacity < path.Capacity) {
 143                                 path.Clear ();
 144                                 cached_path = path;
 145                         }
 146                 }
 147 #else
 148                 static List<Node> alloc_path ()
 149                 {
 150                         return new List<Node> ();
 151                 }
 152
 153                 static void release_path (List<Node> path)
 154                 {
 155                 }
 156 #endif
 157
 158                 public RBTree (object hlp)
 159                 {
 160                         // hlp is INodeHelper<T> for some T
 161                         this.hlp = hlp;
 162                 }
 163
 164                 public void Clear ()
 165                 {
 166                         root = null;
 167                         ++version;
 168                 }
 169
 170                 // if key is already in the tree, return the node associated with it
 171                 // if not, insert new_node into the tree, and return it
 172                 public Node Intern<T> (T key, Node new_node)
 173                 {
 174                         if (root == null) {
 175                                 if (new_node == null)
 176                                         new_node = ((INodeHelper<T>) hlp).CreateNode (key);
 177                                 root = new_node;
 178                                 root.IsBlack = true;
 179                                 ++version;
 180                                 return root;
 181                         }
 182
 183                         List<Node> path = alloc_path ();
 184                         int in_tree_cmp = find_key (key, path);
 185                         Node retval = path [path.Count - 1];
 186                         if (retval == null) {
 187                                 if (new_node == null)
 188                                         new_node = ((INodeHelper<T>) hlp).CreateNode (key);
 189                                 retval = do_insert (in_tree_cmp, new_node, path);
 190                         }
 191                         // no need for a try .. finally, this is only used to mitigate allocations
 192                         release_path (path);
 193                         return retval;
 194                 }
 195
 196                 // returns the just-removed node (or null if the value wasn't in the tree)
 197                 public Node Remove<T> (T key)
 198                 {
 199                         if (root == null)
 200                                 return null;
 201
 202                         List<Node> path = alloc_path ();
 203                         int in_tree_cmp = find_key (key, path);
 204                         Node retval = null;
 205                         if (in_tree_cmp == 0)
 206                                 retval = do_remove (path);
 207                         // no need for a try .. finally, this is only used to mitigate allocations
 208                         release_path (path);
 209                         return retval;
 210                 }
 211
 212                 public Node Lookup<T> (T key)
 213                 {
 214                         INodeHelper<T> hlp = (INodeHelper<T>) this.hlp;
 215                         Node current = root;
 216                         while (current != null) {
 217                                 int c = hlp.Compare (key, current);
 218                                 if (c == 0)
 219                                         break;
 220                                 current = c < 0 ? current.left : current.right;
 221                         }
 222                         return current;
 223                 }
 224
 225                 public int Count {
 226                         get { return root == null ? 0 : (int) root.Size; }
 227                 }
 228
 229                 public Node this [int index] {
 230                         get {
 231                                 if (index < 0 || index >= Count)
 232                                         throw new IndexOutOfRangeException ("index");
 233
 234                                 Node current = root;
 235                                 while (current != null) {
 236                                         int left_size = current.left == null ? 0 : (int) current.left.Size;
 237                                         if (index == left_size)
 238                                                 return current;
 239                                         if (index < left_size) {
 240                                                 current = current.left;
 241                                         } else {
 242                                                 index -= left_size + 1;
 243                                                 current = current.right;
 244                                         }
 245                                 }
 246                                 throw new SystemException ("Internal Error: index calculation");
 247                         }
 248                 }
 249
 250                 public NodeEnumerator GetEnumerator ()
 251                 {
 252                         return new NodeEnumerator (this);
 253                 }
 254
 255                 IEnumerator<Node> IEnumerable<Node>.GetEnumerator ()
 256                 {
 257                         return GetEnumerator ();
 258                 }
 259
 260                 IEnumerator IEnumerable.GetEnumerator ()
 261                 {
 262                         return GetEnumerator ();
 263                 }
 264
 265 #if TEST
 266                 public void VerifyInvariants ()
 267                 {
 268                         if (root != null) {
 269                                 if (!root.IsBlack)
 270                                         throw new SystemException ("Internal Error: root is not black");
 271                                 root.VerifyInvariants ();
 272                         }
 273                 }
 274
 275                 public void Dump ()
 276                 {
 277                         if (root != null)
 278                                 root.Dump ("");
 279                 }
 280 #endif
 281
 282                 // Pre-condition: root != null
 283                 int find_key<T> (T key, List<Node> path)
 284                 {
 285                         INodeHelper<T> hlp = (INodeHelper<T>) this.hlp;
 286                         int c = 0;
 287                         Node sibling = null;
 288                         Node current = root;
 289
 290                         if (path != null)
 291                                 path.Add (root);
 292
 293                         while (current != null) {
 294                                 c = hlp.Compare (key, current);
 295                                 if (c == 0)
 296                                         return c;
 297
 298                                 if (c < 0) {
 299                                         sibling = current.right;
 300                                         current = current.left;
 301                                 } else {
 302                                         sibling = current.left;
 303                                         current = current.right;
 304                                 }
 305
 306                                 if (path != null) {
 307                                         path.Add (sibling);
 308                                         path.Add (current);
 309                                 }
 310                         }
 311
 312                         return c;
 313                 }
 314
 315                 Node do_insert (int in_tree_cmp, Node current, List<Node> path)
 316                 {
 317                         path [path.Count - 1] = current;
 318                         Node parent = path [path.Count - 3];
 319
 320                         if (in_tree_cmp < 0)
 321                                 parent.left = current;
 322                         else
 323                                 parent.right = current;
 324                         for (int i = 0; i < path.Count - 2; i += 2)
 325                                 ++ path [i].Size;
 326
 327                         if (!parent.IsBlack)
 328                                 rebalance_insert (path);
 329
 330                         if (!root.IsBlack)
 331                                 throw new SystemException ("Internal error: root is not black");
 332
 333                         ++version;
 334                         return current;
 335                 }
 336
 337                 Node do_remove (List<Node> path)
 338                 {
 339                         int curpos = path.Count - 1;
 340
 341                         Node current = path [curpos];
 342                         if (current.left != null) {
 343                                 Node pred = right_most (current.left, current.right, path);
 344                                 current.SwapValue (pred);
 345                                 if (pred.left != null) {
 346                                         Node ppred = pred.left;
 347                                         path.Add (null); path.Add (ppred);
 348                                         pred.SwapValue (ppred);
 349                                 }
 350                         } else if (current.right != null) {
 351                                 Node succ = current.right;
 352                                 path.Add (null); path.Add (succ);
 353                                 current.SwapValue (succ);
 354                         }
 355
 356                         curpos = path.Count - 1;
 357                         current = path [curpos];
 358
 359                         if (current.Size != 1)
 360                                 throw new SystemException ("Internal Error: red-black violation somewhere");
 361
 362                         // remove it from our data structures
 363                         path [curpos] = null;
 364                         node_reparent (curpos == 0 ? null : path [curpos-2], current, 0, null);
 365
 366                         for (int i = 0; i < path.Count - 2; i += 2)
 367                                 -- path [i].Size;
 368
 369                         if (current.IsBlack) {
 370                                 current.IsBlack = false;
 371                                 if (curpos != 0)
 372                                         rebalance_delete (path);
 373                         }
 374
 375                         if (root != null && !root.IsBlack)
 376                                 throw new SystemException ("Internal Error: root is not black");
 377
 378                         ++version;
 379                         return current;
 380                 }
 381
 382                 // Pre-condition: current is red
 383                 void rebalance_insert (List<Node> path)
 384                 {
 385                         int curpos = path.Count - 1;
 386                         do {
 387                                 // parent == curpos-2, uncle == curpos-3, grandpa == curpos-4
 388                                 if (path [curpos-3] == null || path [curpos-3].IsBlack) {
 389                                         rebalance_insert__rotate_final (curpos, path);
 390                                         return;
 391                                 }
 392
 393                                 path [curpos-2].IsBlack = path [curpos-3].IsBlack = true;
 394
 395                                 curpos -= 4; // move to the grandpa
 396
 397                                 if (curpos == 0) // => current == root
 398                                         return;
 399                                 path [curpos].IsBlack = false;
 400                         } while (!path [curpos-2].IsBlack);
 401                 }
 402
 403                 // Pre-condition: current is black
 404                 void rebalance_delete (List<Node> path)
 405                 {
 406                         int curpos = path.Count - 1;
 407                         do {
 408                                 Node sibling = path [curpos-1];
 409                                 // current is black => sibling != null
 410                                 if (!sibling.IsBlack) {
 411                                         // current is black && sibling is red
 412                                         // => both sibling.left and sibling.right are black, and are not null
 413                                         curpos = ensure_sibling_black (curpos, path);
 414                                         // one of the nephews became the new sibling -- in either case, sibling != null
 415                                         sibling = path [curpos-1];
 416                                 }
 417
 418                                 if ((sibling.left != null && !sibling.left.IsBlack) ||
 419                                     (sibling.right != null && !sibling.right.IsBlack)) {
 420                                         rebalance_delete__rotate_final (curpos, path);
 421                                         return;
 422                                 }
 423
 424                                 sibling.IsBlack = false;
 425
 426                                 curpos -= 2; // move to the parent
 427
 428                                 if (curpos == 0)
 429                                         return;
 430                         } while (path [curpos].IsBlack);
 431                         path [curpos].IsBlack = true;
 432                 }
 433
 434                 void rebalance_insert__rotate_final (int curpos, List<Node> path)
 435                 {
 436                         Node current = path [curpos];
 437                         Node parent = path [curpos-2];
 438                         Node grandpa = path [curpos-4];
 439
 440                         uint grandpa_size = grandpa.Size;
 441
 442                         Node new_root;
 443
 444                         bool l1 = parent == grandpa.left;
 445                         bool l2 = current == parent.left;
 446                         if (l1 && l2) {
 447                                 grandpa.left = parent.right; parent.right = grandpa;
 448                                 new_root = parent;
 449                         } else if (l1 && !l2) {
 450                                 grandpa.left = current.right; current.right = grandpa;
 451                                 parent.right = current.left; current.left = parent;
 452                                 new_root = current;
 453                         } else if (!l1 && l2) {
 454                                 grandpa.right = current.left; current.left = grandpa;
 455                                 parent.left = current.right; current.right = parent;
 456                                 new_root = current;
 457                         } else { // (!l1 && !l2)
 458                                 grandpa.right = parent.left; parent.left = grandpa;
 459                                 new_root = parent;
 460                         }
 461
 462                         grandpa.FixSize (); grandpa.IsBlack = false;
 463                         if (new_root != parent)
 464                                 parent.FixSize (); /* parent is red already, so no need to set it */
 465
 466                         new_root.IsBlack = true;
 467                         node_reparent (curpos == 4 ? null : path [curpos-6], grandpa, grandpa_size, new_root);
 468                 }
 469
 470                 // Pre-condition: sibling is black, and one of sibling.left and sibling.right is red
 471                 void rebalance_delete__rotate_final (int curpos, List<Node> path)
 472                 {
 473                         //Node current = path [curpos];
 474                         Node sibling = path [curpos-1];
 475                         Node parent = path [curpos-2];
 476
 477                         uint parent_size = parent.Size;
 478                         bool parent_was_black = parent.IsBlack;
 479
 480                         Node new_root;
 481                         if (parent.right == sibling) {
 482                                 // if far nephew is black
 483                                 if (sibling.right == null || sibling.right.IsBlack) {
 484                                         // => near nephew is red, move it up
 485                                         Node nephew = sibling.left;
 486                                         parent.right = nephew.left; nephew.left = parent;
 487                                         sibling.left = nephew.right; nephew.right = sibling;
 488                                         new_root = nephew;
 489                                 } else {
 490                                         parent.right = sibling.left; sibling.left = parent;
 491                                         sibling.right.IsBlack = true;
 492                                         new_root = sibling;
 493                                 }
 494                         } else {
 495                                 // if far nephew is black
 496                                 if (sibling.left == null || sibling.left.IsBlack) {
 497                                         // => near nephew is red, move it up
 498                                         Node nephew = sibling.right;
 499                                         parent.left = nephew.right; nephew.right = parent;
 500                                         sibling.right = nephew.left; nephew.left = sibling;
 501                                         new_root = nephew;
 502                                 } else {
 503                                         parent.left = sibling.right; sibling.right = parent;
 504                                         sibling.left.IsBlack = true;
 505                                         new_root = sibling;
 506                                 }
 507                         }
 508
 509                         parent.FixSize (); parent.IsBlack = true;
 510                         if (new_root != sibling)
 511                                 sibling.FixSize (); /* sibling is already black, so no need to set it */
 512
 513                         new_root.IsBlack = parent_was_black;
 514                         node_reparent (curpos == 2 ? null : path [curpos-4], parent, parent_size, new_root);
 515                 }
 516
 517                 // Pre-condition: sibling is red (=> parent, sibling.left and sibling.right are black)
 518                 int ensure_sibling_black (int curpos, List<Node> path)
 519                 {
 520                         Node current = path [curpos];
 521                         Node sibling = path [curpos-1];
 522                         Node parent = path [curpos-2];
 523
 524                         bool current_on_left;
 525                         uint parent_size = parent.Size;
 526
 527                         if (parent.right == sibling) {
 528                                 parent.right = sibling.left; sibling.left = parent;
 529                                 current_on_left = true;
 530                         } else {
 531                                 parent.left = sibling.right; sibling.right = parent;
 532                                 current_on_left = false;
 533                         }
 534
 535                         parent.FixSize (); parent.IsBlack = false;
 536
 537                         sibling.IsBlack = true;
 538                         node_reparent (curpos == 2 ? null : path [curpos-4], parent, parent_size, sibling);
 539
 540                         // accomodate the rotation
 541                         if (curpos+1 == path.Count) {
 542                                 path.Add (null);
 543                                 path.Add (null);
 544                         }
 545
 546                         path [curpos-2] = sibling;
 547                         path [curpos-1] = current_on_left ? sibling.right : sibling.left;
 548                         path [curpos] = parent;
 549                         path [curpos+1] = current_on_left ? parent.right : parent.left;
 550                         path [curpos+2] = current;
 551
 552                         return curpos + 2;
 553                 }
 554
 555                 void node_reparent (Node orig_parent, Node orig, uint orig_size, Node updated)
 556                 {
 557                         if (updated != null && updated.FixSize () != orig_size)
 558                                 throw new SystemException ("Internal error: rotation");
 559
 560                         if (orig == root)
 561                                 root = updated;
 562                         else if (orig == orig_parent.left)
 563                                 orig_parent.left = updated;
 564                         else if (orig == orig_parent.right)
 565                                 orig_parent.right = updated;
 566                         else
 567                                 throw new SystemException ("Internal error: path error");
 568                 }
 569
 570                 // Pre-condition: current != null
 571                 static Node right_most (Node current, Node sibling, List<Node> path)
 572                 {
 573                         for (;;) {
 574                                 path.Add (sibling);
 575                                 path.Add (current);
 576                                 if (current.right == null)
 577                                         return current;
 578                                 sibling = current.left;
 579                                 current = current.right;
 580                         }
 581                 }
 582
 583                 [Serializable]
 584                 public struct NodeEnumerator : IEnumerator, IEnumerator<Node> {
 585                         RBTree tree;
 586                         uint version;
 587
 588                         Stack<Node> pennants;
 589
 590                         internal NodeEnumerator (RBTree tree)
 591                         {
 592                                 this.tree = tree;
 593                                 version = tree.version;
 594                                 pennants = null;
 595                         }
 596
 597                         public void Reset ()
 598                         {
 599                                 check_version ();
 600                                 pennants = null;
 601                         }
 602
 603                         public Node Current {
 604                                 get { return pennants.Peek (); }
 605                         }
 606
 607                         object IEnumerator.Current {
 608                                 get {
 609                                         check_current ();
 610                                         return Current;
 611                                 }
 612                         }
 613
 614                         public bool MoveNext ()
 615                         {
 616                                 check_version ();
 617
 618                                 Node next;
 619                                 if (pennants == null) {
 620                                         if (tree.root == null)
 621                                                 return false;
 622                                         pennants = new Stack<Node> ();
 623                                         next = tree.root;
 624                                 } else {
 625                                         if (pennants.Count == 0)
 626                                                 return false;
 627                                         Node current = pennants.Pop ();
 628                                         next = current.right;
 629                                 }
 630                                 for (; next != null; next = next.left)
 631                                         pennants.Push (next);
 632
 633                                 return pennants.Count != 0;
 634                         }
 635
 636                         public void Dispose ()
 637                         {
 638                                 tree = null;
 639                                 pennants = null;
 640                         }
 641
 642                         void check_version ()
 643                         {
 644                                 if (tree == null)
 645                                         throw new ObjectDisposedException ("enumerator");
 646                                 if (version != tree.version)
 647                                         throw new InvalidOperationException ("tree modified");
 648                         }
 649
 650                         internal void check_current ()
 651                         {
 652                                 check_version ();
 653                                 if (pennants == null)
 654                                         throw new InvalidOperationException ("state invalid before the first MoveNext()");
 655                         }
 656                 }
 657         }
 658 }
 659
 660 #if TEST
 661 namespace Mono.ValidationTest {
 662         using System.Collections.Generic;
 663
 664         internal class TreeSet<T> : IEnumerable<T>, IEnumerable
 665         {
 666                 public class Node : RBTree.Node {
 667                         public T value;
 668
 669                         public Node (T v)
 670                         {
 671                                 value = v;
 672                         }
 673
 674                         public override void SwapValue (RBTree.Node other)
 675                         {
 676                                 Node o = (Node) other;
 677                                 T v = value;
 678                                 value = o.value;
 679                                 o.value = v;
 680                         }
 681
 682                         public override void Dump (string indent)
 683                         {
 684                                 Console.WriteLine ("{0}{1} {2}({3})", indent, value, IsBlack ? "*" : "", Size);
 685                                 if (left != null)
 686                                         left.Dump (indent + "  /");
 687                                 if (right != null)
 688                                         right.Dump (indent + "  \\");
 689                         }
 690                 }
 691
 692                 public class NodeHelper : RBTree.INodeHelper<T> {
 693                         IComparer<T> cmp;
 694
 695                         public int Compare (T value, RBTree.Node node)
 696                         {
 697                                 return cmp.Compare (value, ((Node) node).value);
 698                         }
 699
 700                         public RBTree.Node CreateNode (T value)
 701                         {
 702                                 return new Node (value);
 703                         }
 704
 705                         private NodeHelper (IComparer<T> cmp)
 706                         {
 707                                 this.cmp = cmp;
 708                         }
 709                         static NodeHelper Default = new NodeHelper (Comparer<T>.Default);
 710                         public static NodeHelper GetHelper (IComparer<T> cmp)
 711                         {
 712                                 if (cmp == null || cmp == Comparer<T>.Default)
 713                                         return Default;
 714                                 return new NodeHelper (cmp);
 715                         }
 716                 }
 717
 718                 public struct Enumerator : IDisposable, IEnumerator, IEnumerator<T> {
 719                         RBTree.NodeEnumerator host;
 720
 721                         internal Enumerator (TreeSet<T> tree)
 722                         {
 723                                 host = new RBTree.NodeEnumerator (tree.tree);
 724                         }
 725
 726                         void IEnumerator.Reset ()
 727                         {
 728                                 host.Reset ();
 729                         }
 730
 731                         public T Current {
 732                                 get { return ((Node) host.Current).value; }
 733                         }
 734
 735                         object IEnumerator.Current {
 736                                 get { return Current; }
 737                         }
 738
 739                         public bool MoveNext ()
 740                         {
 741                                 return host.MoveNext ();
 742                         }
 743
 744                         public void Dispose ()
 745                         {
 746                                 host.Dispose ();
 747                         }
 748                 }
 749
 750                 RBTree tree;
 751
 752                 public TreeSet () : this (null)
 753                 {
 754                 }
 755
 756                 public TreeSet (IComparer<T> cmp)
 757                 {
 758                         tree = new RBTree (NodeHelper.GetHelper (cmp));
 759                 }
 760
 761                 IEnumerator IEnumerable.GetEnumerator ()
 762                 {
 763                         return GetEnumerator ();
 764                 }
 765
 766                 IEnumerator<T> IEnumerable<T>.GetEnumerator ()
 767                 {
 768                         return GetEnumerator ();
 769                 }
 770
 771                 public Enumerator GetEnumerator ()
 772                 {
 773                         return new Enumerator (this);
 774                 }
 775
 776                 // returns true if the value was inserted, false if the value already existed in the tree
 777                 public bool Insert (T value)
 778                 {
 779                         RBTree.Node n = new Node (value);
 780                         return tree.Intern (value, n) == n;
 781                 }
 782
 783                 // returns true if the value was removed, false if the value didn't exist in the tree
 784                 public bool Remove (T value)
 785                 {
 786                         return tree.Remove (value) != null;
 787                 }
 788
 789                 public bool Contains (T value)
 790                 {
 791                         return tree.Lookup (value) != null;
 792                 }
 793
 794                 public T this [int index] {
 795                         get { return ((Node) tree [index]).value; }
 796                 }
 797
 798                 public int Count {
 799                         get { return (int) tree.Count; }
 800                 }
 801
 802                 public void VerifyInvariants ()
 803                 {
 804                         tree.VerifyInvariants ();
 805                 }
 806
 807                 public void Dump ()
 808                 {
 809                         tree.Dump ();
 810                 }
 811         }
 812
 813         class Test {
 814                 static void Main (string [] args)
 815                 {
 816                         Random r = new Random ();
 817                         Dictionary<int, int> d = new Dictionary<int, int> ();
 818                         TreeSet<int> t = new TreeSet<int> ();
 819                         int iters = args.Length == 0 ? 100000 : Int32.Parse (args [0]);
 820                         int watermark = 1;
 821
 822                         for (int i = 0; i < iters; ++i) {
 823                                 if (i >= watermark) {
 824                                         watermark += 1 + watermark/4;
 825                                         t.VerifyInvariants ();
 826                                 }
 827
 828                                 int n = r.Next ();
 829                                 if (d.ContainsKey (n))
 830                                         continue;
 831                                 d [n] = n;
 832
 833                                 try {
 834                                         if (t.Contains (n))
 835                                                 throw new Exception ("tree says it has a number it shouldn't");
 836                                         if (!t.Insert (n))
 837                                                 throw new Exception ("tree says it has a number it shouldn't");
 838                                 } catch {
 839                                         Console.Error.WriteLine ("Exception while inserting {0} in iteration {1}", n, i);
 840                                         throw;
 841                                 }
 842                         }
 843                         t.VerifyInvariants ();
 844                         if (d.Count != t.Count)
 845                                 throw new Exception ("tree count is wrong?");
 846
 847                         Console.WriteLine ("Tree has {0} elements", t.Count);
 848
 849                         foreach (int n in d.Keys)
 850                                 if (!t.Contains (n))
 851                                         throw new Exception ("tree says it doesn't have a number it should");
 852
 853                         Dictionary<int, int> d1 = new Dictionary<int, int> (d);
 854
 855                         int prev = -1;
 856                         foreach (int n in t) {
 857                                 if (n < prev)
 858                                         throw new Exception ("iteration out of order");
 859                                 if (!d1.Remove (n))
 860                                         throw new Exception ("tree has a number it shouldn't");
 861                                 prev = n;
 862                         }
 863
 864                         if (d1.Count != 0)
 865                                 throw new Exception ("tree has numbers it shouldn't");
 866
 867                         for (int i = 0; i < iters; ++i) {
 868                                 int n = r.Next ();
 869                                 if (!d.ContainsKey (n)) {
 870                                         if (t.Contains (n))
 871                                                 throw new Exception ("tree says it doesn't have a number it should");
 872                                 } else if (!t.Contains (n)) {
 873                                         throw new Exception ("tree says it has a number it shouldn't");
 874                                 }
 875                         }
 876
 877                         int count = t.Count;
 878                         foreach (int n in d.Keys) {
 879                                 if (count <= watermark) {
 880                                         watermark -= watermark/4;
 881                                         t.VerifyInvariants ();
 882                                 }
 883                                 try {
 884                                         if (!t.Remove (n))
 885                                                 throw new Exception ("tree says it doesn't have a number it should");
 886                                         --count;
 887                                         if (t.Count != count)
 888                                                 throw new Exception ("Remove didn't remove exactly one element");
 889                                 } catch {
 890                                         Console.Error.WriteLine ("While trying to remove {0} from tree of size {1}", n, t.Count);
 891                                         t.Dump ();
 892                                         t.VerifyInvariants ();
 893                                         throw;
 894                                 }
 895                                 if (t.Contains (n))
 896                                         throw new Exception ("tree says it has a number it shouldn't");
 897                         }
 898                         t.VerifyInvariants ();
 899
 900                         if (t.Count != 0)
 901                                 throw new Exception ("tree claims to have elements");
 902                 }
 903         }
 904 }
 905 #endif
 906
 907 #endif