Merge pull request functionaljava#110 from dobesv/to-mutable-map-comparator

mperry · mperry · commit ce693701579a · 2015-03-22T19:41:44.000+10:00
Preserve ordering function when converting to a TreeMap
diff --git a/core/src/main/java/fj/Ord.java b/core/src/main/java/fj/Ord.java
@@ -12,6 +12,7 @@
 
 import java.math.BigDecimal;
 import java.math.BigInteger;
+import java.util.Comparator;
 
 import static fj.Function.curry;
 
@@ -503,4 +504,14 @@ public static <A> Ord<A> hashEqualsOrd() {
     });
   }
 
+  class OrdComparator implements Comparator<A> {
+	@Override
+    public int compare(A o1, A o2) {
+	    return Ord.this.compare(o1, o2).toInt();
+    }
+  }
+
+  public Comparator<A> toComparator() {
+	  return new OrdComparator();
+  }
 }
diff --git a/core/src/main/java/fj/data/TreeMap.java b/core/src/main/java/fj/data/TreeMap.java
@@ -2,6 +2,7 @@
 
 import fj.*;
 
+import java.util.Comparator;
 import java.util.Iterator;
 import java.util.Map;
 
@@ -172,7 +173,9 @@ public Iterator<P2<K, V>> iterator() {
    * @return A new mutable map isomorphic to this tree map.
    */
   public Map<K, V> toMutableMap() {
-    final Map<K, V> m = new java.util.TreeMap<K, V>();
+    final F<K, P2<K, Option<V>>> fakePair = k -> P.p(k, Option.none());
+	final Comparator<K> comparator = tree.ord().comap(fakePair).toComparator();
+	final Map<K, V> m = new java.util.TreeMap<K, V>(comparator);
     for (final P2<K, V> e : this) {
       m.put(e._1(), e._2());
     }
diff --git a/core/src/test/java/fj/data/TreeMapTest.java b/core/src/test/java/fj/data/TreeMapTest.java
@@ -1,13 +1,17 @@
 package fj.data;
 
+import java.util.Map;
+
 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.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 /**
@@ -74,4 +78,12 @@ public void splitLookup() {
         assertTrue(eq.eq(p3, P.p(leftMap, some(Integer.toString(pivot)), rightMap)));
     }
 
+    @Test
+    public void toMutableMap() {
+        int max = 5;
+        List<List<Integer>> l = List.range(1, max + 1).map(n -> List.single(n));
+        TreeMap<List<Integer>, String> m2 = TreeMap.treeMap(Ord.listOrd(Ord.intOrd), l.zip(l.map(i -> i.toString())));
+        Map<List<Integer>, String> mm = m2.toMutableMap();
+        assertEquals(m2.keys().reverse(), List.iterableList(mm.keySet()));
+    }
 }
