Star.hs

{-# LANGUAGE AllowAmbiguousTypes #-}

-- | 'Star' embeds a functor @f@ as a profunctor with the functor on the target side:
-- @Star f a b = a ~> f b@. It is the representable profunctor of @f@; for a Haskell monad @m@,
-- @Star (Prelude m)@ is its Kleisli promonad.
module Proarrow.Profunctor.Instance.Star where

import Control.Monad qualified as P
import Data.Functor.Compose (Compose (..))
import Data.Kind (Type)
import Prelude qualified as P

import Proarrow.Category.Enriched.Thin (DecidableProfunctor (..), Thin, ThinProfunctor (..), mapDecision)
import Proarrow.Category.Instance.Nat (ApplyAction, Nat' (..), type (.->) (..))
import Proarrow.Category.Instance.Prof (Prof (..))
import Proarrow.Category.Monoidal (Monoidal (..), MonoidalProfunctor (..), Tensor)
import Proarrow.Category.Monoidal.Action (CoprodAction, ProdAction, SubAction)
import Proarrow.Category.Monoidal.Applicative (Alternative (..), Applicative (..))
import Proarrow.Category.Monoidal.Distributive (Distributive, Traversable (..), baseTraverse)
import Proarrow.Category.Monoidal.Strength (Strong (..))
import Proarrow.Colimit.BinaryCoproduct (COPROD (..), Coprod (..), HasBinaryCoproducts (..), HasCoproducts, (++))
import Proarrow.Colimit.Initial (initiate)
import Proarrow.Core (CategoryOf (..), Hom, Profunctor (..), Promonad (..), lmap, obj, (:~>), type (+->))
import Proarrow.Functor (Functor (..), Prelude (..), withObF)
import Proarrow.Profunctor.Instance.Composition ((:.:) (..))
import Proarrow.Profunctor.Instance.Coproduct ((:+:) (..))
import Proarrow.Profunctor.Representable (Representable (..), dimapRep)

type Star' :: j .-> k -> j +-> k
data Star' f a b where
  Star' :: (Ob b) => {unStar :: a ~> f b} -> Star' (NT f) a b

type Star f = Star' (NT f)
pattern Star :: () => (Ob b) => (a ~> f b) -> Star f a b
pattern Star f = Star' f
{-# COMPLETE Star #-}

instance (Functor f) => Profunctor (Star f) where
  dimap = dimapRep
  r \\ Star f = r \\ f

instance (CategoryOf j, CategoryOf k) => Functor (Star' :: (j .-> k) -> j +-> k) where
  map (Nat' n) = Prof \(Star f) -> Star (n . f)

instance (Functor f) => Representable (Star f) where
  type Star f % a = f a
  index = unStar
  tabulate = Star
  repMap = map

instance (Profunctor p) => Promonad (Star ((:+:) p)) where
  id = Star (Prof InjR)
  Star (Prof l) . Star (Prof r) = Star (Prof (\a -> case r a of InjL p -> InjL p; InjR b -> l b))

instance (P.Monad m) => Promonad (Star (Prelude m)) where
  id = Star (Prelude . P.pure)
  Star g . Star f = Star \a -> Prelude (unPrelude (f a) P.>>= (unPrelude . g))

composeStar :: (Functor f) => Star f :.: Star g :~> Star (Compose f g)
composeStar (Star f :.: Star g) = Star (Compose . map g . f)

instance (Applicative f, Monoidal j, Monoidal k) => MonoidalProfunctor (Star (f :: j -> k)) where
  one = Star (pure id)
  Star @a f ** Star @b g = withOb2 @_ @a @b (Star (liftA2 @f @a @b id . (f ** g)))

instance (Functor f, HasCoproducts j, HasCoproducts k) => MonoidalProfunctor (Coprod (Star (f :: j -> k))) where
  one = Coprod (Star initiate)
  Coprod (Star @a f) ** Coprod (Star @b g) = withObCoprod @_ @a @b (Coprod (Star (map (lft @_ @a @b) . f ||| map (rgt @_ @a @b) . g)))

-- Hmm, another wrapper required...

-- | Wraps the domain of @p@ in 'COPROD', letting 'Star' of an 'Alternative' functor be monoidal
-- over coproducts.
type CoprodDom :: j +-> k -> COPROD j +-> k
data CoprodDom p a b where
  Co :: {unCo :: p a b} -> CoprodDom p a (COPR b)

instance (Profunctor p) => Profunctor (CoprodDom p) where
  dimap l (Coprod r) (Co p) = Co (dimap l r p)
  r \\ Co p = r \\ p

instance (Alternative f, Monoidal k, Distributive j) => MonoidalProfunctor (CoprodDom (Star (f :: j -> k))) where
  one = Co (Star empty)
  Co (Star @a f) ** Co (Star @b g) = let ab = obj @a +++ obj @b in Co (Star (alt @f @a @b ab . (f ** g))) \\ ab

instance (P.Functor f) => Strong ProdAction (Star (Prelude f)) where
  act (Star k) = Star (\(a, x) -> P.fmap (a,) (k x))

instance (Functor f) => Strong Tensor (Star (f :: Type -> Type)) where
  act (Star k) = Star (\(a, x) -> map (a,) (k x))
instance (Applicative f) => Strong CoprodAction (Star (f :: Type -> Type)) where
  act (Star k) = Star (f ||| map P.Right . k)
    where
      f a = pure (\() -> P.Left a) ()

instance (P.Applicative f) => Strong (SubAction P.Traversable ApplyAction) (Star (Prelude f)) where
  act (Star f) = Star (P.traverse f)

instance Traversable (Star P.Maybe) where
  traverse (Star a2mb :.: p) = lmap a2mb go :.: Star id
    where
      go =
        dimap
          (P.maybe (P.Left ()) P.Right)
          (P.const P.Nothing ||| P.Just)
          (one ++ p)

instance Traversable (Star []) where
  traverse (Star a2bs :.: p) = lmap a2bs go :.: Star id
    where
      go =
        dimap
          (\case [] -> P.Left (); (x : xs) -> P.Right (x, xs))
          (P.const [] ||| P.uncurry (:))
          (one ++ (p ** go))

-- | The list monad without the 'Prelude' wrapper, so that a Kleisli arrow of @[]@ reads as the
-- plain @a -> [b]@.
instance Promonad (Star []) where
  id = Star P.return
  Star l . Star r = Star (l P.<=< r)

starTraverse
  :: forall t f a b
   . ( Applicative (f :: Type -> Type)
     , Functor t
     , Traversable (Star t)
     , Ob b
     )
  => (a ~> f b) -> t a ~> f (t b)
starTraverse = baseTraverse @(Star t) @(Star f)

instance (Functor f, Thin k) => ThinProfunctor (Star f :: j +-> k) where
  type HasArrow (Star f :: j +-> k) a b = HasArrow (Hom k) a (f b)
  arr = Star arr
  withArr (Star f) r = withArr f r

instance (Functor f, DecidableProfunctor (Hom k)) => DecidableProfunctor (Star f :: j +-> k) where
  type Holds (Star f :: j +-> k) a b = Holds (Hom k) a (f b)
  decide @a @b = withObF @f @b (mapDecision Star (decide @(Hom k) @a @(f b)))
  toHolds (Star f) r = toHolds f r