State.scala

package readerwriter.internal
/* You don't need to edit any of the files inside this package. These are the actual implementations, and you should
 * only use the provided methods, not rely on internal structure
 */

import applicative.Monad

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

case object State:
  def get[S]: State[S, S] = State(s => (s, s))
  def set[S](s: S): State[S, Unit] = State(_ => (s, ()))

  given stateMonad[S]: Monad[[a] =>> State[S, a]] with
    def pure[A](a: A): State[S, A] = State(s => (s, a))
    extension [A](fa: State[S, A])
      def flatMap[B](f: A => State[S, B]): State[S, B] = State(s => 
        val (s1, a) = fa.run(s)
        f(a).run(s1)
      )