From 393ecf2208bd2770bd4dc8efa1279ed488ad922a Mon Sep 17 00:00:00 2001
From: James Strachan <james.strachan@gmail.com>
Date: Sat, 24 Dec 2011 07:46:26 +0000
Subject: [PATCH] added a groupBy method (and stumbled on a couple more little
 compiler bugs which could well be the same as previous issues I've raised
 this week). Happy holidays everyone :)

---
 stdlib/ktSrc/JavaIterables.kt  | 14 ++++++++++++++
 testlib/test/CollectionTest.kt | 18 ++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/stdlib/ktSrc/JavaIterables.kt b/stdlib/ktSrc/JavaIterables.kt
index 438b29dc388..e209250b1fc 100644
--- a/stdlib/ktSrc/JavaIterables.kt
+++ b/stdlib/ktSrc/JavaIterables.kt
@@ -81,6 +81,20 @@ inline fun <T> java.lang.Iterable<T>.foreach(operation: (element: T) -> Unit) {
     operation(elem)
 }
 
+/**
+ * Iterates through the collection performing the transformation on each element and using the result
+ * as the key in a map to group elements by the result
+ */
+inline fun <T,K> java.lang.Iterable<T>.groupBy(result: Map<K,List<T>> = HashMap<K,List<T>>(), toKey: (T)-> K) : Map<K,List<T>> {
+  for (elem in this) {
+    val key = toKey(elem)
+    val list = result.getOrElseUpdate(key){ ArrayList<T>() }
+    list.add(elem)
+  }
+  return result
+}
+
+
 /** Creates a String from all the elements in the collection, using the seperator between them and using the given prefix and postfix if supplied */
 inline fun <T> java.lang.Iterable<T>.join(separator: String, prefix: String = "", postfix: String = "") : String {
   val buffer = StringBuilder(prefix)
diff --git a/testlib/test/CollectionTest.kt b/testlib/test/CollectionTest.kt
index 926c300862c..e26e264f8c5 100644
--- a/testlib/test/CollectionTest.kt
+++ b/testlib/test/CollectionTest.kt
@@ -118,6 +118,24 @@ class CollectionTest() : TestSupport() {
     assertEquals(6, count)
   }
 
+  fun testGroupBy() {
+    val words = arrayList("a", "ab", "abc", "def", "abcd")
+    /*
+     TODO inference engine should not need this type info?
+     */
+    val byLength = words.groupBy<String,Int>{it.length}
+    assertEquals(4, byLength.size())
+
+    println("Grouped by length is: $byLength")
+    /*
+     TODO compiler bug...
+
+    val l3 = byLength.getOrElse(3, {ArrayList<String>()})
+    assertEquals(2, l3.size)
+    */
+
+  }
+
   fun testJoin() {
     val text = data.join("-", "<", ">")
     assertEquals("<foo-bar>", text)