Implementation of Free monad from scratch

free-monad-from-scratch.scala

import $plugin.$ivy.`org.spire-math::kind-projector:0.9.9`
import $ivy.`org.typelevel::cats-core:1.5.0`

import java.nio.file.{Files, Paths}

import cats.{Eval, Monad}
import cats.arrow.FunctionK
import cats.data.State
import scala.io.Source

sealed trait Free[F[_], A] {
    import Free._

    /**
      * Needed for creating sequential computations and so programs
      */
    def flatMap[B](f: A => Free[F, B]): Free[F, B] = FlatMap(this, f)

    /**
      * Needed for creating sequential computations and so programs
      */
    def map[B](f: A => B): Free[F, B] = this.flatMap(a => Pure(f(a)))

    /**
      * Interprets a Free monad into G[A]
      *
      * @param nt Natural transformation from F[_] type constructor to G[_]
      * @tparam G Type constructor with instance of Monad typeclass
      * @return
      */
    def foldMap[G[_]: Monad](nt: FunctionK[F, G]): G[A] = this match {
      case Pure(a)  => Monad[G].pure(a)
      case Lift(fa) => nt(fa)
      case FlatMap(fa, f) =>
        Monad[G].flatMap(fa.foldMap(nt))(a => f(a).foldMap(nt))
    }
  }

object Free {
    final case class Pure[F[_], A](a: A) extends Free[F, A]
    final case class Lift[F[_], A](fa: F[A]) extends Free[F, A]
    final case class FlatMap[F[_], A, B](fa: Free[F, A], f: A => Free[F, B])
        extends Free[F, B]

    /**
      * Wraps a: A into Free[F, A] context
      */
    def pure[F[_], A](a: A): Free[F, A] = Pure(a)

    /**
      * Wraps fa: F[A] into Free[F, A] context
      */
    def liftM[F[_], A](fa: F[A]): Free[F, A] = Lift(fa)
  }

//Our API - we strict ourselves to use only "safe" and "allowed" operations
//we don't work with the disk IO directly

sealed trait Disk[A]
object Disk {
    final case class Read(filename: String) extends Disk[Array[Byte]]
    final case class Write(filename: String, data: Array[Byte])
        extends Disk[Unit]

    def read(filename: String): Free[Disk, Array[Byte]] =
      Free.liftM(Read(filename))
    def write(filename: String, data: Array[Byte]) =
      Free.liftM(Write(filename, data))
  }

import Disk._

//Our program

val program: Free[Disk, String] = for {
    data <- read("test.txt")
    _ <- write("test.log", "Hello ".getBytes ++ data ++ "!".getBytes)
    newData <- read("test.log")
  } yield new String(newData, "UTF-8")

//Natural transformation that knows how to actually run our disk IO operations
//Here we interpret Disk[A] into real-world Eval monad (can be any other monad)
  val nt = new FunctionK[Disk, Eval] {
    override def apply[A](fa: Disk[A]): Eval[A] = {
      fa match {
        case Read(filename) =>
          Eval.always {
            Source.fromFile(filename, "UTF-8").mkString.getBytes
          }
        case Write(filename, data) =>
          Eval.always {
            Files.write(Paths.get(filename), data)
            ()
          }
      }
    }
  }

//Actual running our program in production
//val eval = program.foldMap(nt)
//eval.value

//Fixing S in State leaving 1 hole

type InMemoryState[A] = State[Map[String, Array[Byte]], A]
//Testing interpreter that mocks disk IO operations using in-memory Map[String, Array[Byte]] and State monad

val nt2 = new FunctionK[Disk, InMemoryState] {
    override def apply[A](fa: Disk[A]): InMemoryState[A] =
      fa match {
        case Read(filename) =>
          State[Map[String, Array[Byte]], Array[Byte]](state =>
            (state, state(filename)))
        case Write(filename, data) =>
          State[Map[String, Array[Byte]], Unit](state =>
            (state.updated(filename, data), ()))
      }
  }

//Boiler plate Monad instance

val stateMonadInstance: Monad[InMemoryState] =
    new Monad[InMemoryState] {
      override def pure[A](x: A): State[Map[String, Array[Byte]], A] =
        State.pure(x)

      override def flatMap[A, B](fa: State[Map[String, Array[Byte]], A])(
          f: A => State[Map[String, Array[Byte]], B]) = fa.flatMap(f)

      override def tailRecM[A, B](a: A)(
          f: A => State[Map[String, Array[Byte]], Either[A, B]]) = ???
    }
val eval =
    program
      .foldMap(nt2)
      .run(Map("test.txt" -> "Igor".getBytes))
val (state, result) = eval.value
println(state)
println(result)
//prints
//Map(test.txt -> [B@69d1ad64, test.log -> [B@458b4487)
//Hello Igor!