Stateful monad loops in Scala

mloop.sc

import $ivy.`org.typelevel::cats-core:2.0.0`
import $plugin.$ivy.`org.typelevel::kind-projector:0.10.3`

import cats._
import cats.instances.option._
import cats.syntax.all._

import scala.annotation.tailrec

def next(offset: Int): Option[Vector[Int]] = if (offset < 50) Vector.tabulate(10)(identity).some else Vector.empty[Int].some

@tailrec
def loopM[F[_]: Monad, A](init: F[A])(condition: F[A] => Boolean)(body: A => F[A]): F[A] = {
  val step = init.flatMap(body)
  if (condition(step)) loopM(step)(condition)(body) else step
}

@tailrec
def loopM2[F[_]: Monad, A, R](init: F[A], condition: F[R] => Boolean, recharge: (A, R) => A)(body: A => F[R]): F[R] = {
  val step = init.flatMap(body)
  if (condition(step)) loopM2((init, step).mapN(recharge), condition, recharge)(body) else step
}

// Also stack-safe, because monads in cats are stack-safe
def loopM[F[_]: Monad, A, R](init: A, condition: R => Boolean, recharge: (A, R) => A)(body: A => F[R]): F[R] = {
  val step = body(init)
  step
    .map(condition)
    .ifM(
      step.map(recharge(init, _)).flatMap(loopM(_, condition, recharge)(body)),
      step
    )
}

/*
loopM(0.some)(_.exists(_ > 0)) { offset =>
  for {
    fetched <- next(offset)
    _ = println(fetched)
  } yield fetched.size
}
 */

loopM2(0.some, (fetched: Option[Int]) => fetched.exists(_ > 0), (offset: Int, fetched: Int) => offset + fetched) { offset: Int =>
  for {
    fetched <- next(offset)
    _ = println(fetched)
  } yield fetched.size
}