bundles/org.simantics.utils.datastructures/src/org/simantics/utils/datastructures/persistent/ImmutableSet.java

   1 /*******************************************************************************
   2  * Copyright (c) 2007, 2010 Association for Decentralized Information Management
   3  * in Industry THTH ry.
   4  * All rights reserved. This program and the accompanying materials
   5  * are made available under the terms of the Eclipse Public License v1.0
   6  * which accompanies this distribution, and is available at
   7  * http://www.eclipse.org/legal/epl-v10.html
   8  *
   9  * Contributors:
  10  *     VTT Technical Research Centre of Finland - initial API and implementation
  11  *******************************************************************************/
  12 package org.simantics.utils.datastructures.persistent;
  13
  14 import gnu.trove.procedure.TObjectProcedure;
  15
  16 import java.util.Collection;
  17 import java.util.HashSet;
  18 import java.util.Iterator;
  19 import java.util.Random;
  20 import java.util.Set;
  21
  22 public class ImmutableSet<T extends Comparable<T>> {
  23
  24         @SuppressWarnings({ "rawtypes" })
  25         static final ImmutableSet EMPTY = new EmptyImmutableSet();
  26
  27         private static class EmptyImmutableSet<T extends Comparable<T>> extends ImmutableSet<T> {
  28                 public EmptyImmutableSet() {
  29                         isBlack = true;
  30                 }
  31
  32                 protected org.simantics.utils.datastructures.persistent.ImmutableSet<T> addRec(T obj) {
  33                         return new ImmutableSet<T>(obj);
  34                 }
  35
  36                 protected org.simantics.utils.datastructures.persistent.ImmutableSet<T> removeRec(T obj) {
  37                         return null;
  38                 }
  39
  40                 public boolean contains(T obj) {
  41                         return false;
  42                 }
  43
  44                 @Override
  45                 public boolean forEach(TObjectProcedure<T> arg0) {
  46                         return true;
  47                 }
  48
  49                 @Override
  50                 void print(int arg0) {
  51                 }
  52         }
  53
  54         ImmutableSet<T> left;
  55         T key;
  56         ImmutableSet<T> right;
  57         boolean isBlack;
  58
  59         protected ImmutableSet(
  60                         ImmutableSet<T> left,
  61                         T key,
  62                         ImmutableSet<T> right,
  63                         boolean isBlack) {
  64                 this.left = left;
  65                 this.right = right;
  66                 this.key = key;
  67                 this.isBlack = isBlack;
  68         }
  69
  70         @SuppressWarnings("unchecked")
  71         public ImmutableSet(T key) {
  72                 this(EMPTY, key, EMPTY, false);
  73         }
  74
  75         @SuppressWarnings("unchecked")
  76         public static <T extends Comparable<T>> ImmutableSet<T> of(T ... array) {
  77                 if(array.length == 0)
  78                         return EMPTY;
  79                 ImmutableSet<T> ret = new ImmutableSet<T>(array[0]);
  80                 for(int i=1;i<array.length;++i)
  81                         ret = ret.add(array[i]);
  82                 return ret;
  83         }
  84
  85         @SuppressWarnings("unchecked")
  86         public static <T extends Comparable<T>> ImmutableSet<T> of(Collection<T> c) {
  87                 Iterator<T> it = c.iterator();
  88                 if(!it.hasNext())
  89                         return EMPTY;
  90                 ImmutableSet<T> ret = new ImmutableSet<T>(it.next());
  91                 while(it.hasNext())
  92                         ret = ret.add(it.next());
  93                 return ret;
  94         }
  95
  96         private ImmutableSet() {
  97         }
  98
  99         public boolean contains(T obj) {
 100                 int cmp = obj.compareTo(key);
 101                 if(cmp < 0)
 102                         return left.contains(obj);
 103                 else if(cmp > 0)
 104                         return right.contains(obj);
 105                 else
 106                         return true;
 107         }
 108
 109         protected ImmutableSet<T> addRec(T obj) {
 110                 int cmp = obj.compareTo(key);
 111                 if(cmp < 0) {
 112                         ImmutableSet<T> newLeft = left.addRec(obj);
 113                         if(newLeft == left)
 114                                 return this;
 115                         if(isBlack)
 116                                 return balance(newLeft, key, right);
 117                         else
 118                                 return red(newLeft, key, right);
 119                 }
 120                 else if(cmp > 0) {
 121                         ImmutableSet<T> newRight = right.addRec(obj);
 122                         if(newRight == right)
 123                                 return this;
 124                         if(isBlack)
 125                                 return balance(left, key, newRight);
 126                         else
 127                                 return red(left, key, newRight);
 128                 }
 129                 else
 130                         return this;
 131         }
 132
 133         /**
 134          * Assumes this is a black nonempty node.
 135          *
 136          * Removes obj from the tree. The returned tree has always
 137          * one black node less in every branch (even if nothing is
 138          * removed).
 139          *
 140          * @param obj
 141          * @return
 142          */
 143         protected ImmutableSet<T> removeRec(T obj) {
 144                 int cmp = obj.compareTo(key);
 145                 if(cmp < 0) {
 146                         ImmutableSet<T> newLeft = left.removeRec(obj);
 147                         if(newLeft == null)
 148                                 return null;
 149                         else if(left.isBlack)
 150                                 return balleft(newLeft, key, right);
 151                         else
 152                                 return red(newLeft, key, right);
 153                 }
 154                 else if(cmp > 0) {
 155                         ImmutableSet<T> newRight = right.removeRec(obj);
 156                         if(newRight == null)
 157                                 return null;
 158                         else if(right.isBlack)
 159                                 return balright(left, key, newRight);
 160                         else
 161                                 return red(left, key, newRight);
 162                 }
 163                 else
 164                         return append(left, right);
 165         }
 166
 167         /**
 168          * Assumes a and b have the same black depth and keys in a are strictly less
 169          * than keys in b. Creates a new tree with the same black depth as a and b.
 170          *
 171          * @param <T>
 172          * @param a
 173          * @param b
 174          * @return
 175          */
 176         protected static <T extends Comparable<T>> ImmutableSet<T> append(ImmutableSet<T> a, ImmutableSet<T> b) {
 177                 if(a==EMPTY)
 178                         return b;
 179                 if(b==EMPTY)
 180                         return a;
 181                 if(a.isBlack) {
 182                         if(b.isBlack) {
 183                                 ImmutableSet<T> middle = append(a.right, b.left);
 184                                 if(middle.isBlack)
 185                                         return balleft(a.left, a.key, black(middle, b.key, b.right));
 186                                 else
 187                                         return red(black(a.left, a.key, middle.left), middle.key, black(middle.right, b.key, b.right));
 188                         }
 189                         else
 190                                 return red(append(a, b.left), b.key, b.right);
 191                 }
 192                 else {
 193                         if(b.isBlack)
 194                                 return red(a.left, a.key, append(a.right, b));
 195                         else {
 196                                 ImmutableSet<T> middle = append(a.right, b.left);
 197                                 if(middle.isBlack)
 198                                         return red(a.left, a.key, red(middle, b.key, b.right));
 199                                 else
 200                                         return red(red(a.left, a.key, middle.left), middle.key, red(middle.right, b.key, b.right));
 201                         }
 202                 }
 203         }
 204
 205         public T getFirst() {
 206                 if(left == EMPTY)
 207                         return key;
 208                 else
 209                         return left.getFirst();
 210         }
 211
 212         static private <T extends Comparable<T>> ImmutableSet<T> balance(
 213                         ImmutableSet<T> left,
 214                         T key,
 215                         ImmutableSet<T> right) {
 216                 if(!left.isBlack) {
 217                         if(!left.left.isBlack)
 218                                 return red(
 219                                         toBlack(left.left),
 220                                         left.key,
 221                                         black(left.right, key, right)
 222                                 );
 223                         else if(!left.right.isBlack)
 224                                 return red(
 225                                         black(left.left, left.key, left.right.left),
 226                                         left.right.key,
 227                                         black(left.right.right, key, right)
 228                                 );
 229                 }
 230                 if(!right.isBlack) {
 231                         if(!right.left.isBlack)
 232                                 return red(
 233                                         black(left, key, right.left.left),
 234                                         right.left.key,
 235                                         black(right.left.right, right.key, right.right)
 236                                 );
 237                         else if(!right.right.isBlack)
 238                                 return red(
 239                                         black(left, key, right.left),
 240                                         right.key,
 241                                         toBlack(right.right)
 242                                 );
 243                 }
 244                 return black(left, key, right);
 245         }
 246
 247         static private <T extends Comparable<T>> ImmutableSet<T> black(
 248                         ImmutableSet<T> left,
 249                         T key,
 250                         ImmutableSet<T> right) {
 251                 return new ImmutableSet<T>(left, key, right, true);
 252         }
 253
 254         static private <T extends Comparable<T>> ImmutableSet<T> red(
 255                         ImmutableSet<T> left,
 256                         T key,
 257                         ImmutableSet<T> right) {
 258                 return new ImmutableSet<T>(left, key, right, false);
 259         }
 260
 261         static private <T extends Comparable<T>> ImmutableSet<T> toBlack(ImmutableSet<T> tree) {
 262                 if(tree.isBlack)
 263                         return tree;
 264                 else
 265                         return black(tree.left, tree.key, tree.right);
 266         }
 267
 268         static private <T extends Comparable<T>> ImmutableSet<T> toRed(ImmutableSet<T> tree) {
 269                 if(tree.isBlack)
 270                         return red(tree.left, tree.key, tree.right);
 271                 else
 272                         return tree;
 273         }
 274
 275
 276         static private <T extends Comparable<T>> ImmutableSet<T> balleft(
 277                         ImmutableSet<T> left,
 278                         T key,
 279                         ImmutableSet<T> right) {
 280                 if(left.isBlack) {
 281                         if(right.isBlack)
 282                                 return balance(left, key, toRed(right));
 283                         else
 284                                 return red(black(left, key, right.left.left), right.left.key, balance(right.left.right, right.key, toRed(right.right)));
 285                 }
 286                 else
 287                         return red(toBlack(left), key, right);
 288         }
 289
 290         static private <T extends Comparable<T>> ImmutableSet<T> balright(
 291                         ImmutableSet<T> left,
 292                         T key,
 293                         ImmutableSet<T> right) {
 294                 if(right.isBlack) {
 295                         if(left.isBlack)
 296                                 return balance(toRed(left), key, right);
 297                         else
 298                                 return red(balance(toRed(left.left), left.key, left.right.left), left.right.key, black(left.right.right, key, right));
 299                 }
 300                 else
 301                         return red(left, key, toBlack(right));
 302         }
 303
 304         public ImmutableSet<T> add(T obj) {
 305                 ImmutableSet<T> tree = addRec(obj);
 306                 tree.isBlack = true;
 307                 return tree;
 308         }
 309
 310         public ImmutableSet<T> remove(T obj) {
 311                 ImmutableSet<T> tree = removeRec(obj);
 312                 if(tree == null)
 313                         return this;
 314                 if(tree.isBlack)
 315                         return tree;
 316                 else
 317                         return black(tree.left, tree.key, tree.right);
 318         }
 319
 320         public boolean forEach(TObjectProcedure<T> procedure) {
 321                 if(left.forEach(procedure))
 322                         if(procedure.execute(key))
 323                                 return right.forEach(procedure);
 324                 return false;
 325         }
 326
 327         public ImmutableSet<T> addAll(Iterable<T> c) {
 328                 ImmutableSet<T> ret = this;
 329                 for(T t : c)
 330                         ret = ret.add(t);
 331                 return ret;
 332         }
 333
 334         static class AddAll<T extends Comparable<T>> implements TObjectProcedure<T> {
 335
 336                 ImmutableSet<T> result;
 337
 338                 public AddAll(ImmutableSet<T> result) {
 339                         this.result = result;
 340                 }
 341
 342                 @Override
 343                 public boolean execute(T arg0) {
 344                         result = result.add(arg0);
 345                         return true;
 346                 }
 347
 348         }
 349
 350         public ImmutableSet<T> addAll(ImmutableSet<T> c) {
 351                 if(this==EMPTY)
 352                         return c;
 353                 if(c==EMPTY)
 354                         return this;
 355                 AddAll<T> proc = new AddAll<T>(this);
 356                 c.forEach(proc);
 357                 return proc.result;
 358         }
 359
 360         /**************************************************************************
 361          *
 362          *  Testing
 363          *
 364          **************************************************************************
 365          */
 366
 367         void print(int indentation) {
 368                 left.print(indentation + 1);
 369                 for(int i=0;i<indentation;++i)
 370                         System.out.print("    ");
 371                 System.out.println(key + " " + (isBlack ? "BLACK" : "RED"));
 372                 right.print(indentation + 1);
 373         }
 374
 375         private int validateRec() {
 376                 if(this==EMPTY)
 377                         return 1;
 378                 int lh = left.validateRec();
 379                 int rh = right.validateRec();
 380                 if(lh != rh)
 381                         System.out.println("Unbalanced!");
 382                 if(!isBlack) {
 383                         if(!left.isBlack || !right.isBlack)
 384                                 System.out.println("Red under red");
 385                         return lh;
 386                 }
 387                 else
 388                         return lh+1;
 389         }
 390
 391         @SuppressWarnings("unchecked")
 392         public static <T extends Comparable<T>> ImmutableSet<T> empty() {
 393                 return EMPTY;
 394         }
 395
 396         public void validate() {
 397                 validateRec();
 398                 if(!isBlack)
 399                         System.out.println("Root is not black");
 400         }
 401
 402         @SuppressWarnings("unchecked")
 403         public static void main(String[] args) {
 404                 ImmutableSet<Integer> set = ImmutableSet.EMPTY;
 405                 final Set<Integer> set2 = new HashSet<Integer>();
 406
 407                 Random random = new Random();
 408                 for(int i=0;i<10000;++i) {
 409                         int r1 = random.nextInt(1000);
 410                         int r2 = random.nextInt(1000);
 411                         set = set.add(r1);
 412                         set = set.remove(r2);
 413                         set2.add(r1);
 414                         set2.remove(r2);
 415                         set.validate();
 416
 417                         for(Integer v : set2)
 418                                 if(!set.contains(v))
 419                                         System.out.println("set2 has more elements");
 420                         set.forEach(new TObjectProcedure<Integer>() {
 421
 422                                 @Override
 423                                 public boolean execute(Integer arg0) {
 424                                         if(!set2.contains(arg0))
 425                                                 System.out.println("set has more elements");
 426                                         return true;
 427                                 }
 428
 429                         });
 430                 }
 431
 432                 /*
 433                 map.forEach(new TObjectProcedure<Integer>() {
 434
 435                         @Override
 436                         public boolean execute(Integer arg0) {
 437                                 System.out.println(arg0);
 438                                 return true;
 439                         }
 440
 441                 });
 442                 */
 443         }
 444
 445 }