Merge pull request #780 from adelbertc/state-transformS

Add StateT transformS

core/src/main/scala/cats/state/StateT.scala

@@ -73,6 +73,33 @@ final class StateT[F[_], S, A](val runF: F[S => F[(S, A)]]) extends Serializable
   def transformF[G[_], B](f: F[(S, A)] => G[(S, B)])(implicit F: FlatMap[F], G: Applicative[G]): StateT[G, S, B] =
     StateT(s => f(run(s)))
 
+  /**
+   * Transform the state used.
+   *
+   * This is useful when you are working with many focused `StateT`s and want to pass in a
+   * global state containing the various states needed for each individual `StateT`.
+   *
+   * {{{
+   * scala> import cats.std.option._ // needed for StateT.apply
+   * scala> type GlobalEnv = (Int, String)
+   * scala> val x: StateT[Option, Int, Double] = StateT((x: Int) => Option((x + 1, x.toDouble)))
+   * scala> val xt: StateT[Option, GlobalEnv, Double] = x.transformS[GlobalEnv](_._1, (t, i) => (i, t._2))
+   * scala> val input = 5
+   * scala> x.run(input)
+   * res0: Option[(Int, Double)] = Some((6,5.0))
+   * scala> xt.run((input, "hello"))
+   * res1: Option[(GlobalEnv, Double)] = Some(((6,hello),5.0))
+   * }}}
+   */
+  def transformS[R](f: R => S, g: (R, S) => R)(implicit F: Monad[F]): StateT[F, R, A] =
+    StateT { r =>
+      F.flatMap(runF) { ff =>
+        val s = f(r)
+        val nextState = ff(s)
+        F.map(nextState) { case (s, a) => (g(r, s), a) }
+      }
+    }
+
   /**
    * Modify the state (`S`) component.
    */

tests/src/test/scala/cats/tests/StateTTests.scala

@@ -77,6 +77,23 @@ class StateTTests extends CatsSuite {
     }
   }
 
+  test("StateT#transformS with identity is identity") {
+    forAll { (s: StateT[List, Long, Int]) =>
+      s.transformS[Long](identity, (s, i) => i) should === (s)
+    }
+  }
+
+  test("StateT#transformS modifies state") {
+    final case class Env(int: Int, str: String)
+    val x = StateT((x: Int) => Option((x + 1, x)))
+    val xx = x.transformS[Env](_.int, (e, i) => e.copy(int = i))
+    val input = 5
+
+    val got = x.run(input)
+    val expected = xx.run(Env(input, "hello")).map { case (e, i) => (e.int, i) }
+    got should === (expected)
+  }
+
   {
     implicit val iso = MonoidalTests.Isomorphisms.invariant[StateT[Option, Int, ?]]
     checkAll("StateT[Option, Int, Int]", MonadStateTests[StateT[Option, Int, ?], Int].monadState[Int, Int, Int])