Last active
July 14, 2021 02:50
-
-
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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