diff options
author | Paul Phillips <paulp@improving.org> | 2010-04-07 19:23:08 +0000 |
---|---|---|
committer | Paul Phillips <paulp@improving.org> | 2010-04-07 19:23:08 +0000 |
commit | 6dd51419b8a2565db0ef1d6f36cb57b5ef7654cf (patch) | |
tree | 37fbdacc88c9c3fde2bc414c145fb5c02d707a0f | |
parent | 1b098c643a8df4cdb5e6d0fd84619575ed830266 (diff) | |
download | scala-6dd51419b8a2565db0ef1d6f36cb57b5ef7654cf.tar.gz scala-6dd51419b8a2565db0ef1d6f36cb57b5ef7654cf.tar.bz2 scala-6dd51419b8a2565db0ef1d6f36cb57b5ef7654cf.zip |
Gave Stream a lazy withFilter implementation.
can have a collection containing all the even numbers in the universe
and still be home in time for tea. Threw in some Stream cleanups for
free. Closes #3265, review by community.
-rw-r--r-- | src/library/scala/collection/immutable/Stream.scala | 104 | ||||
-rw-r--r-- | test/files/run/streamWithFilter.check | 5 | ||||
-rw-r--r-- | test/files/run/streamWithFilter.scala | 11 |
3 files changed, 90 insertions, 30 deletions
diff --git a/src/library/scala/collection/immutable/Stream.scala b/src/library/scala/collection/immutable/Stream.scala index 3b10963ddb..dc1f2b6536 100644 --- a/src/library/scala/collection/immutable/Stream.scala +++ b/src/library/scala/collection/immutable/Stream.scala @@ -97,6 +97,12 @@ self => loop(this, "") } + /** It's an imperfect world, but at least we can bottle up the + * imperfection in a capsule. + */ + @inline private def asThat[That](x: AnyRef): That = x.asInstanceOf[That] + @inline private def asStream[B](x: AnyRef): Stream[B] = x.asInstanceOf[Stream[B]] + // Overridden methods from Traversable override def toStream: Stream[A] = this @@ -113,21 +119,23 @@ self => * then StreamBuilder will be chosen for the implicit. * we recognize that fact and optimize to get more laziness. */ - override def ++[B >: A, That](that: TraversableOnce[B])(implicit bf: CanBuildFrom[Stream[A], B, That]): That = { + override def ++[B >: A, That](that: TraversableOnce[B])(implicit bf: CanBuildFrom[Stream[A], B, That]): That = // we assume there is no other builder factory on streams and therefore know that That = Stream[A] - (if (isEmpty) that.toStream - else new Stream.Cons(head, (tail ++ that).asInstanceOf[Stream[A]])).asInstanceOf[That] - } + asThat[That]( + if (isEmpty) that.toStream + else new Stream.Cons(head, asStream[A](tail ++ that)) + ) /** * Create a new stream which contains all intermediate results of applying the operator * to subsequent elements left to right. * @note This works because the target type of the Builder That is a Stream. */ - override final def scanLeft[B, That](z: B)(op: (B, A) => B)(implicit bf: CanBuildFrom[Stream[A], B, That]): That = { - (if (this.isEmpty) Stream(z) - else new Stream.Cons(z, tail.scanLeft(op(z, head))(op).asInstanceOf[Stream[B]])).asInstanceOf[That] - } + override final def scanLeft[B, That](z: B)(op: (B, A) => B)(implicit bf: CanBuildFrom[Stream[A], B, That]): That = + asThat[That]( + if (isEmpty) Stream(z) + else new Stream.Cons(z, asStream[B](tail.scanLeft(op(z, head))(op))) + ) /** Returns the stream resulting from applying the given function * <code>f</code> to each element of this stream. @@ -136,10 +144,11 @@ self => * @return <code>f(a<sub>0</sub>), ..., f(a<sub>n</sub>)</code> if this * sequence is <code>a<sub>0</sub>, ..., a<sub>n</sub></code>. */ - override final def map[B, That](f: A => B)(implicit bf: CanBuildFrom[Stream[A], B, That]): That = { - (if (isEmpty) Stream.Empty - else new Stream.Cons(f(head), (tail map f).asInstanceOf[Stream[B]])).asInstanceOf[That] - } + override final def map[B, That](f: A => B)(implicit bf: CanBuildFrom[Stream[A], B, That]): That = + asThat[That]( + if (isEmpty) Stream.Empty + else new Stream.Cons(f(head), asStream[B](tail map f)) + ) /** Applies the given function <code>f</code> to each element of * this stream, then concatenates the results. @@ -152,20 +161,22 @@ self => // we assume there is no other builder factory on streams and therefore know that That = Stream[B] // optimisations are not for speed, but for functionality // see tickets #153, #498, #2147, and corresponding tests in run/ (as well as run/stream_flatmap_odds.scala) - (if (isEmpty) Stream.Empty - else { - // establish !prefix.isEmpty || nonEmptyPrefix.isEmpty - var nonEmptyPrefix = this - var prefix = f(nonEmptyPrefix.head).toStream - while (!nonEmptyPrefix.isEmpty && prefix.isEmpty) { - nonEmptyPrefix = nonEmptyPrefix.tail - if(!nonEmptyPrefix.isEmpty) - prefix = f(nonEmptyPrefix.head).toStream + asThat[That]( + if (isEmpty) Stream.Empty + else { + // establish !prefix.isEmpty || nonEmptyPrefix.isEmpty + var nonEmptyPrefix = this + var prefix = f(nonEmptyPrefix.head).toStream + while (!nonEmptyPrefix.isEmpty && prefix.isEmpty) { + nonEmptyPrefix = nonEmptyPrefix.tail + if(!nonEmptyPrefix.isEmpty) + prefix = f(nonEmptyPrefix.head).toStream + } + + if (nonEmptyPrefix.isEmpty) Stream.empty + else prefix append asStream[B](nonEmptyPrefix.tail flatMap f) } - - if(nonEmptyPrefix.isEmpty) Stream.empty - else prefix append (nonEmptyPrefix.tail flatMap f).asInstanceOf[Stream[B]] - }).asInstanceOf[That] + ) /** Returns all the elements of this stream that satisfy the * predicate <code>p</code>. The order of the elements is preserved. @@ -180,6 +191,37 @@ self => else new Stream.Cons(rest.head, rest.tail filter p) } + override final def withFilter(p: A => Boolean): StreamWithFilter = new StreamWithFilter(p) + + /** A lazier implementation of WithFilter than TraversableLike's. + */ + final class StreamWithFilter(p: A => Boolean) extends WithFilter(p) { + override def map[B, That](f: A => B)(implicit bf: CanBuildFrom[Stream[A], B, That]): That = { + def tailMap = asStream[B](tail withFilter p map f) + asThat[That]( + if (isEmpty) Stream.Empty + else if (p(head)) new Stream.Cons(f(head), tailMap) + else tailMap + ) + } + + override def flatMap[B, That](f: A => Traversable[B])(implicit bf: CanBuildFrom[Stream[A], B, That]): That = { + def tailFlatMap = asStream[B](tail withFilter p flatMap f) + asThat[That]( + if (isEmpty) Stream.Empty + else if (p(head)) f(head).toStream append tailFlatMap + else tailFlatMap + ) + } + + override def foreach[B](f: A => B) = + for (x <- self) + if (p(x)) f(x) + + override def withFilter(q: A => Boolean): StreamWithFilter = + new StreamWithFilter(x => p(x) && q(x)) + } + /** Apply the given function <code>f</code> to each element of this linear sequence * (while respecting the order of the elements). * @@ -224,11 +266,12 @@ self => * <code>Stream(a<sub>0</sub>, ..., a<sub>m</sub>) * zip Stream(b<sub>0</sub>, ..., b<sub>n</sub>)</code> is invoked. */ - override final def zip[A1 >: A, B, That](that: Iterable[B])(implicit bf: CanBuildFrom[Stream[A], (A1, B), That]): That = { + override final def zip[A1 >: A, B, That](that: Iterable[B])(implicit bf: CanBuildFrom[Stream[A], (A1, B), That]): That = // we assume there is no other builder factory on streams and therefore know that That = Stream[(A1, B)] - (if (this.isEmpty || that.isEmpty) Stream.Empty - else new Stream.Cons((this.head, that.head), (this.tail zip that.tail).asInstanceOf[Stream[(A1, B)]])).asInstanceOf[That] - } + asThat[That]( + if (this.isEmpty || that.isEmpty) Stream.Empty + else new Stream.Cons((this.head, that.head), asStream[(A1, B)](this.tail zip that.tail)) + ) /** Zips this iterable with its indices. `s.zipWithIndex` is equivalent to * `s zip s.indices` @@ -359,7 +402,8 @@ self => def loop(len: Int, these: Stream[A]): Stream[B] = if (these.isEmpty) Stream.fill(len)(elem) else new Stream.Cons(these.head, loop(len - 1, these.tail)) - loop(len, this).asInstanceOf[That] + + asThat[That](loop(len, this)) // was: if (bf.isInstanceOf[Stream.StreamCanBuildFrom[_]]) loop(len, this).asInstanceOf[That] // else super.padTo(len, elem) } diff --git a/test/files/run/streamWithFilter.check b/test/files/run/streamWithFilter.check new file mode 100644 index 0000000000..6b0e91a147 --- /dev/null +++ b/test/files/run/streamWithFilter.check @@ -0,0 +1,5 @@ +15 +30 +45 +60 +75 diff --git a/test/files/run/streamWithFilter.scala b/test/files/run/streamWithFilter.scala new file mode 100644 index 0000000000..cb919d4f55 --- /dev/null +++ b/test/files/run/streamWithFilter.scala @@ -0,0 +1,11 @@ +object Test { + val nums = Stream.from(1) + def isFizz(x: Int) = x % 3 == 0 + def isBuzz(x: Int) = x % 5 == 0 + // next line will run forever if withFilter isn't doing its thing. + val fizzbuzzes = for (n <- nums ; if isFizz(n) ; if isBuzz(n)) yield n + + def main(args: Array[String]): Unit = { + fizzbuzzes take 5 foreach println + } +} |