Add Strong and Arrow instances for IndexedStateT

typelevel · Sep 10, 2017 · 6e25fbb · 6e25fbb
1 parent 3be4238
commit 6e25fbb
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 4 deletions.
diff --git a/core/src/main/scala/cats/data/StateT.scala b/core/src/main/scala/cats/data/StateT.scala
@@ -1,7 +1,8 @@
 package cats
 package data
 
-import cats.functor.{ Contravariant, Bifunctor, Profunctor }
+import cats.functor.{ Contravariant, Bifunctor, Profunctor, Strong }
+import cats.arrow.Arrow
 import cats.syntax.either._
 
 /**
@@ -237,7 +238,7 @@ private[data] sealed abstract class IndexedStateTInstances2 extends IndexedState
     new IndexedStateTMonad[F, S] { implicit def F = F0 }
 }
 
-private[data] sealed abstract class IndexedStateTInstances3 {
+private[data] sealed abstract class IndexedStateTInstances3 extends IndexedStateTInstances4 {
   implicit def catsDataFunctorForIndexedStateT[F[_], SA, SB](implicit F0: Functor[F]): Functor[IndexedStateT[F, SA, SB, ?]] =
     new IndexedStateTFunctor[F, SA, SB] { implicit def F = F0 }
 
@@ -251,6 +252,16 @@ private[data] sealed abstract class IndexedStateTInstances3 {
     new IndexedStateTBifunctor[F, SA] { implicit def F = F0 }
 }
 
+private[data] sealed abstract class IndexedStateTInstances4 extends IndexedStateTInstances5 {
+  implicit def catsDataStrongForIndexedStateT[F[_], V](implicit F0: Monad[F]): Strong[IndexedStateT[F, ?, ?, V]] =
+    new IndexedStateTStrong[F, V] { implicit def F = F0 }
+}
+
+private[data] sealed abstract class IndexedStateTInstances5 {
+  implicit def catsDataArrowForIndexedStateT[F[_], V](implicit F0: Monad[F], V0: Monoid[V]): Arrow[IndexedStateT[F, ?, ?, V]] =
+    new IndexedStateTArrow[F, V] { implicit def F = F0; implicit def V = V0 }
+}
+
 // To workaround SI-7139 `object State` needs to be defined inside the package object
 // together with the type alias.
 private[data] abstract class StateFunctions {
@@ -312,6 +323,42 @@ private[data] sealed abstract class IndexedStateTProfunctor[F[_], V] extends Pro
     fab.dimap(f)(g)
 }
 
+private[data] sealed abstract class IndexedStateTStrong[F[_], V] extends IndexedStateTProfunctor[F, V] with Strong[IndexedStateT[F, ?, ?, V]] {
+  implicit def F: Monad[F]
+
+  def first[A, B, C](fa: IndexedStateT[F, A, B, V]): IndexedStateT[F, (A, C), (B, C), V] =
+    IndexedStateT { case (a, c) =>
+      F.map(fa.run(a)) { case (b, v) =>
+        ((b, c), v)
+      }
+    }
+
+  def second[A, B, C](fa: IndexedStateT[F, A, B, V]): IndexedStateT[F, (C, A), (C, B), V] =
+    first(fa).dimap((_: (C, A)).swap)(_.swap)
+}
+
+private[data] sealed abstract class IndexedStateTArrow[F[_], V] extends IndexedStateTStrong[F, V] with Arrow[IndexedStateT[F, ?, ?, V]] {
+  implicit def F: Monad[F]
+  implicit def V: Monoid[V]
+
+  def lift[A, B](f: A => B): IndexedStateT[F, A, B, V] =
+    IndexedStateT { a =>
+      F.pure((f(a), V.empty))
+    }
+
+  def id[A]: IndexedStateT[F, A, A, V] =
+    IndexedStateT.pure(V.empty)
+
+  def compose[A, B, C](f: IndexedStateT[F, B, C, V], g: IndexedStateT[F, A, B, V]): IndexedStateT[F, A, C, V] =
+    IndexedStateT { a =>
+      F.flatMap(g.run(a)) { case (b, va) =>
+        F.map(f.run(b)) { case (c, vb) =>
+          (c, V.combine(va, vb))
+        }
+      }
+    }
+}
+
 private[data] sealed abstract class IndexedStateTMonad[F[_], S] extends IndexedStateTFunctor[F, S, S] with Monad[IndexedStateT[F, S, S, ?]] {
   implicit def F: Monad[F]
 

diff --git a/tests/src/test/scala/cats/tests/StateTTests.scala b/tests/src/test/scala/cats/tests/StateTTests.scala
@@ -1,8 +1,9 @@
 package cats
 package tests
 
+import cats.arrow.Arrow
 import cats.data.{State, StateT, IndexedStateT, EitherT}
-import cats.functor.{Contravariant, Bifunctor, Profunctor}
+import cats.functor.{Contravariant, Bifunctor, Profunctor, Strong}
 import cats.kernel.instances.tuple._
 import cats.laws.discipline._
 import cats.laws.discipline.eq._
@@ -306,11 +307,31 @@ class IndexedStateTTests extends CatsSuite {
     implicit val FS: Profunctor[IndexedStateT[ListWrapper, ?, ?, Int]] = IndexedStateT.catsDataProfunctorForIndexedStateT
 
     checkAll("IndexedStateT[ListWrapper, String, Int, Int]", ProfunctorTests[IndexedStateT[ListWrapper, ?, ?, Int]].profunctor[String, String, String, Int, Int, Int])
-    checkAll("Profunctor[IndexedStateT[ListWrapper, ?, Int, Int]]", SerializableTests.serializable(Profunctor[IndexedStateT[ListWrapper, ?, ?, Int]]))
+    checkAll("Profunctor[IndexedStateT[ListWrapper, ?, ?, Int]]", SerializableTests.serializable(Profunctor[IndexedStateT[ListWrapper, ?, ?, Int]]))
 
     Profunctor[IndexedStateT[ListWrapper, ?, ?, Int]]
   }
 
+  {
+    implicit val F: Monad[ListWrapper] = ListWrapper.monad
+    implicit val FS: Strong[IndexedStateT[ListWrapper, ?, ?, Int]] = IndexedStateT.catsDataStrongForIndexedStateT
+
+    checkAll("IndexedStateT[ListWrapper, String, Int, Int]", StrongTests[IndexedStateT[ListWrapper, ?, ?, Int]].strong[String, String, String, Int, Int, Int])
+    checkAll("Strong[IndexedStateT[ListWrapper, ?, ?, Int]]", SerializableTests.serializable(Strong[IndexedStateT[ListWrapper, ?, ?, Int]]))
+
+    Strong[IndexedStateT[ListWrapper, ?, ?, Int]]
+  }
+
+  {
+    implicit val F: Monad[ListWrapper] = ListWrapper.monad
+    implicit val FS: Arrow[IndexedStateT[ListWrapper, ?, ?, Int]] = IndexedStateT.catsDataArrowForIndexedStateT
+
+    checkAll("IndexedStateT[ListWrapper, String, Int, Int]", ArrowTests[IndexedStateT[ListWrapper, ?, ?, Int]].arrow[String, String, String, Int, Int, Int])
+    checkAll("Arrow[IndexedStateT[ListWrapper, ?, ?, Int]]", SerializableTests.serializable(Arrow[IndexedStateT[ListWrapper, ?, ?, Int]]))
+
+    Arrow[IndexedStateT[ListWrapper, ?, ?, Int]]
+  }
+
   {
     // F has a Monad
     implicit val F = ListWrapper.monad