001    package fj.data;
002    
003    import static fj.Function.compose;
004    import fj.*;
005    import static fj.F2W.$$;
006    import static fj.FW.$;
007    import static fj.P.p;
008    import static fj.data.IterableW.join;
009    import static fj.data.List.iterableList;
010    import static fj.data.Option.none;
011    import fj.pre.Ord;
012    
013    import java.util.Iterator;
014    import java.util.Map;
015    
016    /**
017     * An immutable, in-memory map, backed by a red-black tree.
018     */
019    public final class TreeMap<K, V> implements Iterable<P2<K, V>> {
020      private final Set<P2<K, Option<V>>> tree;
021    
022      private TreeMap(final Set<P2<K, Option<V>>> tree) {
023        this.tree = tree;
024      }
025    
026      private static <K, V> Ord<P2<K, V>> ord(final Ord<K> keyOrd) {
027        return keyOrd.comap(P2.<K, V>__1());
028      }
029    
030      /**
031       * Constructs an empty tree map.
032       *
033       * @param keyOrd An order for the keys of the tree map.
034       * @return an empty TreeMap with the given key order.
035       */
036      public static <K, V> TreeMap<K, V> empty(final Ord<K> keyOrd) {
037        return new TreeMap<K, V>(Set.empty(TreeMap.<K, Option<V>>ord(keyOrd)));
038      }
039    
040      /**
041       * Returns a potential value that the given key maps to.
042       *
043       * @param k The key to look up in the tree map.
044       * @return A potential value for the given key.
045       */
046      public Option<V> get(final K k) {
047        final Option<P2<K, Option<V>>> x = tree.split(P.p(k, Option.<V>none()))._2();
048        return x.bind(P2.<K, Option<V>>__2());
049      }
050    
051      /**
052       * Inserts the given key and value association into the tree map.
053       * If the given key is already mapped to a value, the old value is replaced with the given one.
054       *
055       * @param k The key to insert.
056       * @param v The value to insert.
057       * @return A new tree map with the given value mapped to the given key.
058       */
059      public TreeMap<K, V> set(final K k, final V v) {
060        final P3<Set<P2<K, Option<V>>>, Option<P2<K, Option<V>>>, Set<P2<K, Option<V>>>> x
061            = tree.split(P.p(k, Option.<V>none()));
062        return new TreeMap<K, V>(x._1().union(x._3().insert(P.p(k, Option.some(v)))));
063      }
064    
065      /**
066       * Deletes the entry in the tree map that corresponds to the given key.
067       *
068       * @param k The key to delete from this tree map.
069       * @return A new tree map with the entry corresponding to the given key removed.
070       */
071      public TreeMap<K, V> delete(final K k) {
072        return new TreeMap<K, V>(tree.delete(P.p(k, Option.<V>none())));
073      }
074    
075      /**
076       * Returns the number of entries in this tree map.
077       *
078       * @return The number of entries in this tree map.
079       */
080      public int size() {
081        return tree.size();
082      }
083    
084      /**
085       * Determines if this tree map has any entries.
086       *
087       * @return <code>true</code> if this tree map has no entries, <code>false</code> otherwise.
088       */
089      public boolean isEmpty() {
090        return tree.isEmpty();
091      }
092    
093      /**
094       * Returns all values in this tree map.
095       *
096       * @return All values in this tree map.
097       */
098      public List<V> values() {
099        return iterableList(join(tree.toList().map(compose(IterableW.<V, Option<V>>wrap(), P2.<K, Option<V>>__2()))));
100      }
101    
102      /**
103       * Returns all keys in this tree map.
104       *
105       * @return All keys in this tree map.
106       */
107      public List<K> keys() {
108        return tree.toList().map(P2.<K, Option<V>>__1());
109      }
110    
111      /**
112       * Determines if the given key value exists in this tree map.
113       *
114       * @param k The key value to look for in this tree map.
115       * @return <code>true</code> if this tree map contains the given key, <code>false</code> otherwise.
116       */
117      public boolean contains(final K k) {
118        return tree.member(P.p(k, Option.<V>none()));
119      }
120    
121      /**
122       * Returns an iterator for this map's key-value pairs.
123       * This method exists to permit the use in a <code>for</code>-each loop.
124       *
125       * @return A iterator for this map's key-value pairs.
126       */
127      public Iterator<P2<K, V>> iterator() {
128        return join(tree.toStream().map(P2.<K, Option<V>, IterableW<V>>map2_(IterableW.<V, Option<V>>wrap())
129        ).map(P2.tuple(compose(IterableW.<V, P2<K, V>>map(), P.<K, V>p2())))).iterator();
130      }
131    
132      /**
133       * A mutable map projection of this tree map.
134       *
135       * @return A new mutable map isomorphic to this tree map.
136       */
137      public Map<K, V> toMutableMap() {
138        final Map<K, V> m = new java.util.TreeMap<K, V>();
139        for (final P2<K, V> e : this) {
140          m.put(e._1(), e._2());
141        }
142        return m;
143      }
144    
145      /**
146       * An immutable projection of the given mutable map.
147       *
148       * @param ord An order for the map's keys.
149       * @param m   A mutable map to project to an immutable one.
150       * @return A new immutable tree map isomorphic to the given mutable map.
151       */
152      public static <K, V> TreeMap<K, V> fromMutableMap(final Ord<K> ord, final Map<K, V> m) {
153        TreeMap<K, V> t = empty(ord);
154        for (final Map.Entry<K, V> e : m.entrySet()) {
155          t = t.set(e.getKey(), e.getValue());
156        }
157        return t;
158      }
159    
160      /**
161       * Returns a first-class version of the get method for this TreeMap.
162       *
163       * @return a functional representation of this TreeMap.
164       */
165      public F<K, Option<V>> get() {
166        return new F<K, Option<V>>() {
167          public Option<V> f(final K k) {
168            return get(k);
169          }
170        };
171      }
172    
173      /**
174       * Modifies the value for the given key, if present, by applying the given function to it.
175       *
176       * @param k The key for the value to modify.
177       * @param f A function with which to modify the value.
178       * @return A new tree map with the value for the given key transformed by the given function,
179       *         paired with True if the map was modified, otherwise False.
180       */
181      public P2<Boolean, TreeMap<K, V>> update(final K k, final F<V, V> f) {
182        final P2<Boolean, Set<P2<K, Option<V>>>> up =
183            tree.update(p(k, Option.<V>none()), P2.<K, Option<V>, Option<V>>map2_(Option.<V, V>map().f(f)));
184        return P.p(up._1(), new TreeMap<K, V>(up._2()));
185      }
186    
187      /**
188       * Modifies the value for the given key, if present, by applying the given function to it, or
189       * inserts the given value if the key is not present.
190       *
191       * @param k The key for the value to modify.
192       * @param f A function with which to modify the value.
193       * @param v A value to associate with the given key if the key is not already present.
194       * @return A new tree map with the value for the given key transformed by the given function.
195       */
196      public TreeMap<K, V> update(final K k, final F<V, V> f, final V v) {
197        final P2<Boolean, TreeMap<K, V>> up = update(k, f);
198        return up._1() ? up._2() : set(k, v);
199      }
200    
201      /**
202       * Splits this TreeMap at the given key. Returns a triple of:
203       * <ul>
204       * <li>A set containing all the values of this map associated with keys less than the given key.</li>
205       * <li>An option of a value mapped to the given key, if it exists in this map, otherwise None.
206       * <li>A set containing all the values of this map associated with keys greater than the given key.</li>
207       * </ul>
208       *
209       * @param k A key at which to split this map.
210       * @return Two sets and an optional value, where all elements in the first set are mapped to keys less than the given
211       *         key in this map, all the elements in the second set are mapped to keys greater than the given key,
212       *         and the optional value is the value associated with the given key if present, otherwise None.
213       */
214      public P3<Set<V>, Option<V>, Set<V>> split(final K k) {
215        final FW<Set<P2<K, Option<V>>>, Set<V>> getSome = $(Option.<V>fromSome()).o(P2.<K, Option<V>>__2())
216            .mapSet(tree.ord().comap($$(P.<K, Option<V>>p2()).f(k).<V>o(Option.<V>some_())));
217        return tree.split(p(k, Option.<V>none())).map1(getSome).map3(getSome)
218            .map2($(Option.<V>join()).o($(P2.<K, Option<V>>__2()).mapOption()));
219      }
220    
221      /**
222       * Maps the given function across the values of this TreeMap.
223       *
224       * @param f A function to apply to the values of this TreeMap.
225       * @return A new TreeMap with the values transformed by the given function.
226       */
227      public <W> TreeMap<K, W> map(final F<V, W> f) {
228        final F<P2<K, Option<V>>, P2<K, Option<W>>> g = P2.map2_($(f).mapOption());
229        final FW<K, P2<K, Option<V>>> coord = $$(P.<K, Option<V>>p2()).flip().f(Option.<V>none());
230        final Ord<K> o = tree.ord().comap(coord);
231        return new TreeMap<K, W>(tree.map(TreeMap.<K, Option<W>>ord(o), g));
232      }
233    
234    }