implement catchSome

Author: mikearnaldi

src/main/scala/IO.scala

@@ -1,6 +1,21 @@
 enum IO[-R, +E, +A]:
   def map[B](f: A => B): IO[R, E, B] = this.flatMap(a => IO.succeed(f(a)))
 
+  def flatMap[R2, E2, B](f: A => IO[R2, E2, B]): IO[R2 & R, E2 | E, B] = IO.FlatMap(this, f)
+
+  def foldM[R2, E2, A2, R3, E3, A3](f: A => IO[R2, E2, A2], g: E => IO[R3, E3, A3]): IO[R3 & R2 & R, E3 | E2, A3 | A2] = IO.Fold(this, f, g)
+
+  def inject[R2, R3](r: R2)(implicit tag: Tag[R2], ev: R2 & R3 => R): IO[R3, E, A] = IO.Provide(this, r, tag, ev)
+
+  def catchAll[R2, E2, A2](f: E => IO[R2, E2, A2]) = this.foldM(IO.succeed, f)
+  
+  def catchSome[R2, E2, A2, E3, E4](f: PartialFunction[E, (E3, IO[R2, E2, A2])])(implicit ev: E => E3 | E4): IO[R & R2, E2 | E4, A | A2] = this.foldM(IO.succeed, e => {
+    f.lift(e) match {
+      case Some((e3, io)) if e3 == e => io
+      case _ => IO.fail(ev(e).asInstanceOf[E4])
+    }
+  })
+
   case Succeed[A](
     val a: () => A
   ) extends IO[Any, Nothing, A]
@@ -29,10 +44,6 @@ enum IO[-R, +E, +A]:
 
 class Tag[A]()
 
-extension [R, E, A](zio: IO[R, E, A])
-  def flatMap[R2, E2, B](f: A => IO[R2, E2, B]): IO[R2 & R, E2 | E, B] = IO.FlatMap(zio, f)
-  def inject[R2, R3](r: R2)(implicit tag: Tag[R2], ev: R2 & R3 => R): IO[R3, E, A] = IO.Provide(zio, r, tag, ev)
-
 enum Exit[+E, +A]:
   case Fail(e: E)
   case Succeed(a: A)
@@ -49,7 +60,6 @@ extension [E, A](zio: IO[Any, E, A])
     var cont = List[ContiuationFrame]()
     var services = Map[Tag[Any], Any]()
 
-    // recursing
     var recursing = true
 
     while (recursing) {
@@ -136,6 +146,8 @@ extension [E, A](zio: IO[Any, E, A])
 
 object IO {
   def succeed[A](a: => A): IO[Any, Nothing, A] = Succeed(() => a)
+
   def fail[E](e: => E): IO[Any, E, Nothing] = Fail(() => e)
+
   def use[R, R2, E2, A](f: R => IO[R2, E2, A])(implicit tag: Tag[R]): IO[R & R2, E2, A] = Use(f, tag)
 }

src/main/scala/Main.scala

@@ -9,7 +9,7 @@ trait Console {
 object Console {
   implicit val tag: Tag[Console] = Tag()
 
-  def putStrLn(msg: => String) = IO use ((c: Console) => c putStrLn msg)
+  def putStrLn(msg: => String) = IO use ((T: Console) => T putStrLn msg)
 }
 
 trait Math {
@@ -19,7 +19,7 @@ trait Math {
 object Math {
   implicit val tag: Tag[Math] = Tag()
 
-  def add(x: Int, y: Int) = IO use ((c: Math) => c add(x, y))
+  def add(x: Int, y: Int) = IO use ((T: Math) => T add(x, y))
 }
 
 case class ErrorA(a: String)
@@ -30,9 +30,14 @@ val program = for {
   _ <- Console putStrLn "world"
   y <- Math add(2, 3)
   _ <- Console putStrLn s"result: $y"
+  _ <- IO fail ErrorA("a")
+  _ <- IO fail ErrorB("b")
 } yield ()
 
 val main = program
+  .catchSome(e => e match {
+    case x: ErrorA => (x, Console putStrLn "recovered from A")
+  })
   .inject(new Console {
     def putStrLn(msg: => String) = IO.succeed(println(msg))
   })