monads.scala

sealed trait List[+A]
case object Nil extends List[Nothing]
case class Cons[+A](head: A, tail: List[A]) extends List[A]

object List {
  def apply[A](as: A*): List[A] =
    if (as.isEmpty) Nil
    else Cons(as.head, apply(as.tail: _*))

  def reverse[A](list: List[A]): List[A] = {
    foldLeft(list, Nil: List[A])((h, t) => Cons(t, h))
  }

  @annotation.tailrec
  def foldLeft[A,B](as: List[A], z: B)(f: (B, A) => B): B =
    as match {
      case Nil => z
      case Cons(x, xs) => foldLeft(xs, f(z, x))(f)
    }

  def foldRight[A,B](as: List[A], z: B)(f: (A, B) => B): B =
    foldLeft(reverse(as), z)((a, b) => f(b, a))

  def concat[A](a: List[A], b: List[A]): List[A] =
    foldRight(a, b)(Cons.apply)

  def flatten[A](list: List[List[A]]): List[A] =
    foldLeft(list, List[A]())(concat)

  def filter[A](as: List[A])(f: A => Boolean): List[A] =
    as match {
      case Nil => Nil
      case Cons(x, xs) =>
        val rest = filter(xs)(f)
        if (f(x)) {
          Cons(x, rest)
        } else {
          rest
        }
    }

  def map[A, B](list: List[A])(f: A => B): List[B] =
    list match {
      case Nil => Nil
      case Cons(x, xs) => Cons(f(x), map(xs)(f))
    }

  def flatMap[A,B](as: List[A])(f: A => List[B]): List[B] =
    flatten(map(as)(f))

  def filter2[A](as: List[A])(f: A => Boolean): List[A] =
    flatMap(as)(x => if (f(x)) { List(x)} else {List()})

  @annotation.tailrec
  def nth[A](as: List[A], n: Int): A =
    as match {
      case Cons(x, xs) => if (n > 1) nth(xs, n - 1) else x
      case Nil => throw null // we don't know about Option yet
    }

  def zipWith2[A](xs: List[A], ys: List[A])(f: (A, A) => A): List[A] = {
    val (_, result) = foldLeft(xs, (ys, List[A]()))((pair, x) => {
      pair._1 match {
        case Cons(y, ytail) => (ytail, Cons(f(x, y), pair._2))
        case Nil => (Nil, Nil)
      }
    })
    reverse(result)
  }

  def zipWith[A](xs: List[A], ys: List[A])(f: (A, A) => A): List[A] = {
    @annotation.tailrec
    def zipWithImpl(xs: List[A], ys: List[A], result: List[A])(f: (A, A) => A): List[A] =
      (xs, ys) match {
        case (Cons(x, xtail), Cons(y, ytail)) => zipWithImpl(xtail, ytail, Cons(f(x, y), result))(f)
        case _ => result
      }
    reverse(zipWithImpl(xs, ys, Nil)(f))
  }

  @annotation.tailrec
  def hasPrefix[A](sup: List[A], sub: List[A]): Boolean =
    (sup, sub) match {
      case (Cons(x, xs), Cons(y, ys)) if x == y => hasPrefix(xs, ys)
      case (_, Nil) => true
      case _ => false
    }

  @annotation.tailrec
  def hasSubsequence[A](sup: List[A], sub: List[A]): Boolean =
    sup match {
      case Nil => sub == Nil
      case list if hasPrefix(list, sub) => true
      case Cons(_, xs) => hasSubsequence(xs, sub)
    }
}

trait Functor[F[_]] {
  def map[A,B](fa: F[A])(f: A => B): F[B]
}

trait Applicative[F[_]] extends Functor[F] {
  // primitive combinators
  def map2[A,B,C](fa: F[A], fb: F[B])(f: (A, B) => C): F[C]
  def unit[A](a: => A): F[A]
  // derived combinators
  def map[A, B](fa: F[A])(f: A => B): F[B] =
    map2(fa, unit(()))((a, _) => f(a))
  def traverse[A,B](as: List[A])(f: A => F[B]): F[List[B]] =
    List.foldRight(as, unit(List[B]()))((a, fbs) => map2(f(a), fbs)(Cons(_,  _)))

  def sequence[A](fs: List[F[A]]): F[List[A]] =
    traverse(fs)(x => x)

  def replicateM[A](n: Int, ma: F[A]): F[List[A]] = {
    val list: List[F[A]] = List.apply(Range(0, n).map(_ => ma): _*)
    sequence(list)
  }

  def product[A,B](ma: F[A], mb: F[B]): F[(A, B)] = map2(ma, mb)((_, _))

  def apply[A,B](fab: F[A => B])(fa: F[A]): F[B] =
    map2(fab, fa)((f, x) => f(x))
}

trait Applicative2[F[_]] extends Functor[F] {
  def apply[A,B](fab: F[A => B])(fa: F[A]): F[B]
  def unit[A](a: => A): F[A]

  def map2[A,B,C](fa: F[A], fb: F[B])(f: (A, B) => C): F[C] =
    apply(apply(unit(f.curried))(fa))(fb)

  def map[A,B](fa: F[A])(f: A => B): F[B] =
    apply(unit(f))(fa)

  def map3[A,B,C,D](fa: F[A],
                    fb: F[B],
                    fc: F[C])(f: (A, B, C) => D): F[D] =
    apply(apply(apply(unit(f.curried))(fa))(fb))(fc)

  def map4[A,B,C,D,E](fa: F[A],
                      fb: F[B],
                      fc: F[C],
                      fd: F[D])(f: (A, B, C, D) => E): F[E] =
    apply(apply(apply(apply(unit(f.curried))(fa))(fb))(fc))(fd)
}

trait Monad[F[_]] extends Applicative[F] {
  // primitive combinators
  def unit[A](a: => A): F[A]
  def flatMap[A,B](ma: F[A])(f: A => F[B]): F[B]

  // derived combinators
  def map2[A,B,C](ma: F[A], mb: F[B])(f: (A, B) => C): F[C] =
    flatMap(ma)(a => flatMap(mb)(b => unit(f(a, b))))

  def filterM[A](ms: List[A])(f: A => F[Boolean]): F[List[A]] = {
    val lists = traverse(ms)(a => map(f(a))(b => if (b) List(a) else List()))
    map(lists)(l => List.flatten(l))
  }

  def compose[A,B,C](f: A => F[B], g: B => F[C]): A => F[C] =
    a => flatMap(f(a))(g)

  def flatMap2[A,B](ma: F[A])(f: A => F[B]): F[B] =
    compose[F[A], A, B](identity, f)(ma)

  def join[A](mma: F[F[A]]): F[A] =
    flatMap(mma)(identity)
}

object Monad {
  val optionMonad: Monad[Option] = new Monad[Option] {
    def unit[A](a: => A): Option[A] = Some(a)
    def flatMap[A,B](ma: Option[A])(f: A => Option[B]): Option[B] =
      ma flatMap f
  }

  val listMonad: Monad[List] = new Monad[List] {
    def unit[A](a: => A): List[A] = List(a)
    def flatMap[A,B](ma: List[A])(f: A => List[B]): List[B] =
      List.flatMap(ma)(f)
  }

  def stateMonad[S]: Monad[({type L[A] = State[S, A]})#L] = {
    type StateS[A] = State[S, A]

    new Monad[StateS] {
      override def unit[A](a: => A): StateS[A] = State(s => (a, s))
      override def flatMap[A, B](f: StateS[A])(g: A => StateS[B]): StateS[B] =
        State[S, B](
          s1 => {
            val (a, s2) = f.run(s1)
            g(a).run(s2)
          }
        )
    }
  }
}

case class State[S,+A](run: S => (A,S))

case class Id[A](value: A) extends Monad[Id] {
  override def unit[A](a: => A): Id[A] = Id(a)
  override def flatMap[A, B](ma: Id[A])(f: A => Id[B]): Id[B] = f(ma.value)
}

case class Reader[R, A](run: R => A)
object Reader {
  def readerMonad[R] = new Monad[({type f[x] = Reader[R,x]})#f] {
    def unit[A](a: => A): Reader[R,A] = Reader(_ => a)
    def flatMap[A,B](st: Reader[R,A])(f: A => Reader[R,B]): Reader[R,B] = Reader[R,B](
      r => {
        val a = st.run(r)
        f(a).run(r)
      }
    )
  }
}

object Hello extends App {
//  println(Monad.listMonad.replicateM(3, List(1,2)))
//  println(Monad.optionMonad.replicateM(3, Some("a")))
//  println(Monad.optionMonad.filterM(List("a", "b", "a"))(a => Some(a == "a")))
//  println(Monad.listMonad.filterM(List(1, 2))(a => List(true, false)))
  val state = State[Int, Int](i => (i, i + 1))
  val state2 = State[Int, Int](i => (i, i * 3))
  val state3 = State[Int, Int](i => (i, i - 1))
  println(Monad.stateMonad.replicateM(3, state).run(2))
  println(Monad.stateMonad.sequence(List(state, state2, state3)).run(3))
  println(Monad.stateMonad.sequence(List(state, state, state)).run(2))
}