Writing generic takeWhile with an offset in Scala

Question

I have a use case for takeWhile, but where I want to keep a fixed number of items after predicate is true. I'm also trying to write it to be as generic as possible for the type of collection. So something like:

def takeWhileWithOffset[A, Iter[_] <: Iterable[A]](iter: Iter[A], p: A => Boolean, offset: Int)

I chose Iterable as the bound, because want it to work with Stream. But I'm having a tough time figuring out how to make this work. If I were using a strict collection, I could use dropRight if offset isn't positive. But Iterable doesn't have dropRight.

The positive case is trickier. I could use sliding to effectively grab future items and then use init and lastOption after the takeWhile terminates. But Iterable doesn't have init and lastOption.

So what's tricky is that I want my method to be lazy, yet take advantage of the fact that iterated items can be treated as a strict collection -- but only if the takeWhile terminates. Is there a way to do this?

Answer 1

span and take:

scala> val it = (1 to 100).iterator
it: Iterator[Int] = non-empty iterator                              

scala> val (a, b) = it.span(_ < 10)
a: Iterator[Int] = non-empty iterator
b: Iterator[Int] = unknown-if-empty iterator

scala> val res = a ++ b.take(5)
res: Iterator[Int] = non-empty iterator

scala> res.toList
res0: List[Int] = List(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14)