Add basic effectful Store

StoreF, or Monomial[F[S], S, Y], for when S must be accessed by an effectful function.
julianpeeters · Mar 22, 2024 · ac857c4 · ac857c4
1 parent d5de9b9
commit ac857c4
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 4 deletions.
diff --git a/modules/poly/shared/src/main/scala/polynomial/morphism/PolyMap.scala b/modules/poly/shared/src/main/scala/polynomial/morphism/PolyMap.scala
@@ -2,7 +2,7 @@ package polynomial.morphism
 
 import cats.Id
 import polynomial.`object`.Binomial.BiInterface
-import polynomial.`object`.Monomial.{Interface, Store}
+import polynomial.`object`.Monomial.{Interface, Store, StoreF}
 import polynomial.product.{Cartesian, Composition, Tensor}
 
 type ~>[P[_], Q[_]] = PolyMap[P, Q, _]
@@ -31,6 +31,7 @@ object PolyMap:
     case (PolyMap[o, p, Y], Tensor[q, r, Y])                              => Phi[o, Tensor[q, r, _], Y]
     case (Store[s, Y], BiInterface[a1, b1, a2, b2, Y])                    => (s => b1, s => b2)
     case (Store[s, Y], Interface[a, b, Y])                                => s => b
+    case (StoreF[f, s, Y], Interface[a, b, Y])                            => s => b
     case (Store[s, Y], Composition[p, q, Y])                              => (Phi[Store[s, _], p, Y], Phi[p, q, Y])
 
   type PhiSharp[P[_], Q[_], Y] = (Eval[P, Y], Eval[Q, Y]) match
@@ -40,6 +41,7 @@ object PolyMap:
     case (Interface[a1, b1, Y], Interface[a2, b2, Y])                     => (b1, a2) => a1
     case (Store[s, Y], BiInterface[a1, b1, a2, b2, Y])                    => ((s, a1) => s, (s, a2) => s)
     case (Store[s, Y], Interface[a, b, Y])                                => (s, a) => s
+    case (StoreF[f, s, Y], Interface[a, b, Y])                            => (s, a) => f[s]
     case (Store[s, Y], Composition[p, q, Y])                              => (s, Composition.DecomposeSharp[p, q, Y]) => s
     case (PolyMap[p, q, Y], BiInterface[a3, b3, a4, b4, Y])               => PhiSharp[p, BiInterface[a3, b3, a4, b4, _], Y]
     case (PolyMap[p, q, Y], Interface[a2, b2, Y])                         => PhiSharp[p, Interface[a2, b2, _], Y]
@@ -51,6 +53,7 @@ object PolyMap:
     case Interface[a, b, Y]             => Interface[a, b, Y]
     case PolyMap[p, q, Y]               => PolyMap[p, q, Y]
     case Store[s, Y]                    => Store[s, Y]
+    case StoreF[f, s, Y]                => StoreF[f, s, Y]
     case Tensor[p, q, Y]                => Tensor.DayConvolution[Eval[p, _], Eval[q, _], Y]
 
   export polynomial.morphism.methods.andThen.mS_bI.*

diff --git a/modules/poly/shared/src/main/scala/polynomial/morphism/methods/andThen/mS_mI.scala b/modules/poly/shared/src/main/scala/polynomial/morphism/methods/andThen/mS_mI.scala
@@ -2,8 +2,8 @@ package polynomial.morphism.methods.andThen
 
 import cats.Id
 import polynomial.morphism.PolyMap
-import polynomial.`object`.Monomial.{Interface, Store}
-import cats.FlatMap
+import polynomial.`object`.Monomial.{Interface, Store, StoreF}
+import cats.{FlatMap, Monad}
 
 object mS_mI:
 
@@ -27,4 +27,26 @@ object mS_mI:
         def φ: PolyMap.Phi[PolyMap[Store[F[S], _], Interface[A, Id[B], _], _], Interface[A, A => F[B], _], Y] =
           s => a2 => FlatMap[F].flatMap(p.φ(s))(b1 => w.φ(b1)(a2))
         def `φ#`: PolyMap.PhiSharp[PolyMap[Store[F[S], _], Interface[A, Id[B], _], _], Interface[A, A => F[B], _], Y] =
-          (s, a) => FlatMap[F].flatMap(p.φ(s))(b1 => p.`φ#`(s, w.`φ#`(b1, a)))
+          (s, a) => FlatMap[F].flatMap(p.φ(s))(b1 => p.`φ#`(s, w.`φ#`(b1, a)))
+
+  extension [F[_]: Monad, S, A, B, Y] (p: PolyMap[StoreF[F, S, _], Interface[A, F[B], _], Y])
+    @scala.annotation.targetName("andThenStoreFMonoToMonoFPrism")
+    def andThen(
+      w: PolyMap[Interface[A, F[B], _], Interface[A, A => F[B], _], Y]
+    ): PolyMap[PolyMap[StoreF[F, S, _], Interface[A, F[B], _], _], Interface[A, A => F[B], _], Y] =
+      new PolyMap[PolyMap[StoreF[F, S, _], Interface[A, F[B], _], _], Interface[A, A => F[B], _], Y]:
+        def φ: PolyMap.Phi[PolyMap[StoreF[F, S, _], Interface[A, F[B], _], _], Interface[A, A => F[B], _], Y] =
+          s => a2 => Monad[F].flatMap(p.φ(s))(b1 => w.φ(Monad[F].pure(b1))(a2))
+        def `φ#`: PolyMap.PhiSharp[PolyMap[StoreF[F, S, _], Interface[A, F[B], _], _], Interface[A, A => F[B], _], Y] =
+          (s, a) => Monad[F].flatMap(p.φ(s))(b1 => p.`φ#`(s, w.`φ#`(Monad[F].pure(b1), a)))
+
+  extension [F[_]: Monad, S, A, B, Y] (p: PolyMap[Interface[F[S], S, _], Interface[A, F[B], _], Y])
+    @scala.annotation.targetName("andThenFStoreMonoToMonoFPrism")
+    def andThen(
+      w: PolyMap[Interface[A, F[B], _], Interface[A, A => F[B], _], Y]
+    ): PolyMap[PolyMap[Interface[F[S], S, _], Interface[A, F[B], _], _], Interface[A, A => F[B], _], Y] =
+      new PolyMap[PolyMap[Interface[F[S], S, _], Interface[A, F[B], _], _], Interface[A, A => F[B], _], Y]:
+        def φ: PolyMap.Phi[PolyMap[Interface[F[S], S, _], Interface[A, F[B], _], _], Interface[A, A => F[B], _], Y] =
+          s => a2 => Monad[F].flatMap(p.φ(s))(b1 => w.φ(Monad[F].pure(b1))(a2))
+        def `φ#`: PolyMap.PhiSharp[PolyMap[Interface[F[S], S, _], Interface[A, F[B], _], _], Interface[A, A => F[B], _], Y] =
+          (s, a) => Monad[F].flatMap(p.φ(s))(b1 => p.`φ#`(s, w.`φ#`(Monad[F].pure(b1), a)))
diff --git a/modules/poly/shared/src/main/scala/polynomial/object.scala b/modules/poly/shared/src/main/scala/polynomial/object.scala
@@ -7,6 +7,7 @@ object `object`:
   object Monomial:
     case class Interface[A, B, Y](yᵉˣᵖ: A => Y, coeff: B) extends Monomial[A, B, Y]
     case class Store[S, Y](yᵉˣᵖ: S => Y, coeff: S) extends Monomial[S, S, Y]
+    case class StoreF[F[_], S, Y](yᵉˣᵖ: F[S] => Y, coeff: S) extends Monomial[F[S], S, Y]
 
   // B₁yᴬ¹ + B₂yᴬ²
   sealed trait Binomial[+A1, +B1, +A2, +B2, Y]