/core/src/main/scala/scalaz/StreamT.scala

  1package scalaz
  2
  3import Id._
  4
  5/**
  6 * StreamT monad transformer.
  7 */
  8sealed class StreamT[M[_], A](val step: M[StreamT.Step[A, StreamT[M, A]]]) {
  9
 10  import StreamT._
 11
 12  def uncons(implicit M: Monad[M]): M[Option[(A, StreamT[M, A])]] =
 13    M.bind(step) {
 14      case Yield(a, s) => M.point(Some((a, s())))
 15      case Skip(s)     => s().uncons
 16      case Done()      => M.point(None)
 17    }
 18
 19  def unconsRec(implicit M: BindRec[M]): M[Option[(A, StreamT[M, A])]] =
 20    M.tailrecM(this)(s =>
 21      M.map(s.step) {
 22        case Yield(a, s1) => \/-(Some((a, s1())))
 23        case Skip(s1)     => -\/(s1())
 24        case Done()       => \/-(None)
 25      })
 26
 27  def ::(a: A)(implicit M: Applicative[M]): StreamT[M, A] = StreamT[M, A](M.point(Yield(a, this)))
 28
 29  def isEmpty(implicit M: Monad[M]): M[Boolean] = M.map(uncons)(_.isEmpty)
 30  def isEmptyRec(implicit M: BindRec[M]): M[Boolean] = M.map(unconsRec)(_.isEmpty)
 31
 32  def head(implicit M: Monad[M]): M[A] = M.map(uncons)(_.getOrElse(sys.error("head: empty StreamT"))._1)
 33  def headRec(implicit M: BindRec[M]): M[A] = M.map(unconsRec)(_.getOrElse(sys.error("head: empty StreamT"))._1)
 34
 35  def headOption(implicit M: Monad[M]): M[Option[A]] = M.map(uncons)(_.map(_._1))
 36  def headOptionRec(implicit M: BindRec[M]): M[Option[A]] = M.map(unconsRec)(_.map(_._1))
 37
 38  def tailM(implicit M: Monad[M]): M[StreamT[M, A]] = M.map(uncons)(_.getOrElse(sys.error("tailM: empty StreamT"))._2)
 39  def tailMRec(implicit M: BindRec[M]): M[StreamT[M, A]] = M.map(unconsRec)(_.getOrElse(sys.error("tailM: empty StreamT"))._2)
 40
 41  def tailOption(implicit M: Monad[M]): M[Option[StreamT[M, A]]] =
 42    M.map(uncons)(_.map(_._2))
 43  def tailOptionRec(implicit M: BindRec[M]): M[Option[StreamT[M, A]]] =
 44    M.map(unconsRec)(_.map(_._2))
 45
 46  def trans[N[_]](t: M ~> N)(implicit M: Functor[M], N: Functor[N]): StreamT[N, A] =
 47    StreamT(t(M.map[Step[A, StreamT[M, A]], Step[A, StreamT[N, A]]](this.step) {
 48      case Yield(a, s) => Yield(a, s() trans t)
 49      case Skip(s)     => Skip(s() trans t)
 50      case Done()      => Done()
 51    }))
 52
 53  def filter(p: A => Boolean)(implicit m: Functor[M]): StreamT[M, A] = stepMap[A] {
 54    case Yield(a, s) => if (p(a)) Yield(a, s() filter p) else Skip(s() filter p)
 55    case Skip(s)     => Skip(s() filter p)
 56    case Done()      => Done()
 57  }
 58
 59  def drop(n: Int)(implicit M: Functor[M]): StreamT[M, A] = stepMap[A] {
 60    case Yield(a, s) => if (n > 0) Skip(s() drop (n - 1)) else Yield(a, s)
 61    case Skip(s)     => Skip(s() drop n)
 62    case Done()      => Done()
 63  }
 64
 65  def dropWhile(p: A => Boolean)(implicit m: Functor[M]): StreamT[M, A] = stepMap[A] {
 66    case Yield(a, s) => if (p(a)) Skip(s() dropWhile p) else Yield(a, s)
 67    case Skip(s)     => Skip(s() dropWhile p)
 68    case Done()      => Done()
 69  }
 70
 71  def take(n: Int)(implicit M: Functor[M]): StreamT[M, A] = stepMap[A] {
 72    case Yield(a, s) => if (n <= 0) Done() else Yield(a, s() take (n - 1))
 73    case Skip(s)     => Skip(s() take n)
 74    case Done()      => Done()
 75  }
 76
 77  def takeWhile(p: A => Boolean)(implicit m: Functor[M]): StreamT[M, A] = stepMap[A] {
 78    case Yield(a, s) => if (p(a)) Yield(a, s() takeWhile p) else Done()
 79    case Skip(s)     => Skip(s() takeWhile p)
 80    case Done()      => Done()
 81  }
 82
 83  def ++(bs: => StreamT[M, A])(implicit m: Functor[M]): StreamT[M, A] = stepMap[A] {
 84    case Yield(a, s) => Yield(a, s() ++ bs)
 85    case Skip(s)     => Skip(s() ++ bs)
 86    case Done()      => Skip(bs)
 87  }
 88
 89  def flatMap[B](f: A => StreamT[M, B])(implicit m: Functor[M]): StreamT[M, B] = stepMap[B] {
 90    case Yield(a, s) => Skip(f(a) ++ (s() flatMap f))
 91    case Skip(s)     => Skip(s() flatMap f)
 92    case Done()      => Done()
 93  }
 94
 95  def map[B](f: A => B)(implicit m: Functor[M]): StreamT[M, B] = stepMap[B] {
 96    case Yield(a, s) => Yield(f(a), s() map f)
 97    case Skip(s)     => Skip(s() map f)
 98    case Done()      => Done()
 99  }
100
101  /** @since 7.0.1 */
102  def mapM[B](f: A => M[B])(implicit M: Monad[M]): StreamT[M, B] = stepBind[B] {
103    case Yield(a, s) => M.map(f(a))(Yield(_, s() mapM f))
104    case Skip(s)     => M.point(Skip(s() mapM f))
105    case Done()      => M.point(Done())
106  }
107
108  def foldLeft[B](z: B)(f: (B, A) => B)(implicit M: Monad[M]): M[B] =
109    M.bind(step) {
110      case Yield(a, s) => s().foldLeft(f(z, a))(f)
111      case Skip(s)     => s().foldLeft(z)(f)
112      case Done()      => M.point(z)
113    }
114
115  def foldLeftRec[B](z: B)(f: (B, A) => B)(implicit M: BindRec[M]): M[B] =
116    M.tailrecM((() => this, z))(sb => M.map(sb._1().step) {
117      case Yield(a, s) => -\/((s, f(sb._2, a)))
118      case Skip(s)     => -\/((s, sb._2))
119      case Done()      => \/-(sb._2)
120    })
121
122  /**
123   * **Warning:** Requires evaluation of the whole stream. Depending on
124   * the monad `M`, the evaluation will happen either immediately, or
125   * will be deferred until the resulting `LazyList` is extracted from the
126   * returned `M`.
127   */
128  def toLazyList(implicit M: Monad[M]): M[LazyList[A]] = M.map(rev)(_.reverse)
129
130  /**
131   * **Warning:** Requires evaluation of the whole stream. Depending on
132   * the monad `M`, the evaluation will happen either immediately, or
133   * will be deferred until the resulting `LazyList` is extracted from the
134   * returned `M`.
135   */
136  def toLazyListRec(implicit M: BindRec[M]): M[LazyList[A]] = M.map(revRec)(_.reverse)
137
138  /**
139   * Converts this `StreamT` to a lazy `LazyList`, i.e. without forcing
140   * evaluation of all elements. Note, however, that at least one element
141   * of this stream will be evaluated, and depending on the structure of
142   * this stream, up to two elements might be evaluated.
143   */
144  def asLazyList(implicit ev: M[Step[A, StreamT[M, A]]] === Id[Step[A, StreamT[Id, A]]]): LazyList[A] = {
145    def go(s: StreamT[Id, A]): LazyList[A] = s.unconsRec match {
146      case None          => LazyList.empty[A]
147      case Some((a, s1)) => LazyList.cons(a, go(s1))
148    }
149
150    go(StreamT(ev(step)))
151  }
152
153  def foldRight[B](z: => B)(f: (=> A, => B) => B)(implicit M: Monad[M]): M[B] =
154    M.map(rev) {
155      _.foldLeft(z)((a, b) => f(b, a))
156    }
157
158  def foldRightRec[B](z: => B)(f: (=> A, => B) => B)(implicit M: BindRec[M]): M[B] =
159    M.map(revRec) {
160      _.foldLeft(z)((a, b) => f(b, a))
161    }
162
163  /**
164   * `foldRight` with potential to terminate early, e.g. on an infinite stream.
165   */
166  def foldRightM[B](z: => M[B])(f: (=> A, => M[B]) => M[B])(implicit M: Monad[M]): M[B] =
167    M.bind(step) {
168      case Yield(a, s) => f(a, s().foldRightM(z)(f))
169      case Skip(s)     => s().foldRightM(z)(f)
170      case Done()      => z
171    }
172
173  def foldMap[B](f: A => B)(implicit M: Foldable[M], B: Monoid[B]): B =
174    M.foldMap(step) {
175      case Yield(a, s) => B.append(f(a), s().foldMap(f))
176      case Skip(s)     => s().foldMap(f)
177      case Done()      => B.zero
178    }
179
180  def length(implicit m: Monad[M]): M[Int] =
181    foldLeft(0)((c, a) => 1 + c)
182
183  def lengthRec(implicit M: BindRec[M]): M[Int] =
184    foldLeftRec(0)((c, a) => 1 + c)
185
186  def foreach(f: A => M[Unit])(implicit M: Monad[M]): M[Unit] = M.bind(step) {
187    case Yield(a,s) => M.bind(f(a))(_ => s() foreach f)
188    case Skip(s)    => s().foreach(f)
189    case Done()     => M.pure(())
190  }
191
192  def foreachRec(f: A => M[Unit])(implicit M: Monad[M], B: BindRec[M]): M[Unit] =
193    B.tailrecM(() => this)(s => M.bind(s().step) {
194      case Yield(a, s1) => M.map(f(a))(_ => -\/(s1))
195      case Skip(s1)     => M.pure(-\/(s1))
196      case Done()       => M.pure(\/-(()))
197    })
198
199  private def stepMap[B](f: Step[A, StreamT[M, A]] => Step[B, StreamT[M, B]])(implicit M: Functor[M]): StreamT[M, B] = StreamT(M.map(step)(f))
200
201  private def stepBind[B](f: Step[A, StreamT[M, A]] => M[Step[B, StreamT[M, B]]])(implicit M: Monad[M]): StreamT[M, B] = StreamT(M.bind(step)(f))
202
203  private def rev(implicit M: Monad[M]): M[LazyList[A]] = {
204    def loop(xs: StreamT[M, A], ys: LazyList[A]): M[LazyList[A]] =
205      M.bind(xs.step) {
206        case Yield(a, s) => loop(s(), a #:: ys)
207        case Skip(s)     => loop(s(), ys)
208        case Done()      => M.point(ys)
209      }
210    loop(this, LazyList.empty)
211  }
212
213  private def revRec(implicit M: BindRec[M]): M[LazyList[A]] =
214    M.tailrecM((() => this, LazyList.empty[A])) { case (xs, ys) =>
215      M.map(xs().step) {
216        case Yield(a, s) => -\/((s, a #:: ys))
217        case Skip(s)     => -\/((s, ys))
218        case Done()      => \/-(ys)
219      }
220    }
221
222}
223
224//
225// Prioritized Implicits for type class instances
226//
227
228sealed abstract class StreamTInstances0 {
229  implicit def StreamTInstance1[F[_]](implicit F0: Functor[F]): Bind[StreamT[F, *]] with Plus[StreamT[F, *]] =
230    new StreamTInstance1[F] {
231      implicit def F: Functor[F] = F0
232    }
233
234  implicit def StreamTSemigroup[F[_], A](implicit F0: Functor[F]): Semigroup[StreamT[F, A]] =
235    new StreamTSemigroup[F, A] {
236      implicit def F: Functor[F] = F0
237    }
238}
239
240sealed abstract class StreamTInstances extends StreamTInstances0 {
241  implicit def StreamTMonoid[F[_], A](implicit F0: Applicative[F]): Monoid[StreamT[F, A]] =
242    new StreamTMonoid[F, A] {
243      implicit def F: Applicative[F] = F0
244    }
245  implicit def StreamTMonadPlus[F[_]](implicit F0: Applicative[F]): MonadPlus[StreamT[F, *]] =
246    new StreamTMonadPlus[F] {
247      implicit def F: Applicative[F] = F0
248    }
249  implicit def StreamTEqual[F[_], A](implicit E: Equal[F[LazyList[A]]], F: Monad[F]): Equal[StreamT[F, A]] = E.contramap((_: StreamT[F, A]).toLazyList)
250  implicit def StreamTShow[F[_], A](implicit E: Show[F[LazyList[A]]], F: Monad[F]): Show[StreamT[F, A]] = Contravariant[Show].contramap(E)((_: StreamT[F, A]).toLazyList)
251  implicit val StreamTHoist: Hoist[StreamT] = new StreamTHoist {}
252  implicit def StreamTFoldable[F[_]: Foldable]: Foldable[StreamT[F, *]] =
253    new Foldable[StreamT[F, *]] with Foldable.FromFoldMap[StreamT[F, *]] {
254      override def foldMap[A, M: Monoid](s: StreamT[F, A])(f: A => M) = s.foldMap(f)
255    }
256}
257
258object StreamT extends StreamTInstances {
259  def apply[M[_], A](step: M[Step[A, StreamT[M, A]]]): StreamT[M, A] = new StreamT[M, A](step)
260
261  def empty[M[_], A](implicit M: Applicative[M]): StreamT[M, A] = new StreamT[M, A](M point Done())
262
263  def fromLazyList[M[_], A](mas: M[LazyList[A]])(implicit M: Applicative[M]): StreamT[M, A] = {
264    def loop(as: LazyList[A]): Step[A, StreamT[M, A]] = as match {
265      case head #:: tail => Yield(head, apply(M.point(loop(tail))))
266      case _ => Done()
267    }
268
269    apply[M, A](M.map(mas)(loop))
270  }
271
272  def unfoldM[M[_],A,B](start: B)(f: B => M[Option[(A,B)]])(implicit M: Functor[M]): StreamT[M,A] =
273    StreamT[M,A](M.map(f(start)) {
274      case Some((a, b)) => Yield(a, unfoldM(b)(f))
275      case None => Done()
276    })
277
278  def unfold[A,B](b: B)(f: B => Option[(A,B)]): StreamT[Id,A] = unfoldM[Id,A,B](b)(f)
279
280  def fromIterable[A](s: Iterable[A]): StreamT[Id,A] = {
281    def stepper(b: Iterable[A]): Option[(A,Iterable[A])] = if (b.isEmpty) None else Some((b.head, b.tail))
282    unfold(s)(stepper)
283  }
284
285  def wrapEffect[M[_]:Functor,A](m: M[StreamT[M,A]]): StreamT[M,A] = StreamT(Functor[M].map(m)(Skip(_)))
286
287  def runStreamT[S,A](stream : StreamT[State[S, *],A], s0: S): StreamT[Id,A] =
288    StreamT[Id,A]({
289      val (a, s) = stream.step(s0)
290      s match {
291        case Yield(a1, s1) => Yield(a1, runStreamT(s1(), a))
292        case Skip(s1)      => Skip(runStreamT(s1(), a))
293        case Done()        => Done()
294      }
295    })
296
297
298  sealed abstract class Step[A, S] extends Product with Serializable
299
300  final case class Yield[A, S](a: A, s: () => S) extends Step[A, S]
301  object Yield {
302    def apply[A, S](a: A, s: => S): Step[A, S] = new Yield(a, () => s)
303  }
304  final case class Skip[A, S](s: () => S) extends Step[A, S]
305  object Skip {
306    def apply[A, S](s: => S): Step[A, S] = new Skip(() => s)
307  }
308  sealed abstract case class Done[A, S]() extends Step[A, S]
309  object Done {
310    def apply[A, S](): Step[A, S] = done_.asInstanceOf[Done[A, S]]
311    // https://github.com/scala/bug/issues/11953
312    private[this] final val done_ : Done[Nothing, Nothing] = new Done[Nothing, Nothing]{}
313  }
314}
315
316//
317// Implementation traits for type class instances
318//
319
320private trait StreamTInstance1[F[_]] extends Bind[StreamT[F, *]] with Plus[StreamT[F, *]] {
321  implicit def F: Functor[F]
322
323  override final def map[A, B](fa: StreamT[F, A])(f: A => B) =
324    fa map f
325
326  override final def bind[A, B](fa: StreamT[F, A])(f: A => StreamT[F, B]) =
327    fa flatMap f
328
329  override final def plus[A](a: StreamT[F, A], b: => StreamT[F, A]) =
330    a ++ b
331}
332
333private trait StreamTSemigroup[F[_], A] extends Semigroup[StreamT[F, A]] {
334  implicit def F: Functor[F]
335
336  def append(f1: StreamT[F, A], f2: => StreamT[F, A]): StreamT[F, A] = f1 ++ f2
337}
338
339private trait StreamTMonoid[F[_], A] extends Monoid[StreamT[F, A]] with StreamTSemigroup[F, A] {
340  implicit def F: Applicative[F]
341
342  def zero: StreamT[F, A] = StreamT.empty[F, A]
343}
344
345private trait StreamTMonadPlus[F[_]] extends MonadPlus[StreamT[F, *]] with StreamTInstance1[F] {
346  implicit def F: Applicative[F]
347
348  def point[A](a: => A): StreamT[F, A] = a :: StreamT.empty[F, A]
349
350  def empty[A]: StreamT[F, A] = StreamT.empty
351}
352
353private trait StreamTHoist extends Hoist[StreamT] {
354  import StreamT._
355
356  implicit def apply[G[_] : Monad]: Monad[StreamT[G, *]] = StreamTMonadPlus[G]
357
358  def liftM[G[_], A](a: G[A])(implicit G: Monad[G]): StreamT[G, A] =
359    StreamT[G, A](G.map(a)(Yield(_, empty[G, A](G))))
360
361  override def wrapEffect[G[_]: Monad, A](a: G[StreamT[G, A]]): StreamT[G, A] = StreamT.wrapEffect(a)
362
363  def hoist[M[_], N[_]](f: M ~> N)(implicit M: Monad[M]): StreamT[M, *] ~> StreamT[N, *] =
364    λ[StreamT[M, *] ~> StreamT[N, *]](a =>
365      StreamT(f(M.map(a.step) {
366        case Yield(a, s) => Yield(a, hoist(f).apply(s()))
367        case Skip(s)     => Skip(hoist(f).apply(s()))
368        case Done()      => Done()
369      }
370    )))
371}