diff options
-rw-r--r-- | src/library/scala/collection/immutable/Stream.scala | 32 | ||||
-rw-r--r-- | test/files/run/streams.check | 1 | ||||
-rw-r--r-- | test/files/run/streams.scala | 5 |
3 files changed, 31 insertions, 7 deletions
diff --git a/src/library/scala/collection/immutable/Stream.scala b/src/library/scala/collection/immutable/Stream.scala index 78c4d76eda..5566806c55 100644 --- a/src/library/scala/collection/immutable/Stream.scala +++ b/src/library/scala/collection/immutable/Stream.scala @@ -181,6 +181,7 @@ import scala.language.implicitConversions * @define coll stream * @define orderDependent * @define orderDependentFold + * @define willTerminateInf Note: lazily evaluated; will terminate for infinite-sized collections. */ abstract class Stream[+A] extends AbstractSeq[A] with LinearSeq[A] @@ -286,9 +287,8 @@ self => len } - /** It's an imperfect world, but at least we can bottle up the - * imperfection in a capsule. - */ + // 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]] @inline private def isStreamBuilder[B, That](bf: CanBuildFrom[Stream[A], B, That]) = @@ -725,10 +725,15 @@ self => * // produces: "5, 6, 7, 8, 9" * }}} */ - override def take(n: Int): Stream[A] = + override def take(n: Int): Stream[A] = ( + // Note that the n == 1 condition appears redundant but is not. + // It prevents "tail" from being referenced (and its head being evaluated) + // when obtaining the last element of the result. Such are the challenges + // of working with a lazy-but-not-really sequence. if (n <= 0 || isEmpty) Stream.empty else if (n == 1) cons(head, Stream.empty) else cons(head, tail take n-1) + ) @tailrec final override def drop(n: Int): Stream[A] = if (n <= 0 || isEmpty) this @@ -784,8 +789,23 @@ self => these } - // there's nothing we can do about dropRight, so we just keep the definition - // in LinearSeq + /** + * @inheritdoc + * $willTerminateInf + */ + override def dropRight(n: Int): Stream[A] = { + // We make dropRight work for possibly infinite streams by carrying + // a buffer of the dropped size. As long as the buffer is full and the + // rest is non-empty, we can feed elements off the buffer head. When + // the rest becomes empty, the full buffer is the dropped elements. + def advance(stub0: List[A], stub1: List[A], rest: Stream[A]): Stream[A] = { + if (rest.isEmpty) Stream.empty + else if (stub0.isEmpty) advance(stub1.reverse, Nil, rest) + else cons(stub0.head, advance(stub0.tail, rest.head :: stub1, rest.tail)) + } + if (n <= 0) this + else advance((this take n).toList, Nil, this drop n) + } /** Returns the longest prefix of this `Stream` whose elements satisfy the * predicate `p`. diff --git a/test/files/run/streams.check b/test/files/run/streams.check index 7f894052d9..032057d4a1 100644 --- a/test/files/run/streams.check +++ b/test/files/run/streams.check @@ -23,3 +23,4 @@ Stream(100001, ?) true true 705082704 +6 diff --git a/test/files/run/streams.scala b/test/files/run/streams.scala index 51b4e5d76c..dc5d0204ac 100644 --- a/test/files/run/streams.scala +++ b/test/files/run/streams.scala @@ -29,7 +29,7 @@ object Test extends App { def powers(x: Int) = if ((x&(x-1)) == 0) Some(x) else None println(s3.flatMap(powers).reverse.head) - // large enough to generate StackOverflows (on most systems) + // large enough to generate StackOverflows (on most systems) // unless the following methods are tail call optimized. val size = 100000 @@ -43,4 +43,7 @@ object Test extends App { println(Stream.from(1).take(size).foldLeft(0)(_ + _)) val arr = new Array[Int](size) Stream.from(1).take(size).copyToArray(arr, 0) + + // dropRight terminates + println(Stream from 1 dropRight 1000 take 3 sum) } |