StateF.scala

package scalaz

import scalaz.{State => _, _}
import scalaz.Free.FreeC
import scalaz.Id.Id
import FreeState.StateF

sealed abstract class FreeState[S, A] {
  def map[B](f: A => B): FreeState[S, B]
}

object Main extends App{
  import std.list._

  val size = args.headOption.fold(100000)(_.toInt)
  println("size = " + size)
  val list = (1 to size).toList
  val F = Foldable[List]

  def a = F.traverse_[({type l[a] = StateF[Int, a]})#l, Int, Unit](list)(a => FreeState.put(a))(FreeState.stateMonad)
  def b = F.traverse_[({type l[a] = scalaz.State[Int, a]})#l, Int, Unit](list)(a => scalaz.State.put(a))

  println(
    FreeState.eval(a, 0)
  )

//  println(b.eval(0))
}

object FreeState {

  implicit class StateFOps[S, A](val self: StateF[S, A]) extends AnyVal {
    def eval(init: S): (S, A) =
      FreeState.eval(self, init)

    def evalZero(implicit S: Monoid[S]): (S, A) =
      FreeState.evalZero(self)

    def contramap[B](f: B => S): StateF[B, A] =
      FreeState.contramap(self)(f)
  }

  private final case class SetState[S, A](s: S, k: A) extends FreeState[S, A] {
    def map[B](f: A => B) =
      SetState(s, f(k))
  }
  private final case class GetState[S, A](k: S => A) extends FreeState[S, A] {
    def map[B](f: A => B) =
      GetState(k andThen f)
  }
 
  type StateF[S, A] = Free[({type l[a] = FreeState[S, a]})#l, A]

  private def liftFreeState[S, A](s: FreeState[S, A]): StateF[S, A] =
    Free.liftF[({type l[a] = FreeState[S, a]})#l,A](s)

  implicit def stateFFunctor[S]: Functor[({type l[a] = FreeState[S, a]})#l] =
    new Functor[({type l[a] = FreeState[S, a]})#l] {
      def map[A, B](fa: FreeState[S, A])(f: A => B) =
        fa map f
    }

  implicit def stateMonad[S]: Monad[({type l[a] = StateF[S, a]})#l] =
    Free.freeMonad[({type l[a] = FreeState[S, a]})#l](stateFFunctor)

  def gets[S, A](f: S => A): StateF[S, A] =
    liftFreeState(GetState(f))

  def init[S]: StateF[S, S] =
    liftFreeState(GetState(identity))

  def put[S](s: S): StateF[S, Unit] =
    liftFreeState(SetState(s, () => ()))

  def modify[S](f: S => S): StateF[S, Unit] =
    init[S].flatMap(x => put(f(x)))

  def state[S, A](f: S => (S, A)): StateF[S, A] =
    modify[S](s => f(s)._1).flatMap(_ => gets(s => f(s)._2))

  def eval[S, A](state: StateF[S, A], init: S): (S, A) = {
    state.foldRun(init){
      case (s, GetState(f)) =>
        s -> f(s)
      case (_, SetState(ss, a)) =>
        ss -> a
    }
  }

  def evalZero[S, A](state: StateF[S, A])(implicit S: Monoid[S]): (S, A) =
    eval(state, S.zero)

  def contramap[S, A, X](state: StateF[S, A])(g: X => S): StateF[X, A] =
    gets(g).map(eval(state, _)._2)
}