001 package fj.data; 002 003 import fj.F; 004 import fj.F2; 005 import fj.Function; 006 import fj.P; 007 import fj.P2; 008 import fj.P3; 009 import static fj.Function.*; 010 import static fj.data.Either.right; 011 import static fj.data.Option.some; 012 import static fj.function.Booleans.not; 013 import fj.pre.Monoid; 014 import fj.pre.Ord; 015 import fj.pre.Ordering; 016 import static fj.pre.Ordering.GT; 017 import static fj.pre.Ordering.LT; 018 019 import java.util.Iterator; 020 021 /** 022 * Provides an in-memory, immutable set, implemented as a red/black tree. 023 */ 024 public abstract class Set<A> implements Iterable<A> { 025 private Set(final Ord<A> ord) { 026 this.ord = ord; 027 } 028 029 private enum Color { 030 R, B 031 } 032 033 private final Ord<A> ord; 034 035 public boolean isEmpty() { 036 return this instanceof Empty; 037 } 038 039 @SuppressWarnings({"ClassEscapesDefinedScope"}) 040 protected abstract Color color(); 041 042 abstract Set<A> l(); 043 044 abstract A head(); 045 046 abstract Set<A> r(); 047 048 /** 049 * Returns the order of this Set. 050 * 051 * @return the order of this Set. 052 */ 053 public Ord<A> ord() { 054 return ord; 055 } 056 057 private static final class Empty<A> extends Set<A> { 058 private Empty(final Ord<A> ord) { 059 super(ord); 060 } 061 062 public Color color() { 063 return Color.B; 064 } 065 066 public Set<A> l() { 067 throw new Error("Left on empty set."); 068 } 069 070 public Set<A> r() { 071 throw new Error("Right on empty set."); 072 } 073 074 public A head() { 075 throw new Error("Head on empty set."); 076 } 077 } 078 079 private static final class Tree<A> extends Set<A> { 080 private final Color c; 081 private final Set<A> a; 082 private final A x; 083 private final Set<A> b; 084 085 private Tree(final Ord<A> ord, final Color c, final Set<A> a, final A x, final Set<A> b) { 086 super(ord); 087 this.c = c; 088 this.a = a; 089 this.x = x; 090 this.b = b; 091 } 092 093 public Color color() { 094 return c; 095 } 096 097 public Set<A> l() { 098 return a; 099 } 100 101 public A head() { 102 return x; 103 } 104 105 public Set<A> r() { 106 return b; 107 } 108 } 109 110 /** 111 * Updates, with the given function, the first element in the set that is equal to the given element, 112 * according to the order. 113 * 114 * @param a An element to replace. 115 * @param f A function to transforms the found element. 116 * @return A pair of: (1) True if an element was found that matches the given element, otherwise false. 117 * (2) A new set with the given function applied to the first set element 118 * that was equal to the given element. 119 */ 120 public P2<Boolean, Set<A>> update(final A a, final F<A, A> f) { 121 return isEmpty() 122 ? P.p(false, this) 123 : tryUpdate(a, f).either(new F<A, P2<Boolean, Set<A>>>() { 124 public P2<Boolean, Set<A>> f(final A a2) { 125 return P.p(true, delete(a).insert(a2)); 126 } 127 }, Function.<P2<Boolean, Set<A>>>identity()); 128 } 129 130 private Either<A, P2<Boolean, Set<A>>> tryUpdate(final A a, final F<A, A> f) { 131 if (isEmpty()) 132 return right(P.p(false, this)); 133 else if (ord.isLessThan(a, head())) 134 return l().tryUpdate(a, f).right().map(new F<P2<Boolean, Set<A>>, P2<Boolean, Set<A>>>() { 135 public P2<Boolean, Set<A>> f(final P2<Boolean, Set<A>> set) { 136 return set._1() ? P.p(true, (Set<A>) new Tree<A>(ord, color(), set._2(), head(), r())) : set; 137 } 138 }); 139 else if (ord.eq(a, head())) { 140 final A h = f.f(head()); 141 return ord.eq(head(), h) ? Either 142 .<A, P2<Boolean, Set<A>>>right(P.p(true, (Set<A>) new Tree<A>(ord, color(), l(), h, r()))) 143 : Either.<A, P2<Boolean, Set<A>>>left(h); 144 } else return r().tryUpdate(a, f).right().map(new F<P2<Boolean, Set<A>>, P2<Boolean, Set<A>>>() { 145 public P2<Boolean, Set<A>> f(final P2<Boolean, Set<A>> set) { 146 return set._1() ? P.p(true, (Set<A>) new Tree<A>(ord, color(), l(), head(), set._2())) : set; 147 } 148 }); 149 } 150 151 /** 152 * The empty set. 153 * 154 * @param ord An order for the type of elements. 155 * @return the empty set. 156 */ 157 public static <A> Set<A> empty(final Ord<A> ord) { 158 return new Empty<A>(ord); 159 } 160 161 /** 162 * Checks if the given element is a member of this set. 163 * 164 * @param x An element to check for membership in this set. 165 * @return true if the given element is a member of this set. 166 */ 167 public boolean member(final A x) { 168 return !isEmpty() && (ord.isLessThan(x, head()) && l().member(x) || ord.eq(head(), x) || r().member(x)); 169 } 170 171 172 /** 173 * First-class membership check. 174 * 175 * @return A function that returns true if the given element if a member of the given set. 176 */ 177 public static <A> F<Set<A>, F<A, Boolean>> member() { 178 return curry(new F2<Set<A>, A, Boolean>() { 179 public Boolean f(final Set<A> s, final A a) { 180 return s.member(a); 181 } 182 }); 183 } 184 185 /** 186 * Inserts the given element into this set. 187 * 188 * @param x An element to insert into this set. 189 * @return A new set with the given element inserted. 190 */ 191 public Set<A> insert(final A x) { 192 return ins(x).makeBlack(); 193 } 194 195 /** 196 * First-class insertion function. 197 * 198 * @return A function that inserts a given element into a given set. 199 */ 200 public static <A> F<A, F<Set<A>, Set<A>>> insert() { 201 return curry(new F2<A, Set<A>, Set<A>>() { 202 public Set<A> f(final A a, final Set<A> set) { 203 return set.insert(a); 204 } 205 }); 206 } 207 208 private Set<A> ins(final A x) { 209 return isEmpty() 210 ? new Tree<A>(ord, Color.R, empty(ord), x, empty(ord)) 211 : ord.isLessThan(x, head()) 212 ? balance(ord, color(), l().ins(x), head(), r()) 213 : ord.eq(x, head()) 214 ? new Tree<A>(ord, color(), l(), x, r()) 215 : balance(ord, color(), l(), head(), r().ins(x)); 216 } 217 218 private Set<A> makeBlack() { 219 return new Tree<A>(ord, Color.B, l(), head(), r()); 220 } 221 222 private static <A> Tree<A> tr(final Ord<A> o, 223 final Set<A> a, final A x, final Set<A> b, 224 final A y, 225 final Set<A> c, final A z, final Set<A> d) { 226 return new Tree<A>(o, Color.R, new Tree<A>(o, Color.B, a, x, b), y, new Tree<A>(o, Color.B, c, z, d)); 227 } 228 229 private static <A> Set<A> balance(final Ord<A> ord, final Color c, final Set<A> l, final A h, final Set<A> r) { 230 if (c == Color.B && l.isTR() && l.l().isTR()) { 231 return tr(ord, l.l().l(), l.l().head(), l.l().r(), l.head(), l.r(), h, r); 232 } else if (c == Color.B && l.isTR() && l.r().isTR()) { 233 return tr(ord, l.l(), l.head(), l.r().l(), l.r().head(), l.r().r(), h, r); 234 } else if (c == Color.B && r.isTR() && r.l().isTR()) { 235 return tr(ord, l, h, r.l().l(), r.l().head(), r.l().r(), r.head(), r.r()); 236 } else if (c == Color.B && r.isTR() && r.r().isTR()) { 237 return tr(ord, l, h, r.l(), r.head(), r.r().l(), r.r().head(), r.r().r()); 238 } else { 239 return new Tree<A>(ord, c, l, h, r); 240 } 241 } 242 243 private boolean isTR() { 244 return !isEmpty() && color() == Color.R; 245 } 246 247 /** 248 * Returns an iterator over this set. 249 * 250 * @return an iterator over this set. 251 */ 252 public Iterator<A> iterator() { 253 return toStream().iterator(); 254 } 255 256 /** 257 * Returns a set with a single element. 258 * 259 * @param o An order for the type of element. 260 * @param a An element to put in a set. 261 * @return A new set with the given element in it. 262 */ 263 public static <A> Set<A> single(final Ord<A> o, final A a) { 264 return empty(o).insert(a); 265 } 266 267 /** 268 * Maps the given function across this set. 269 * 270 * @param o An order for the elements of the new set. 271 * @param f A function to map across this set. 272 * @return The set of the results of applying the given function to the elements of this set. 273 */ 274 public <B> Set<B> map(final Ord<B> o, final F<A, B> f) { 275 return iterableSet(o, toStream().map(f)); 276 } 277 278 /** 279 * Folds this Set using the given monoid. 280 * 281 * @param f A transformation from this Set's elements, to the monoid. 282 * @param m The monoid to fold this Set with. 283 * @return The result of folding the Set with the given monoid. 284 */ 285 public <B> B foldMap(final F<A, B> f, final Monoid<B> m) { 286 return isEmpty() ? 287 m.zero() : 288 m.sum(m.sum(r().foldMap(f, m), f.f(head())), l().foldMap(f, m)); 289 } 290 291 /** 292 * Returns a list representation of this set. 293 * 294 * @return a list representation of this set. 295 */ 296 public List<A> toList() { 297 return foldMap(List.cons(List.<A>nil()), Monoid.<A>listMonoid()); 298 } 299 300 /** 301 * Returns a stream representation of this set. 302 * 303 * @return a stream representation of this set. 304 */ 305 public Stream<A> toStream() { 306 return foldMap(Stream.<A>single(), Monoid.<A>streamMonoid()); 307 } 308 309 /** 310 * Binds the given function across this set. 311 * 312 * @param o An order for the elements of the target set. 313 * @param f A function to bind across this set. 314 * @return A new set after applying the given function and joining the resulting sets. 315 */ 316 public <B> Set<B> bind(final Ord<B> o, final F<A, Set<B>> f) { 317 return join(o, map(Ord.setOrd(o), f)); 318 } 319 320 /** 321 * Add all the elements of the given set to this set. 322 * 323 * @param s A set to add to this set. 324 * @return A new set containing all elements of both sets. 325 */ 326 public Set<A> union(final Set<A> s) { 327 return iterableSet(ord, s.toStream().append(toStream())); 328 } 329 330 /** 331 * Filters elements from this set by returning only elements which produce <code>true</code> 332 * when the given function is applied to them. 333 * 334 * @param f The predicate function to filter on. 335 * @return A new set whose elements all match the given predicate. 336 */ 337 public Set<A> filter(final F<A, Boolean> f) { 338 return iterableSet(ord, toStream().filter(f)); 339 } 340 341 /** 342 * Deletes the given element from this set. 343 * 344 * @param a an element to remove. 345 * @return A new set containing all the elements of this set, except the given element. 346 */ 347 public Set<A> delete(final A a) { 348 return minus(single(ord, a)); 349 } 350 351 /** 352 * First-class deletion function. 353 * 354 * @return A function that deletes a given element from a given set. 355 */ 356 public F<A, F<Set<A>, Set<A>>> delete() { 357 return curry(new F2<A, Set<A>, Set<A>>() { 358 public Set<A> f(final A a, final Set<A> set) { 359 return set.delete(a); 360 } 361 }); 362 } 363 364 /** 365 * Remove all elements from this set that do not occur in the given set. 366 * 367 * @param s A set of elements to retain. 368 * @return A new set which is the intersection of this set and the given set. 369 */ 370 public Set<A> intersect(final Set<A> s) { 371 return filter(Set.<A>member().f(s)); 372 } 373 374 /** 375 * Remove all elements from this set that occur in the given set. 376 * 377 * @param s A set of elements to delete. 378 * @return A new set which contains only the elements of this set that do not occur in the given set. 379 */ 380 public Set<A> minus(final Set<A> s) { 381 return filter(compose(not, Set.<A>member().f(s))); 382 } 383 384 /** 385 * Returns the size of this set. 386 * 387 * @return The number of elements in this set. 388 */ 389 public int size() { 390 final F<A, Integer> one = constant(1); 391 return foldMap(one, Monoid.intAdditionMonoid); 392 } 393 394 /** 395 * Splits this set at the given element. Returns a product-3 of: 396 * <ul> 397 * <li>A set containing all the elements of this set which are less than the given value.</li> 398 * <li>An option of a value equal to the given value, if one was found in this set, otherwise None. 399 * <li>A set containing all the elements of this set which are greater than the given value.</li> 400 * </ul> 401 * 402 * @param a A value at which to split this set. 403 * @return Two sets and an optional value, where all elements in the first set are less than the given value 404 * and all the elements in the second set are greater than the given value, and the optional value is the 405 * given value if found, otherwise None. 406 */ 407 public P3<Set<A>, Option<A>, Set<A>> split(final A a) { 408 if (isEmpty()) 409 return P.p(empty(ord), Option.<A>none(), empty(ord)); 410 else { 411 final A h = head(); 412 final Ordering i = ord.compare(a, h); 413 if (i == LT) { 414 final P3<Set<A>, Option<A>, Set<A>> lg = l().split(a); 415 return P.p(lg._1(), lg._2(), lg._3().insert(h).union(r())); 416 } else if (i == GT) { 417 final P3<Set<A>, Option<A>, Set<A>> lg = r().split(a); 418 return P.p(lg._1().insert(h).union(l()), lg._2(), lg._3()); 419 } else 420 return P.p(l(), some(h), r()); 421 } 422 } 423 424 /** 425 * Returns true if this set is a subset of the given set. 426 * 427 * @param s A set which is a superset of this set if this method returns true. 428 * @return true if this set is a subset of the given set. 429 */ 430 public boolean subsetOf(final Set<A> s) { 431 if (isEmpty() || s.isEmpty()) 432 return isEmpty(); 433 else { 434 final P3<Set<A>, Option<A>, Set<A>> find = s.split(head()); 435 return find._2().isSome() && l().subsetOf(find._1()) && r().subsetOf(find._3()); 436 } 437 } 438 439 /** 440 * Join a set of sets into a single set. 441 * 442 * @param s A set of sets. 443 * @param o An order for the elements of the new set. 444 * @return A new set which is the join of the given set of sets. 445 */ 446 public static <A> Set<A> join(final Ord<A> o, final Set<Set<A>> s) { 447 final F<Set<A>, Set<A>> id = identity(); 448 return s.foldMap(id, Monoid.<A>setMonoid(o)); 449 } 450 451 /** 452 * Return the elements of the given iterable as a set. 453 * 454 * @param o An order for the elements of the new set. 455 * @param as An iterable of elements to add to a set. 456 * @return A new set containing the elements of the given iterable. 457 */ 458 public static <A> Set<A> iterableSet(final Ord<A> o, final Iterable<A> as) { 459 Set<A> s = empty(o); 460 for (final A a : as) 461 s = s.insert(a); 462 return s; 463 } 464 465 }