functionaljava · mperry · Jan 11, 2015 · Jan 11, 2015 · Jan 11, 2015 · Jan 11, 2015
diff --git a/core/src/main/java/fj/Equal.java b/core/src/main/java/fj/Equal.java
@@ -522,6 +522,10 @@ public static <A> Equal<Set<A>> setEqual(final Equal<A> e) {
     return equal(curry((Set<A> a, Set<A> b) -> streamEqual(e).eq(a.toStream(), b.toStream())));
   }
 
+  public static <K, V> Equal<TreeMap<K, V>> treeMapEqual(Equal<K> k, Equal<V> v) {
+    return equal(t1 -> t2 -> Equal.streamEqual(p2Equal(k, v)).eq(t1.toStream(), t2.toStream()));
+  }
+
   public static <A, B> Equal<Writer<A, B>> writerEqual(Equal<A> eq1, Equal<B> eq2) {
     return equal(w1 -> w2 -> p2Equal(eq1, eq2).eq(w1.run(), w2.run()));
   }

diff --git a/core/src/main/java/fj/data/Set.java b/core/src/main/java/fj/data/Set.java
@@ -474,4 +474,18 @@ public static <A> Set<A> set(final Ord<A> o, final A ... as) {
     return s;
   }
 
+  /**
+   * Constructs a set from the given elements.
+   *
+   * @param o  An order for the elements of the new set.
+   * @param list The elements to add to a set.
+   * @return A new set containing the elements of the given list.
+   */
+  public static <A> Set<A> set(final Ord<A> o, List<A> list) {
+    Set<A> s = empty(o);
+    for (final A a : list)
+      s = s.insert(a);
+    return s;
+  }
+
 }
diff --git a/core/src/main/java/fj/data/TreeMap.java b/core/src/main/java/fj/data/TreeMap.java
@@ -40,6 +40,32 @@ public static <K, V> TreeMap<K, V> empty(final Ord<K> keyOrd) {
     return new TreeMap<K, V>(Set.empty(TreeMap.<K, Option<V>>ord(keyOrd)));
   }
 
+  /**
+   * Constructs a tree map from the given elements.
+   *
+   * @param keyOrd An order for the keys of the tree map.
+   * @param p2s The elements to construct the tree map with.
+   * @return a TreeMap with the given elements.
+   */
+  public static <K, V> TreeMap<K, V> treeMap(final Ord<K> keyOrd, final P2<K, V>... p2s) {
+    return treeMap(keyOrd, List.list(p2s));
+  }
+
+  /**
+   * Constructs a tree map from the given elements.
+   *
+   * @param keyOrd An order for the keys of the tree map.
+   * @param list The elements to construct the tree map with.
+   * @return a TreeMap with the given elements.
+   */
+  public static <K, V> TreeMap<K, V> treeMap(final Ord<K> keyOrd, final List<P2<K, V>> list) {
+    TreeMap<K, V> tm = empty(keyOrd);
+    for (final P2<K, V> p2 : list) {
+      tm = tm.set(p2._1(), p2._2());
+    }
+    return tm;
+  }
+
   /**
    * Returns a potential value that the given key maps to.
    *
@@ -143,6 +169,9 @@ public Map<K, V> toMutableMap() {
     return m;
   }
 
+  public Stream<P2<K, V>> toStream() {
+    return Stream.iteratorStream(iterator());
+  }
   /**
    * An immutable projection of the given mutable map.
    *
@@ -208,13 +237,50 @@ public TreeMap<K, V> update(final K k, final F<V, V> f, final V v) {
    *         key in this map, all the elements in the second set are mapped to keys greater than the given key,
    *         and the optional value is the value associated with the given key if present, otherwise None.
    */
-  public P3<Set<V>, Option<V>, Set<V>> split(final K k) {
-    final F<Set<P2<K, Option<V>>>, Set<V>> getSome = F1Functions.mapSet(F1Functions.o(Option.<V>fromSome(), P2.<K, Option<V>>__2())
-        , tree.ord().comap(F1Functions.o(P.<K, Option<V>>p2().f(k), Option.<V>some_())));
+  public P3<Set<V>, Option<V>, Set<V>> split(Ord<V> ord, final K k) {
+    final F<Set<P2<K, Option<V>>>, Set<V>> getSome = F1Functions.mapSet(F1Functions.o(Option.<V>fromSome(), P2.<K, Option<V>>__2()), ord);
     return tree.split(p(k, Option.<V>none())).map1(getSome).map3(getSome)
         .map2(F1Functions.o(Option.<V>join(), F1Functions.mapOption(P2.<K, Option<V>>__2())));
   }
 
+  /**
+   * Internal construction of a TreeMap from the given set.
+   * @param ord An order for the keys of the tree map.
+   * @param s The elements to construct the tree map with.
+   * @return a TreeMap with the given elements.
+   */
+  private static <K, V> TreeMap<K, V> treeMap(Ord<K> ord, Set<P2<K, Option<V>>> s) {
+    TreeMap<K, V> empty = TreeMap.<K, V>empty(ord);
+    TreeMap<K, V> tree = s.toList().foldLeft((tm, p2) -> {
+      Option<V> opt = p2._2();
+      if (opt.isSome()) {
+        return tm.set(p2._1(), opt.some());
+      }
+      return tm;
+    }, empty);
+    return tree;
+  }
+
+  /**
+   * Splits this TreeMap at the given key. Returns a triple of:
+   * <ul>
+   * <li>A tree map containing all the values of this map associated with keys less than the given key.</li>
+   * <li>An option of a value mapped to the given key, if it exists in this map, otherwise None.
+   * <li>A tree map containing all the values of this map associated with keys greater than the given key.</li>
+   * </ul>
+   *
+   * @param k A key at which to split this map.
+   * @return Two tree maps and an optional value, where all keys in the first tree map are mapped
+   * to keys less than the given key in this map, all the keys in the second tree map are mapped
+   * to keys greater than the given key, and the optional value is the value associated with the
+   * given key if present, otherwise None.
+   */
+  public P3<TreeMap<K, V>, Option<V>, TreeMap<K, V>> splitLookup(final K k) {
+    P3<Set<P2<K, Option<V>>>, Option<P2<K, Option<V>>>, Set<P2<K, Option<V>>>> p3 = tree.split(P.p(k, get(k)));
+    Ord<K> o = tree.ord().<K>comap(k2 -> P.p(k2, Option.<V>none()));
+    return P.p(treeMap(o, p3._1()), get(k), treeMap(o, p3._3()));
+  }
+
   /**
    * Maps the given function across the values of this TreeMap.
    *

diff --git a/core/src/test/java/fj/data/TreeMapTest.java b/core/src/test/java/fj/data/TreeMapTest.java
@@ -0,0 +1,77 @@
+package fj.data;
+
+import fj.Equal;
+import fj.Ord;
+import fj.P3;
+import fj.Show;
+import fj.P;
+import org.junit.Test;
+
+import static fj.data.Option.some;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * Created by MarkPerry on 11/01/2015.
+ */
+public class TreeMapTest {
+
+    @Test
+    public void split() {
+        // do the split
+        int pivot = 4;
+        int max = 5;
+        List<Integer> l = List.range(1, max + 1);
+        TreeMap<Integer, String> m2 = TreeMap.treeMap(Ord.intOrd, l.zip(l.map(i -> i.toString())));
+        P3<Set<String>, Option<String>, Set<String>> p = m2.split(Ord.stringOrd, pivot);
+
+        // print debug info
+        Show<TreeMap<Integer, String>> st = Show.treeMapShow(Show.intShow, Show.stringShow);
+        Show<Set<String>> ss = Show.setShow(Show.stringShow);
+        Show<Option<String>> so = Show.optionShow(Show.stringShow);
+        Show<P3<Set<String>, Option<String>, Set<String>>> sp3 = Show.p3Show(ss, so, ss);
+        if (true) {
+            st.println(m2);
+            sp3.println(p);
+        }
+
+        // assert equals
+        Equal<Set<String>> seq = Equal.setEqual(Equal.stringEqual);
+        Set<String> left = toSetString(List.range(1, pivot));
+        Set<String> right = toSetString(List.range(pivot + 1, max + 1));
+        P3<Set<String>, Option<String>, Set<String>> expected = P.p(left, some(Integer.toString(pivot)), right);
+        assertTrue(Equal.p3Equal(seq, Equal.optionEqual(Equal.stringEqual), seq).eq(p, expected));
+    }
+
+    private static Set<String> toSetString(List<Integer> list) {
+        return Set.set(Ord.stringOrd, list.map(i -> i.toString()));
+    }
+
+    @Test
+    public void splitLookup() {
+        // do the split
+        int pivot = 4;
+        int max = 5;
+        List<Integer> l = List.range(1, max + 1);
+        TreeMap<Integer, String> m2 = TreeMap.treeMap(Ord.intOrd, l.zip(l.map(i -> i.toString())));
+        P3<TreeMap<Integer, String>, Option<String>, TreeMap<Integer, String>> p3 = m2.splitLookup(pivot);
+
+        // create expected output
+        List<Integer> leftList = List.range(1, pivot);
+        TreeMap<Integer, String> leftMap = TreeMap.treeMap(Ord.intOrd, leftList.zip(leftList.map(i -> i.toString())));
+        List<Integer> rightList = List.range(pivot + 1, max + 1);
+        TreeMap<Integer, String> rightMap = TreeMap.treeMap(Ord.intOrd, rightList.zip(rightList.map(i -> i.toString())));
+
+        // debug info
+        if (true) {
+            Show<TreeMap<Integer, String>> st = Show.treeMapShow(Show.intShow, Show.stringShow);
+            Show<P3<TreeMap<Integer, String>, Option<String>, TreeMap<Integer, String>>> sp3 = Show.p3Show(st, Show.optionShow(Show.stringShow), st);
+            sp3.println(p3);
+        }
+
+        // do the assert
+        Equal<TreeMap<Integer, String>> tme = Equal.treeMapEqual(Equal.intEqual, Equal.stringEqual);
+        Equal<P3<TreeMap<Integer, String>, Option<String>, TreeMap<Integer, String>>> eq = Equal.p3Equal(tme, Equal.optionEqual(Equal.stringEqual), tme);
+        assertTrue(eq.eq(p3, P.p(leftMap, some(Integer.toString(pivot)), rightMap)));
+    }
+
+}