Trace any stateful computation in 10 lines.

trace.scala

/* Lift a stateful computation State[S,A] to State[(S, Seq[(S,A)]), A]
 * where the (S, Seq[(S,A)]) is a history of the computation's state at
 * each step: (currentState, Seq[(stateBeforeStep, resultOfStep)])
 */
def trace[S,A](iter: State[S,A]): State[(S, Seq[(S,A)]),A] = for {
  preStateAndHistory <- get[(S, Seq[(S,A)])]
  preState = preStateAndHistory._1
  history = preStateAndHistory._2
  (postState, result) = runState(for {
    result <- iter
    postState <- get
  } yield (postState, result), preState)
  _ <- set((postState, (preState, result) +: history))
} yield result