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    }