State.scala

import scala.language.implicitConversions

case class State[S,R](f: S => (R,S)) {
  def apply(s: S): (R,S) = f(s)
  def calculate(s: S): R = f(s)._1
  def perform(s: S): S = f(s)._2

  def map[Q](g: R => Q): State[S,Q] = State(s => {
    val (res,_) = f(s)
    (g(res),s)
  })

  def flatMap[Q](g: R => State[S,Q]): State[S,Q] = State(s => {
    val (res,state): (R,S) = f(s)
    g(res)(state)
  })
}

object State {
  def pure[S,R](x: R): State[S,R] = State(s => (x,s))

  def gets[S,R](f: S => R): State[S,R] = State(s => (f(s),s))
  def get[S]: State[S,S] = gets(identity)

  def modify[S](f: S => S): State[S,Unit] = State(s => ((),f(s)))
  def put[S](x: S): State[S,Unit] = modify(_ => x)

  def sequence[S,R](ss: List[State[S,R]]): State[S,List[R]] = State(s => {
      val accu: (List[R], S) = (Nil,s)
      ss.foldLeft(accu) { case ((results,accuState),stateM) => { 
          val (res,state): (R,S) = stateM(accuState) 
          (results :+ res, state) 
        }
      }
    })

  implicit def function2State[S,R](f: S => (R,S)): State[S,R] = State(f)
  implicit def state2Function[S,R](state: State[S,R]): S => (R,S) = state.f
}

type Stack = List[Int]

def push(i: Int): State[Stack,Unit] = State(s => {
  ((),i :: s)
})

def pop: State[Stack,Int] = State(s => {
  (s.head,s.tail)
})

def manipulate: State[Stack,Int] = for {
  _ <- push(5)
  _ <- push(6)
  v <- pop
} yield v 

println(manipulate.calculate(List()))

def pop3: State[Stack,List[Int]] = State.sequence(List(pop,pop,pop))
println(s"Result is: ${pop3.calculate((1 to 20).toList)}")

def push3: State[Stack,List[Unit]] = State.sequence(List(push(3),push(3),push(3)))
println(s"State is: ${push3.perform(List())}")

println(State.pure(5).map(_ * 2)(()))