Skip to content

Instantly share code, notes, and snippets.

@johnynek
Last active July 14, 2021 02:50
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save johnynek/07e5f4b12eaa7d6f3cb833359a2c7588 to your computer and use it in GitHub Desktop.
Save johnynek/07e5f4b12eaa7d6f3cb833359a2c7588 to your computer and use it in GitHub Desktop.
An example of a Free-monad which also includes ability to do general recursion.
package ioloop
object IOLoop {
enum Res[+A] {
case Done(a: A)
case FlatMap[A, B](prev: Res[B], fn: B => Res[A]) extends Res[A]
case Recurse[A, B](arg: A, loop: (A => Res[B]) => (A => Res[B])) extends Res[B]
def flatMap[B](fn: A => Res[B]): Res[B] =
FlatMap(this, fn)
def map[B](fn: A => B): Res[B] = flatMap { a => Done(fn(a)) }
def run: A = IOLoop.run(this)
}
object Res {
val Unit: Res[Unit] = Res(())
def apply[A](a: A): Res[A] = Res.Done(a)
def defer[A](fa: => Res[A]): Res[A] = Unit.flatMap(_ => fa)
def delay[A](a: => A): Res[A] = defer(Done(a))
def loop[A, B](fn: (A => Res[B]) => (A => Res[B])): A => Res[B] =
{ (a: A) => Res.Recurse(a, fn) }
def tailRecM[A, B](init: A)(fn: A => Res[Either[A, B]]): Res[B] =
(loop[Either[A, B], B] { rec =>
{
case Left(a) => fn(a).flatMap(rec)
case Right(b) => Res(b)
}
})(Left(init))
}
enum Stack[-A, B] {
case Ident[A, B](ev: A <:< B) extends Stack[A, B]
case App[A, B, C](first: A => Res[B], rest: Stack[B, C]) extends Stack[A, C]
}
def identStack[A]: Stack[A, A] = Stack.Ident(summon[A <:< A])
def run[A](res: Res[A]): A = {
// Inside this method is the only place where we actually do recursion
// outside of here we can do recursion without explicitly using it
@annotation.tailrec
def loop[B](left: Res[B], stack: Stack[B, A]): A =
left match {
case Res.Done(a) =>
stack match {
case Stack.App(first, rest) => loop(first(a), rest)
case Stack.Ident(ev) => ev(a)
}
case Res.FlatMap(prev, fn) =>
// by matching here we can avoid allocating App just to remove it
prev match {
case Res.Done(a) => loop(fn(a), stack)
case _ => loop(prev, Stack.App(fn, stack))
}
case Res.Recurse(arg, loopFn) =>
// fix f = f (fix f) (thanks to Steven Noble who noticed a nice simplification here)
val step = loopFn(Res.loop(loopFn))
loop(step(arg), stack)
}
loop(res, identStack)
}
def main(args: Array[String]): Unit = {
val fn = Res.loop[Int, Long] { rec =>
{ i =>
if i < 0 then Res(i.toLong)
else rec(i - 1).map(_ + 1L)
}
}
println(fn(1000000).run)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment