Merge pull request scala#7124 from joshlemer/tap-each

[11098] Add IterableOps#tapEach method

Author: SethTisue

src/library/scala/collection/Iterable.scala

@@ -819,6 +819,8 @@ trait IterableOps[+A, +CC[_], +C] extends Any with IterableOnce[A] with Iterable
     */
   def inits: Iterator[C] = iterateUntilEmpty(_.init)
 
+  override def tapEach[U](f: A => U): C = fromSpecific(new View.Map(this, { a: A => f(a); a }))
+
   // A helper for tails and inits.
   private[this] def iterateUntilEmpty(f: Iterable[A] => Iterable[A]): Iterator[C] = {
     val it = Iterator.iterate(toIterable)(f).takeWhile(x => !x.isEmpty)

src/library/scala/collection/IterableOnce.scala

@@ -458,6 +458,17 @@ trait IterableOnceOps[+A, +CC[_], +C] extends Any { this: IterableOnce[A] =>
     */
   def span(p: A => Boolean): (C, C)
 
+  /** Applies a side-effecting function to each element in this collection.
+    * Strict collections will apply `f` to their elements immediately, while lazy collections
+    * like Views and LazyLists will only apply `f` on each element if and when that element
+    * is evaluated, and each time that element is evaluated.
+    *
+    * @param f a function to apply to each element in this $coll
+    * @tparam U the return type of f
+    * @return The same logical collection as this
+    */
+  def tapEach[U](f: A => U): C
+
   /////////////////////////////////////////////////////////////// Concrete methods based on iterator
 
   def knownSize: Int = -1

src/library/scala/collection/Iterator.scala

@@ -863,6 +863,16 @@ trait Iterator[+A] extends IterableOnce[A] with IterableOnceOps[A, Iterator, Ite
       }
     }
 
+  override def tapEach[U](f: A => U): Iterator[A] = new AbstractIterator[A] {
+    override def knownSize = self.knownSize
+    override def hasNext = self.hasNext
+    override def next() = {
+      val _next = self.next()
+      f(_next)
+      _next
+    }
+  }
+
   /** Converts this iterator to a string.
    *
    *  @return `"<iterator>"`

src/library/scala/collection/StrictOptimizedIterableOps.scala

@@ -249,4 +249,10 @@ trait StrictOptimizedIterableOps[+A, +CC[_], +C]
     (l.result(), r.result())
   }
 
+  // Optimization avoids creation of second collection
+  override def tapEach[U](f: A => U): C  = {
+    foreach(f)
+    coll
+  }
+
 }

src/library/scala/collection/immutable/LazyList.scala

@@ -380,6 +380,8 @@ final class LazyList[+A] private(private[this] var lazyState: () => LazyList.Sta
     if (knownIsEmpty) LazyList.empty
     else (mapImpl(f): @inline)
 
+  override def tapEach[U](f: A => U): LazyList[A] = map { a => f(a); a}
+
   private def mapImpl[B](f: A => B): LazyList[B] =
     newLL {
       if (isEmpty) State.Empty

test/junit/scala/collection/IteratorTest.scala

@@ -535,24 +535,33 @@ class IteratorTest {
     // Avoid reaching seq1 through test class. Avoid testing Array.iterator.
     class C extends Iterable[String] {
       val ss = Array("first", "second")
+
       def iterator = new Iterator[String] {
         var i = 0
+
         def hasNext = i < ss.length
+
         def next() =
-          if (hasNext) { val res = ss(i) ; i += 1 ; res }
+          if (hasNext) {
+            val res = ss(i); i += 1; res
+          }
           else Iterator.empty.next()
       }
+
       def apply(i: Int) = ss(i)
     }
     val seq1 = new WeakReference(new C)
     val seq2 = List("third")
     val it0: Iterator[Int] = Iterator(1, 2)
     lazy val it: Iterator[String] = it0.flatMap {
       case 1 => seq1.get
-      case _ => check() ; seq2
+      case _ => check(); seq2
     }
+
     def check() = assertNotReachable(seq1.get, it)(())
+
     def checkHasElement() = assertNotReachable(seq1.get.apply(1), it)(())
+
     assert(it.hasNext)
     assertEquals("first", it.next())
 
@@ -568,4 +577,27 @@ class IteratorTest {
     }
     assert(!it.hasNext)
   }
+
+  @Test def tapEach(): Unit = {
+    locally {
+      var i = 0
+      val tapped = Iterator(-1, -1, -1).tapEach(_ => i += 1)
+      assertEquals(true, tapped.hasNext)
+      assertEquals(0, i)
+    }
+
+    locally {
+      var i = 0
+      val tapped = Iterator(-1, -1, -1).tapEach(_ => i += 1)
+      assertEquals(-3, tapped.sum)
+      assertEquals(3, i)
+    }
+
+    locally {
+      var i = 0
+      val tapped = Iterator(-1, -1, -1).tapEach(_ => i += 1)
+      assertEquals(-1, tapped.next())
+      assertEquals(1, i)
+    }
+  }
 }

test/junit/scala/collection/ViewTest.scala

@@ -1,12 +1,13 @@
 package scala.collection
 
 import scala.collection.immutable.List
-
 import org.junit.Assert._
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
+
 import language.postfixOps
+import scala.collection.mutable.ListBuffer
 
 @RunWith(classOf[JUnit4])
 class ViewTest {
@@ -17,7 +18,8 @@ class ViewTest {
 
     import scala.language.postfixOps
     assertEquals(Iterable.empty[Int], iter.view take Int.MinValue to Iterable)
-    assertEquals(Iterable.empty[Int], iter.view takeRight Int.MinValue to Iterable)
+    assertEquals(Iterable.empty[Int],
+                 iter.view takeRight Int.MinValue to Iterable)
     assertEquals(iter, iter.view drop Int.MinValue to Iterable)
     assertEquals(iter, iter.view dropRight Int.MinValue to Iterable)
   }
@@ -69,4 +71,22 @@ class ViewTest {
     check(immutable.Map(1 -> "a", 2 -> "b")) // MapView
   }
 
+  @Test
+  def tapEach: Unit = {
+    val lb = ListBuffer[Int]()
+
+    val v =
+      View(1, 2, 3)
+        .tapEach(lb += _)
+        .map(_ => 10)
+        .tapEach(lb += _)
+        .tapEach(_ => lb += -1)
+
+    assertEquals(ListBuffer[Int](), lb)
+
+    val strict = v.to(Seq)
+    assertEquals(strict, Seq(10, 10, 10))
+    assertEquals(lb, Seq(1, 10, -1, 2, 10, -1, 3, 10, -1))
+  }
+
 }

test/junit/scala/collection/immutable/LazyListTest.scala

@@ -6,6 +6,7 @@ import org.junit.Test
 import org.junit.Assert._
 
 import scala.collection.Iterator
+import scala.collection.mutable.ListBuffer
 import scala.ref.WeakReference
 import scala.util.Try
 
@@ -358,6 +359,15 @@ class LazyListTest {
     assertLazyHeadWhenNextHeadEvaluated(op)
   }
 
+  @Test
+  def tapEach_properlyLazy(): Unit = {
+    val op = lazyListOp(_.tapEach(_ + 1))
+    assertRepeatedlyFullyLazy(op)
+    assertLazyTailWhenHeadEvaluated(op)
+    assertLazyHeadWhenTailEvaluated(op)
+    assertLazyHeadWhenNextHeadEvaluated(op)
+  }
+
   @Test
   def collect_properlyLazy(): Unit = {
     val op = lazyListOp(_ collect { case i if i % 2 != 0 => i})
@@ -1000,4 +1010,24 @@ class LazyListTest {
       assertEquals(s"mkString 3 $i", goal(3,i), precyc(3,i).mkString)
     }
   }
+  @Test
+  def tapEach: Unit = {
+
+    /** @param makeLL must make a lazylist that evaluates to Seq(1,2,3,4,5) */
+    def check(makeLL: => LazyList[Int]): Unit = {
+      val lb = ListBuffer[Int]()
+      val ll = makeLL.tapEach(lb += _)
+      assertEquals(ListBuffer[Int](), lb)
+      assertEquals(Vector(1, 2), ll.take(2).to(Vector))
+      assertEquals(ListBuffer(1, 2), lb)
+      assertEquals(4, ll(3))
+      assertEquals(ListBuffer(1, 2, 4), lb)
+      assertEquals(Vector(1,2,3,4,5), ll.to(Vector))
+      assertEquals(ListBuffer(1, 2, 4, 3, 5), lb)
+    }
+
+    check(LazyList.from(Iterator(1, 2, 3, 4, 5)))
+    check(LazyList.from(Vector(1, 2, 3, 4, 5)))
+    check(LazyList.tabulate(5)(_ + 1))
+  }
 }

test/junit/scala/collection/immutable/VectorTest.scala

@@ -5,6 +5,8 @@ import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
 import org.junit.Test
 
+import scala.collection.mutable.ListBuffer
+
 @RunWith(classOf[JUnit4])
 class VectorTest {
 
@@ -161,4 +163,22 @@ class VectorTest {
     assertEquals(-1, test.indexOf(1000))
   }
 
+  @Test
+  def tapEach(): Unit = {
+    val lb = ListBuffer[Int]()
+
+    val v =
+      Vector(1,2,3)
+      .tapEach(lb += _)
+      .tapEach(lb += _)
+
+    assertEquals(ListBuffer(1,2,3,1,2,3), lb)
+    assertEquals(Vector(1,2,3), v)
+
+
+    val f: Any => Unit = println
+
+    // test that type info is not lost
+    val x: Vector[Char] = Vector[Char]().tapEach(f)
+  }
 }