Tail recursive implementation of mapConserve, s...

Tail recursive implementation of mapConserve, submitted by Eric
Willigers. Closes #2411, review by malayeri.

2 files changed, 74 insertions, 18 deletions

diff --git a/src/library/scala/collection/immutable/List.scala b/src/library/scala/collection/immutable/List.scala
index 7785d73175..93c3134fd7 100644
--- a/src/library/scala/collection/immutable/List.scala
+++ b/src/library/scala/collection/immutable/List.scala
@@ -107,29 +107,32 @@ sealed abstract class List[+A] extends LinearSeq[A]
    *                `f` to each element of this list and collecting the results.
    *  @usecase def mapConserve(f: A => A): List[A]
    */
-  def mapConserve[B >: A <: AnyRef] (f: A => B): List[B] = {
-    def loop(ys: List[A]): List[B] =
-      if (ys.isEmpty) this
+  def mapConserve[B >: A <: AnyRef](f: A => B): List[B] = {
+    @tailrec
+    def loop(mapped: ListBuffer[B], unchanged: List[A], pending: List[A]): List[B] =
+      if (pending.isEmpty) {
+        if (mapped eq null) unchanged
+        else mapped.prependToList(unchanged)
+      }
       else {
-        val head0 = ys.head
+        val head0 = pending.head
         val head1 = f(head0)
-        if (head1 eq head0.asInstanceOf[AnyRef]) {
-          loop(ys.tail)
-        } else {
-          val ys1 = head1 :: ys.tail.mapConserve(f)
-          if (this eq ys) ys1
-          else {
-            val b = new ListBuffer[B]
-            var xc = this
-            while (xc ne ys) {
-              b += xc.head
-              xc = xc.tail
-            }
-            b.prependToList(ys1)
+
+        if (head1 == head0)
+          loop(mapped, unchanged, pending.tail)
+        else {
+          val b = if (mapped eq null) new ListBuffer[B] else mapped
+          var xc = unchanged
+          while (xc ne pending) {
+            b += xc.head
+            xc = xc.tail
           }
+          b += head1
+          val tail0 = pending.tail
+          loop(b, tail0, tail0)
         }
       }
-    loop(this)
+    loop(null, this, this)
   }
 
   // Overridden methods from IterableLike and SeqLike or overloaded variants of such methods
diff --git a/test/files/run/mapConserve.scala b/test/files/run/mapConserve.scala
new file mode 100644
index 0000000000..a285113b11
--- /dev/null
+++ b/test/files/run/mapConserve.scala
@@ -0,0 +1,53 @@
+import scala.annotation.tailrec
+import scala.collection.mutable.ListBuffer
+
+object Test {
+  val maxListLength = 7 // up to 16, but larger is slower
+  var testCount = 0
+
+  def checkStackOverflow() = {
+    var xs: List[String] = Nil
+    for (i <- 0 until 250000)
+        xs = "X" :: xs
+
+    val lowers = xs.mapConserve(_.toLowerCase)
+    assert(xs.mapConserve(x => x) eq xs)
+  }
+
+  def checkBehaviourUnchanged(input: List[_], oldOutput: List[_], newOutput: List[_]) {
+    if (oldOutput eq input)
+      assert(newOutput eq oldOutput)
+    else {
+      assert(newOutput.head == oldOutput.head)
+      checkBehaviourUnchanged(input.tail, oldOutput.tail, newOutput.tail)
+    }
+    testCount += 1
+  }
+
+    var callCount = 0
+    val lastHexDigit: Function1[BigInt, AnyRef] = { x: BigInt => callCount+=1; if (x < 16) x else x % 16 }
+
+    def main(args: Array[String]) {
+        for (length <- 0 to maxListLength;
+             bitmap <- 0 until (1 << length);
+             data = List.range(0, length) map { x: Int =>
+               if ((bitmap & (1 << x)) != 0) BigInt(x+16)
+               else BigInt(x)
+             })
+        {
+            // Behaves like map with respect to  ==
+            callCount = 0
+            val numUnconserved = data.reverse.dropWhile(_ < 16).length
+            val result = data mapConserve lastHexDigit
+            val mapResult = data map lastHexDigit
+            assert(result == mapResult)
+            assert((result drop numUnconserved) eq (data drop numUnconserved))
+            assert(callCount == 2 * length) // map, mapConserve call transform for each element in the list
+
+            // Behaves like existing mapConserve with respect to  eq
+            checkBehaviourUnchanged(data, data mapConserve lastHexDigit, data mapConserve lastHexDigit)
+        }
+
+        checkStackOverflow();
+    }
+}
\ No newline at end of file