bundles/org.simantics.utils.datastructures/src/org/simantics/utils/datastructures/persistent/ImmutableMap.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.TObjectObjectProcedure;
  15
  16 public class ImmutableMap<T extends Comparable<T>, V> {
  17
  18         @SuppressWarnings({ "rawtypes" })
  19         static final ImmutableMap EMPTY = new EmptyImmutableSet();
  20
  21         private static class EmptyImmutableSet<T extends Comparable<T>, V> extends ImmutableMap<T, V> {
  22                 public EmptyImmutableSet() {
  23                         isBlack = true;
  24                 }
  25
  26                 protected ImmutableMap<T, V> putRec(T key, V value) {
  27                         return new ImmutableMap<T, V>(key, value);
  28                 }
  29
  30                 protected ImmutableMap<T, V> removeRec(T obj) {
  31                         return null;
  32                 }
  33
  34                 public boolean contains(T obj) {
  35                         return false;
  36                 }
  37
  38                 public V get(T key) {
  39                         return null;
  40                 }
  41
  42                 @Override
  43                 public boolean forEach(TObjectObjectProcedure<T, V> arg0) {
  44                         return true;
  45                 }
  46         }
  47
  48         ImmutableMap<T, V> left;
  49         T key;
  50         V value;
  51         ImmutableMap<T, V> right;
  52         boolean isBlack;
  53
  54         protected ImmutableMap(
  55                         ImmutableMap<T, V> left,
  56                         T key,
  57                         V value,
  58                         ImmutableMap<T, V> right,
  59                         boolean isBlack) {
  60                 this.left = left;
  61                 this.right = right;
  62                 this.key = key;
  63                 this.value = value;
  64                 this.isBlack = isBlack;
  65         }
  66
  67         @SuppressWarnings("unchecked")
  68         public ImmutableMap(T key, V value) {
  69                 this(EMPTY, key, value, EMPTY, false);
  70         }
  71
  72         private ImmutableMap() {
  73         }
  74
  75         public boolean contains(T obj) {
  76                 int cmp = obj.compareTo(key);
  77                 if(cmp < 0)
  78                         return left.contains(obj);
  79                 else if(cmp > 0)
  80                         return right.contains(obj);
  81                 else
  82                         return true;
  83         }
  84
  85         public V get(T key) {
  86                 int cmp = key.compareTo(key);
  87                 if(cmp < 0)
  88                         return left.get(key);
  89                 else if(cmp > 0)
  90                         return right.get(key);
  91                 else
  92                         return value;
  93         }
  94
  95         protected ImmutableMap<T, V> putRec(T obj, V value) {
  96                 int cmp = obj.compareTo(key);
  97                 if(cmp < 0) {
  98                         ImmutableMap<T, V> newLeft = left.putRec(obj, value);
  99                         if(newLeft == left)
 100                                 return this;
 101                         if(isBlack)
 102                                 return balance(newLeft, key, value, right);
 103                         else
 104                                 return red(newLeft, key, value, right);
 105                 }
 106                 else if(cmp > 0) {
 107                         ImmutableMap<T, V> newRight = right.putRec(obj, value);
 108                         if(newRight == right)
 109                                 return this;
 110                         if(isBlack)
 111                                 return balance(left, key, value, newRight);
 112                         else
 113                                 return red(left, key, value, newRight);
 114                 }
 115                 else
 116                         return this;
 117         }
 118
 119         /**
 120          * Assumes this is a black nonempty node.
 121          *
 122          * Removes obj from the tree. The returned tree has always
 123          * one black node less in every branch (even if nothing is
 124          * removed).
 125          *
 126          * @param obj
 127          * @return
 128          */
 129         protected ImmutableMap<T, V> removeRec(T obj) {
 130                 int cmp = obj.compareTo(key);
 131                 if(cmp < 0) {
 132                         ImmutableMap<T, V> newLeft = left.removeRec(obj);
 133                         if(newLeft == null)
 134                                 return null;
 135                         else if(left.isBlack)
 136                                 return balleft(newLeft, key, value, right);
 137                         else
 138                                 return red(newLeft, key, value, right);
 139                 }
 140                 else if(cmp > 0) {
 141                         ImmutableMap<T, V> newRight = right.removeRec(obj);
 142                         if(newRight == null)
 143                                 return null;
 144                         else if(right.isBlack)
 145                                 return balright(left, key, value, newRight);
 146                         else
 147                                 return red(left, key, value, newRight);
 148                 }
 149                 else
 150                         return append(left, right);
 151         }
 152
 153         /**
 154          * Assumes a and b have the same black depth and keys in a are strictly less
 155          * than keys in b. Creates a new tree with the same black depth as a and b.
 156          *
 157          * @param <T>
 158          * @param a
 159          * @param b
 160          * @return
 161          */
 162         protected static <T extends Comparable<T>, V> ImmutableMap<T, V> append(ImmutableMap<T, V> a, ImmutableMap<T, V> b) {
 163                 if(a==EMPTY)
 164                         return b;
 165                 if(b==EMPTY)
 166                         return a;
 167                 if(a.isBlack) {
 168                         if(b.isBlack) {
 169                                 ImmutableMap<T, V> middle = append(a.right, b.left);
 170                                 if(middle.isBlack)
 171                                         return balleft(a.left, a.key, a.value, black(middle, b.key, b.value, b.right));
 172                                 else
 173                                         return red(black(a.left, a.key, a.value, middle.left), middle.key, middle.value, black(middle.right, b.key, b.value, b.right));
 174                         }
 175                         else
 176                                 return red(append(a, b.left), b.key, b.value, b.right);
 177                 }
 178                 else {
 179                         if(b.isBlack)
 180                                 return red(a.left, a.key, a.value, append(a.right, b));
 181                         else {
 182                                 ImmutableMap<T, V> middle = append(a.right, b.left);
 183                                 if(middle.isBlack)
 184                                         return red(a.left, a.key, a.value, red(middle, b.key, b.value, b.right));
 185                                 else
 186                                         return red(red(a.left, a.key, a.value, middle.left), middle.key, middle.value, red(middle.right, b.key, b.value, b.right));
 187                         }
 188                 }
 189         }
 190
 191         public T getFirst() {
 192                 if(left == EMPTY)
 193                         return key;
 194                 else
 195                         return left.getFirst();
 196         }
 197
 198         static private <T extends Comparable<T>, V> ImmutableMap<T, V> balance(
 199                         ImmutableMap<T, V> left,
 200                         T key,
 201                         V value,
 202                         ImmutableMap<T, V> right) {
 203                 if(!left.isBlack) {
 204                         if(!left.left.isBlack)
 205                                 return red(
 206                                         toBlack(left.left),
 207                                         left.key,
 208                                         left.value,
 209                                         black(left.right, key, value, right)
 210                                 );
 211                         else if(!left.right.isBlack)
 212                                 return red(
 213                                         black(left.left, left.key, left.value, left.right.left),
 214                                         left.right.key, left.right.value,
 215                                         black(left.right.right, key, value, right)
 216                                 );
 217                 }
 218                 if(!right.isBlack) {
 219                         if(!right.left.isBlack)
 220                                 return red(
 221                                         black(left, key, value, right.left.left),
 222                                         right.left.key, right.left.value,
 223                                         black(right.left.right, right.key, right.value, right.right)
 224                                 );
 225                         else if(!right.right.isBlack)
 226                                 return red(
 227                                         black(left, key, value, right.left),
 228                                         right.key,
 229                                         right.value,
 230                                         toBlack(right.right)
 231                                 );
 232                 }
 233                 return black(left, key, value, right);
 234         }
 235
 236         static private <T extends Comparable<T>, V> ImmutableMap<T, V> black(
 237                         ImmutableMap<T, V> left,
 238                         T key,
 239                         V value,
 240                         ImmutableMap<T, V> right) {
 241                 return new ImmutableMap<T, V>(left, key, value, right, true);
 242         }
 243
 244         static private <T extends Comparable<T>, V> ImmutableMap<T, V> red(
 245                         ImmutableMap<T, V> left,
 246                         T key,
 247                         V value,
 248                         ImmutableMap<T, V> right) {
 249                 return new ImmutableMap<T, V>(left, key, value, right, false);
 250         }
 251
 252         static private <T extends Comparable<T>, V> ImmutableMap<T, V> toBlack(ImmutableMap<T, V> tree) {
 253                 if(tree.isBlack)
 254                         return tree;
 255                 else
 256                         return black(tree.left, tree.key, tree.value, tree.right);
 257         }
 258
 259         static private <T extends Comparable<T>, V> ImmutableMap<T, V> toRed(ImmutableMap<T, V> tree) {
 260                 if(tree.isBlack)
 261                         return red(tree.left, tree.key, tree.value, tree.right);
 262                 else
 263                         return tree;
 264         }
 265
 266
 267         static private <T extends Comparable<T>, V> ImmutableMap<T, V> balleft(
 268                         ImmutableMap<T, V> left,
 269                         T key,
 270                         V value,
 271                         ImmutableMap<T, V> right) {
 272                 if(left.isBlack) {
 273                         if(right.isBlack)
 274                                 return balance(left, key, value, toRed(right));
 275                         else
 276                                 return red(black(left, key, value, right.left.left), right.left.key, right.left.value,
 277                                         balance(right.left.right, right.key, right.value, toRed(right.right)));
 278                 }
 279                 else
 280                         return red(toBlack(left), key, value, right);
 281         }
 282
 283         static private <T extends Comparable<T>, V> ImmutableMap<T, V> balright(
 284                         ImmutableMap<T, V> left,
 285                         T key,
 286                         V value,
 287                         ImmutableMap<T, V> right) {
 288                 if(right.isBlack) {
 289                         if(left.isBlack)
 290                                 return balance(toRed(left), key, value, right);
 291                         else
 292                                 return red(balance(toRed(left.left), left.key, left.value, left.right.left), left.right.key, left.right.value,
 293                                         black(left.right.right, key, value, right));
 294                 }
 295                 else
 296                         return red(left, key, value, toBlack(right));
 297         }
 298
 299         public ImmutableMap<T, V> put(T key, V value) {
 300                 ImmutableMap<T, V> tree = putRec(key, value);
 301                 tree.isBlack = true;
 302                 return tree;
 303         }
 304
 305         public ImmutableMap<T, V> remove(T obj) {
 306                 ImmutableMap<T, V> tree = removeRec(obj);
 307                 if(tree == null)
 308                         return this;
 309                 if(tree.isBlack)
 310                         return tree;
 311                 else
 312                         return black(tree.left, tree.key, tree.value, tree.right);
 313         }
 314
 315         public boolean forEach(TObjectObjectProcedure<T, V> proc) {
 316                 if(left.forEach(proc))
 317                         if(proc.execute(key, value))
 318                                 return right.forEach(proc);
 319                 return false;
 320         }
 321
 322 }