PowerGrate.hs

{-# LANGUAGE AllowAmbiguousTypes #-}

-- | A __power grate__ is a 'Proarrow.Optic.Grate.Grate' whose exponent is a fixed tensor /power/ of
-- the focus: the witness @'Pow' n@ presents @s@ as @a ** ... ** a@ (@n@ times), the reader
-- applicative for @n@ readers. A fixed finite shape is a decomposition, so a power grate is also a
-- __fixed-arity 'Proarrow.Optic.Traversal.Traversal'__
-- (@'PowerGrateFl' <: 'GrateFl', 'KaleidoFl', 'MonTravFl'@).
--
-- It adds an eliminator: 'powerGrateP' distributes any 'MonoidalProfunctor', using only 'one' and
-- '**'. A traversal needs a 'Proarrow.Category.Monoidal.Distributive.StrongDistributiveProfunctor',
-- and a kaleidoscope needs a traversable carrier. With a fixed arity /any/ @Costar f@ distributes,
-- by unzipping @f (a ** ... ** a)@ into @f a ** ... ** f a@. So 'powerGrateOf' works at any
-- monoidal profunctor carrier: the hom @('~>')@ gives 'Proarrow.Optic.Setter.over', and an
-- applicative @'Proarrow.Profunctor.Instance.Star.Star' f@ combines the foci through @f@.
module Proarrow.Optic.PowerGrate
  ( PowerGrateFl (..)
  , PowerGrate
  , PowerGrate'
  , powerGrateOf
  , zipWithOf

    -- * @n@-ary aggregation
  , Pow (..)
  , CoPow (..)
  , powerGrate
  ) where

import Data.Type.Nat (Nat, Nat2, SNat (..), SNatI, snat)
import Proarrow.Adjunction (Proadjunction (..))
import Proarrow.Category.Monoidal
  ( Monoidal (..)
  , MonoidalProfunctor (..)
  , NFold
  , SymMonoidal
  , swapInner
  , withObNFold
  , type (**)
  )
import Proarrow.Category.Monoidal qualified as M
import Proarrow.Category.Monoidal.Action (CoprodAction)
import Proarrow.Category.Monoidal.Cartesian (Cartesian)
import Proarrow.Category.Monoidal.Closed (Closed (..), mkExponential)
import Proarrow.Category.Monoidal.CopyDiscard (CopyDiscard (..), fst, snd, (&&&))
import Proarrow.Category.Monoidal.Distributive (Traversable (..))
import Proarrow.Category.Monoidal.Strength (Strong (..))
import Proarrow.Colimit.BinaryCoproduct (COPROD (..), Coprod (..), HasBinaryCoproducts (..), HasCoproducts)
import Proarrow.Colimit.Initial (HasInitialObject (..))
import Proarrow.Core (CategoryOf (..), Profunctor (..), Promonad (..), obj, (//), (\\), type (+->))
import Proarrow.Functor (Functor)
import Proarrow.Monoid (fanIn, fanOut)
import Proarrow.Object (pattern Objs)
import Proarrow.Optic
  ( ExOptic
  , FLAVOR
  , Optic
  , Prostrong (..)
  , legs2prof
  , withLegs
  )
import Proarrow.Optic.Fold (FoldFl (..))
import Proarrow.Optic.Glass (GlassFl (..), Mod, withObSel)
import Proarrow.Optic.Grate (GrateFl (..))
import Proarrow.Optic.Kaleidoscope (CotravFl, KaleidoFl (..), Kaleidoscopic (..), kaleidoscopeOf)
import Proarrow.Optic.Setter (SetterFl (..))
import Proarrow.Optic.Traversal (MonTravFl (..), TravFl (..))
import Proarrow.Profunctor.Instance.Composition ((:.:) (..))
import Proarrow.Profunctor.Instance.Costar (Costar)
import Proarrow.Profunctor.Instance.Identity (Id (..))
import Proarrow.Profunctor.Representable (RepCostar (..), Representable (..))

-- | The power-grate flavor: distribute any 'MonoidalProfunctor' @r@ through the witness
-- pair. 'Proarrow.Optic.Traversal.TravFl' is a superclass: every power-grate witness is a
-- traversal witness (instantiate @r@ at a 'Proarrow.Category.Monoidal.Distributive.StrongDistributiveProfunctor',
-- a special 'MonoidalProfunctor'), so it folds, sets, and traverses. The extra power is
-- distributing the /non/-SDP monoidal profunctors as well. 'Proarrow.Optic.Kaleidoscope.KaleidoFl'
-- is a superclass too: a tensor power is an applicative functor (the reader applicative).
type PowerGrateFl :: forall {k}. FLAVOR k k
class (MonTravFl p q, GrateFl p q) => PowerGrateFl (p :: k +-> k) (q :: k +-> k) where
  powerGrateP :: (MonoidalProfunctor r) => p s a -> q b t -> r a b -> r s t

instance (CategoryOf k) => PowerGrateFl (Id :: k +-> k) (Id :: k +-> k) where
  powerGrateP (Id l) (Id r) = dimap l r
instance (PowerGrateFl f g, PowerGrateFl f' g') => PowerGrateFl (f :.: f') (g' :.: g) where
  powerGrateP (f :.: f') (g' :.: g) = powerGrateP @f @g f g . powerGrateP @f' @g' f' g'

-- | The carrier of the literature's kaleidoscope eliminator (@>-@): @'Costar' f@, i.e. @f a -> b@ for
-- any functor @f@ on a cartesian category. Power grates distribute any 'MonoidalProfunctor', and
-- @'Costar' f@ is one, so this is 'powerGrateP' at that carrier. It is an instance (and not only
-- reachable through 'powerGrateOf') so that a power grate composed with another flavor that also
-- runs at @Costar f@, an algebraic lens say, can be eliminated there directly.
instance (Cartesian k, Functor (f :: k -> k)) => Prostrong PowerGrateFl (Costar f :: k +-> k) where
  proact (f :.: c :.: g) = powerGrateP f g c

type PowerGrate (s :: k) (t :: k) a b = Optic (Prostrong PowerGrateFl) s t a b
type PowerGrate' s a = PowerGrate s s a a

-- | Distribute any 'MonoidalProfunctor' through a power grate (or any stronger optic). At the
-- hom @('~>')@ this is 'Proarrow.Optic.Setter.over'; at an applicative @'Proarrow.Profunctor.Instance.Star.Star' f@ the foci are
-- combined through @f@.
--
-- Accepts any encoding (cf. 'Proarrow.Optic.Traversal.traverseOf'), including '(%)'-composites.
powerGrateOf
  :: forall {k} c (s :: k) (t :: k) a b r
   . (Monoidal k, MonoidalProfunctor r, (Ob a, Ob b) => c (ExOptic PowerGrateFl a b))
  => Optic c s t a b -> r a b -> r s t
powerGrateOf o rab = withLegs @PowerGrateFl o \l r -> powerGrateP l r rab

-- * @n@-ary aggregation via tensor powers

-- | Distribute a 'MonoidalProfunctor' over the @n@-fold tensor power, by combining @n@ copies of
-- the carrier value with 'one' (at 'Z') and '**' (at 'S'). This is the profunctor-general core
-- of the @n@-ary power grate.
powDist :: forall n r a b. (SNatI n, MonoidalProfunctor r) => r a b -> r (NFold n a) (NFold n b)
powDist rab = case snat @n of
  SZ -> one
  SS @m -> rab ** powDist @m rab

-- | Distribute the internal hom over the tensor power: split @x ~~> aⁿ@ into @(x ~~> a)ⁿ@ using
-- 'CopyDiscard' projections. This makes an @n@-ary power grate a 'Proarrow.Optic.Grate.Grate'.
splitPow
  :: forall n k (x :: k) a. (SNatI n, Closed k, CopyDiscard k, Ob x, Ob a) => (x ~~> NFold n a) ~> NFold n (x ~~> a)
splitPow = case snat @n of
  SZ -> withObExp @k @x @Unit (discard @k @(x ~~> Unit))
  SS @m ->
    withObNFold @m @a
      ((fst @a @(NFold m a) ^^^ obj @x) &&& (splitPow @m @k @x @a . (snd @a @(NFold m a) ^^^ obj @x)))

-- | Zip two tensor powers into the tensor power of the tensor: the @<*>@ of the reader
-- applicative @NFold n@.
powZip
  :: forall n k (a :: k) c. (SNatI n, SymMonoidal k, Ob a, Ob c) => (NFold n a ** NFold n c) ~> NFold n (a ** c)
powZip = case snat @n of
  SZ -> leftUnitor @k @Unit
  SS @m ->
    withObNFold @m @a
      (withObNFold @m @c (((obj @a ** obj @c) ** powZip @m @k @a @c) . swapInner @a @(NFold m a) @c @(NFold m c)))

-- | The tensor power of the unit is (isomorphic to) the unit.
powUnit :: forall n k. (SNatI n, Monoidal k) => Unit ~> NFold n (Unit :: k)
powUnit = case snat @n of
  SZ -> id
  SS @m -> (obj @(Unit :: k) ** powUnit @m @k) . leftUnitorInv @k @Unit

-- | The arity-@n@ aggregation witness: @s@ presents @n@ foci via the tensor power.
--
-- @'Pow' n@ is the representable profunctor of the tensor power @NFold n@, which is the reader
-- applicative for @n@ readers. Its instances make it a
-- 'Proarrow.Category.Monoidal.Distributive.StrongDistributiveProfunctor', hence a kaleidoscope
-- witness.
type Pow :: forall {k}. Nat -> k +-> k
data Pow n s a where
  Pow :: forall (n :: Nat) {k} (s :: k) (a :: k). (Ob a) => (s ~> NFold n a) -> Pow n s a

-- | The dual of 'Pow': @t@ is rebuilt from @n@ foci.
type CoPow :: forall {k}. Nat -> k +-> k
data CoPow n b t where
  CoPow :: forall (n :: Nat) {k} (b :: k) (t :: k). (Ob b) => (NFold n b ~> t) -> CoPow n b t

instance (Monoidal k, SNatI n) => Profunctor (Pow n :: k +-> k) where
  dimap l r (Pow sa) = Pow (powDist @n r . sa . l) \\ l \\ r
  r \\ Pow sa = r \\ sa
instance (Monoidal k, SNatI n) => Profunctor (CoPow n :: k +-> k) where
  dimap l r (CoPow bt) = CoPow (r . bt . powDist @n l) \\ l \\ r
  r \\ CoPow bt = r \\ bt

instance (Monoidal k, SNatI n) => SetterFl (Pow n :: k +-> k) (CoPow n :: k +-> k) where
  overP (Pow sl) (CoPow rt) f = rt . powDist @n f . sl
instance (Monoidal k, SNatI n) => FoldFl (Pow n :: k +-> k) (CoPow n :: k +-> k) where
  foldMapP (Pow sl) am = fanIn @n . powDist @n am . sl
instance (Monoidal k, SNatI n) => TravFl (Pow n :: k +-> k) (CoPow n :: k +-> k)
instance (Monoidal k, SNatI n) => MonTravFl (Pow n :: k +-> k) (CoPow n :: k +-> k) where
  monTravP (Pow sl) (CoPow rt) rab = dimap sl rt (powDist @n rab)

-- | A power grate is a glass: ignore the source, and for each of the @n@ positions feed the
-- consumer the selector "project this focus". The selectors come from @splitPow@ of @sl@, the
-- consumer is copied @n@ times with 'fanOut', @powZip@ pairs them, and @powDist@ applies each.
-- Everything is stated with the 'CopyDiscard' structure that 'CCC' provides, so the tensor
-- and the product never have to be identified by hand.
instance (Monoidal k, HasCoproducts k, SNatI n) => GlassFl (Pow n :: k +-> k) (CoPow n :: k +-> k) where
  glassP @s @a @b (Pow sl@Objs) (CoPow rt@Objs) =
    withObSel @s @a @b
      ( withOb2 @k @s @(Mod s a b)
          ( rt
              . powDist @n (apply @k @(s ~~> a) @b)
              . powZip @n @k @(Mod s a b) @(s ~~> a)
              . ( (fanOut @n @(Mod s a b) . snd @s @(Mod s a b))
                    &&& (splitPow @n @k @s @a . mkExponential sl . discard @k @(s ** Mod s a b))
                )
          )
      )

instance (CopyDiscard k, HasCoproducts k, SNatI n) => GrateFl (Pow n :: k +-> k) (CoPow n :: k +-> k) where
  zipWithP (Pow @_ @_ @a sl) (CoPow rt) @x kk = rt . powDist @n kk . splitPow @n @_ @x @a . (sl ^^^ obj @x)
instance (CopyDiscard k, HasCoproducts k, SNatI n) => PowerGrateFl (Pow n :: k +-> k) (CoPow n :: k +-> k) where
  powerGrateP (Pow sl) (CoPow rt) rab = dimap sl rt (powDist @n rab)

-- | A tensor power is a fixed-shape traversable: distribute the carrier over the @n@ copies.
instance (Monoidal k, SNatI n) => Traversable (Pow n :: k +-> k) where
  traverse @_ @_ @b (Pow f :.: p) = p // withObNFold @n @b (lmap f (powDist @n p) :.: Pow id)

instance (Monoidal k, SNatI n) => Representable (Pow n :: k +-> k) where
  type Pow n % a = NFold n a
  index (Pow f) = f
  tabulate = Pow
  repMap = powDist @n

instance (SymMonoidal k, SNatI n) => MonoidalProfunctor (Pow n :: k +-> k) where
  one = Pow (powUnit @n)
  Pow @_ @_ @a f ** Pow @_ @_ @c g = f // g // withOb2 @k @a @c (Pow (powZip @n @k @a @c . (f ** g)))
instance (SymMonoidal k, HasCoproducts k, SNatI n) => MonoidalProfunctor (Coprod (Pow n :: k +-> k)) where
  one = withObNFold @n @(InitialObject :: k) (Coprod (Pow initiate))
  Coprod (Pow @_ @_ @a f) ** Coprod (Pow @_ @_ @c g) =
    withObCoprod @k @a @c (Coprod (Pow (powDist @n (lft @k @a @c) . f ||| powDist @n (rgt @k @a @c) . g)))
instance (CopyDiscard k, SNatI n) => Strong M.Tensor (Pow n :: k +-> k) where
  act @x (Pow @_ @_ @a f) = f // withOb2 @k @x @a (Pow (powZip @n @k @x @a . (fanOut @n @x ** f)))
instance (CopyDiscard k, HasCoproducts k, SNatI n) => Strong CoprodAction (Pow n :: k +-> k) where
  act @(COPR x) (Pow @_ @_ @a f) =
    f // withObCoprod @k @x @a (Pow (powDist @n (lft @k @x @a) . fanOut @n @x ||| powDist @n (rgt @k @x @a) . f))
instance (CopyDiscard k, HasCoproducts k, SNatI n) => CotravFl (Pow n :: k +-> k) (CoPow n :: k +-> k)
instance (CopyDiscard k, HasCoproducts k, SNatI n) => KaleidoFl (Pow n :: k +-> k) (CoPow n :: k +-> k) where
  kaleidoP (Pow sl) (CoPow rt) rab = dimap sl rt (kaleidoAct @_ @(Pow n) rab)
instance (Monoidal k, SNatI n) => Proadjunction (Pow n :: k +-> k) (CoPow n) where
  unit @x = (CoPow id :.: Pow id) \\ powDist @n (id :: x ~> x)
  counit (Pow sl :.: CoPow rt) = rt . sl

-- | Build an @n@-ary power grate from a tensor-power decomposition of @s@ and recomposition
-- of @t@.
powerGrate
  :: forall {k} (n :: Nat) (s :: k) (t :: k) a b
   . (CopyDiscard k, HasCoproducts k, SNatI n, Ob a, Ob b)
  => (s ~> NFold n a) -> (NFold n b ~> t) -> PowerGrate s t a b
powerGrate sl rt = legs2prof @PowerGrateFl (Pow @n sl) (CoPow @n rt)

-- | Zip two sources through a 'Proarrow.Optic.Kaleidoscope.Kaleidoscope' (or any stronger optic, a
-- 'Proarrow.Optic.Grate.Grate' in particular, in any encoding): combine the foci pairwise. This is
-- 'kaleidoscopeOf' at the carrier @'RepCostar' ('Pow' 2)@, the costar of the binary tensor power:
-- a binary combination @(a ** a) ~> b@ of foci, which the optic's applicative lifts by @liftA2@.
zipWithOf
  :: forall {k} c (s :: k) (t :: k) a b
   . (Monoidal k, Ob a, (Ob a, Ob b) => c (ExOptic KaleidoFl a b))
  => Optic c s t a b -> ((a ** a) ~> b) -> (s ** s) ~> t
zipWithOf o f =
  case kaleidoscopeOf o (RepCostar @_ @(Pow Nat2) (f . (obj @a ** rightUnitor @k @a))) of
    RepCostar @s' g -> g . (obj @s' ** rightUnitorInv @k @s')