{-# LANGUAGE PolyKinds    #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wno-orphans #-}
module Barbies.Internal.DistributiveB
  ( DistributiveB(..)
  , bdistribute'
  , bcotraverse
  , bdecompose
  , brecompose
  , gbdistributeDefault
  , CanDeriveDistributiveB
  )

where

import Barbies.Internal.FunctorB (FunctorB(..))
import Barbies.Generics.Distributive (GDistributive(..))

import Data.Functor.Compose   (Compose (..))
import Data.Functor.Identity  (Identity (..))
import Data.Functor.Product   (Product (..))
import Data.Generics.GenericN
import Data.Proxy             (Proxy (..))
import Data.Distributive
import Data.Kind              (Type)

-- | A 'FunctorB' where the effects can be distributed to the fields:
--  `bdistribute` turns an effectful way of building a Barbie-type
--  into a pure Barbie-type with effectful ways of computing the
--  values of its fields.
--
--  This class is the categorical dual of `Barbies.Internal.TraversableB.TraversableB`,
--  with `bdistribute` the dual of `Barbies.Internal.TraversableB.bsequence`
--  and `bcotraverse` the dual of `Barbies.Internal.TraversableB.btraverse`. As such,
--  instances need to satisfy these laws:
--
-- @
-- 'bdistribute' . h = 'bmap' ('Compose' . h . 'getCompose') . 'bdistribute'    -- naturality
-- 'bdistribute' . 'Data.Functor.Identity' = 'bmap' ('Compose' . 'Data.Functor.Identity')                 -- identity
-- 'bdistribute' . 'Compose' = 'bmap' ('Compose' . 'Compose' . 'fmap' 'getCompose' . 'getCompose') . 'bdistribute' . 'fmap' 'bdistribute' -- composition
-- @
--
-- By specializing @f@ to @((->) a)@ and @g@ to 'Identity', we can define a function that
-- decomposes a function on distributive barbies into a collection of simpler functions:
--
-- @
-- 'bdecompose' :: 'DistributiveB' b => (a -> b 'Identity') -> b ((->) a)
-- 'bdecompose' = 'bmap' ('fmap' 'runIdentity' . 'getCompose') . 'bdistribute'
-- @
--
-- Lawful instances of the class can then be characterized as those that satisfy:
--
-- @
-- 'brecompose' . 'bdecompose' = 'id'
-- 'bdecompose' . 'brecompose' = 'id'
-- @
--
-- This means intuitively that instances need to have a fixed shape (i.e. no sum-types can be involved).
-- Typically, this means record types, as long as they don't contain fields where the functor argument is not applied.
--
--
-- There is a default implementation of 'bdistribute' based on
-- 'Generic'.  Intuitively, it works on product types where the shape
-- of a pure value is uniquely defined and every field is covered by
-- the argument @f@.
class (FunctorB b) => DistributiveB (b :: (k -> Type) -> Type) where
  bdistribute :: Functor f => f (b g) -> b (Compose f g)

  default bdistribute
    :: forall f g
    .  CanDeriveDistributiveB b f g
    => Functor f => f (b g) -> b (Compose f g)
  bdistribute = forall {k1} (b :: (k1 -> *) -> *) (f :: * -> *) (g :: k1 -> *).
(CanDeriveDistributiveB b f g, Functor f) =>
f (b g) -> b (Compose f g)
gbdistributeDefault


-- | A version of `bdistribute` with @g@ specialized to `Identity`.
bdistribute' :: (DistributiveB b, Functor f) => f (b Identity) -> b f
bdistribute' :: forall (b :: (* -> *) -> *) (f :: * -> *).
(DistributiveB b, Functor f) =>
f (b Identity) -> b f
bdistribute' = forall k (b :: (k -> *) -> *) (f :: k -> *) (g :: k -> *).
FunctorB b =>
(forall (a :: k). f a -> g a) -> b f -> b g
bmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute

-- | Dual of `Barbies.Internal.TraversableB.btraverse`
bcotraverse :: (DistributiveB b, Functor f) => (forall a . f (g a) -> f a) -> f (b g) -> b f
bcotraverse :: forall (b :: (* -> *) -> *) (f :: * -> *) (g :: * -> *).
(DistributiveB b, Functor f) =>
(forall a. f (g a) -> f a) -> f (b g) -> b f
bcotraverse forall a. f (g a) -> f a
h = forall k (b :: (k -> *) -> *) (f :: k -> *) (g :: k -> *).
FunctorB b =>
(forall (a :: k). f a -> g a) -> b f -> b g
bmap (forall a. f (g a) -> f a
h forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute

-- | Decompose a function returning a distributive barbie, into
--   a collection of simpler functions.
bdecompose :: DistributiveB b => (a -> b Identity) -> b ((->) a)
bdecompose :: forall (b :: (* -> *) -> *) a.
DistributiveB b =>
(a -> b Identity) -> b ((->) a)
bdecompose = forall (b :: (* -> *) -> *) (f :: * -> *).
(DistributiveB b, Functor f) =>
f (b Identity) -> b f
bdistribute'

-- | Recompose a decomposed function.
brecompose :: FunctorB b => b ((->) a) -> a -> b Identity
brecompose :: forall (b :: (* -> *) -> *) a.
FunctorB b =>
b ((->) a) -> a -> b Identity
brecompose b ((->) a)
bfs = \a
a -> forall k (b :: (k -> *) -> *) (f :: k -> *) (g :: k -> *).
FunctorB b =>
(forall (a :: k). f a -> g a) -> b f -> b g
bmap (forall a. a -> Identity a
Identity forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a b. (a -> b) -> a -> b
$ a
a)) b ((->) a)
bfs

-- | @'CanDeriveDistributiveB' B f g@ is in practice a predicate about @B@ only.
--   Intuitively, it says the the following holds  for any arbitrary @f@:
--
--     * There is an instance of @'Generic' (B f)@.
--
--     * @(B f)@ has only one constructor, and doesn't contain "naked" fields
--       (that is, not covered by `f`).
--
--     * @B f@ can contain fields of type @b f@ as long as there exists a
--       @'DistributiveB' b@ instance. In particular, recursive usages of @B f@
--       are allowed.
--
--     * @B f@ can also contain usages of @b f@ under a @'Distributive' h@.
--       For example, one could use @a -> (B f)@ as a field of @B f@.
type CanDeriveDistributiveB b f g
  = ( GenericP 0 (b g)
    , GenericP 0 (b (Compose f g))
    , GDistributive 0 f (RepP 0 (b g)) (RepP 0 (b (Compose f g)))
    )

-- | Default implementation of 'bdistribute' based on 'Generic'.
gbdistributeDefault
  :: CanDeriveDistributiveB b f g
  => Functor f => f (b g) -> b (Compose f g)
gbdistributeDefault :: forall {k1} (b :: (k1 -> *) -> *) (f :: * -> *) (g :: k1 -> *).
(CanDeriveDistributiveB b f g, Functor f) =>
f (b g) -> b (Compose f g)
gbdistributeDefault
  = forall (n :: Nat) a x. GenericP n a => Proxy n -> RepP n a x -> a
toP (forall {k} (t :: k). Proxy t
Proxy @0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (n :: Nat) (f :: * -> *) (repbg :: k -> *)
       (repbfg :: k -> *) (x :: k).
GDistributive n f repbg repbfg =>
Proxy n -> f (repbg x) -> repbfg x
gdistribute (forall {k} (t :: k). Proxy t
Proxy @0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (n :: Nat) a x. GenericP n a => Proxy n -> a -> RepP n a x
fromP (forall {k} (t :: k). Proxy t
Proxy @0))
{-# INLINE gbdistributeDefault #-}

-- ------------------------------------------------------------
-- Generic derivation: Special cases for DistributiveB
-- -----------------------------------------------------------

type P = Param

instance
  ( Functor f
  , DistributiveB b
  ) => GDistributive 0 f (Rec (b' (P 0 g)) (b g)) (Rec (b' (P 0 (Compose f g))) (b (Compose f g)))
  where
  gdistribute :: forall (x :: k).
Proxy 0
-> f (Rec (b' (P 0 g)) (b g) x)
-> Rec (b' (P 0 (Compose f g))) (b (Compose f g)) x
gdistribute Proxy 0
_ = forall {k} p a (x :: k). K1 R a x -> Rec p a x
Rec forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k i c (p :: k). c -> K1 i c p
K1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall k i c (p :: k). K1 i c p -> c
unK1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} p a (x :: k). Rec p a x -> K1 R a x
unRec)
  {-# INLINE gdistribute #-}


instance
  ( Functor f
  , Distributive h
  , DistributiveB b
  ) =>
  GDistributive n f (Rec (h (b (P n g))) (h (b g))) (Rec (h (b (P n (Compose f g)))) (h (b (Compose f g))))
  where
  gdistribute :: forall (x :: k).
Proxy n
-> f (Rec (h (b (P n g))) (h (b g)) x)
-> Rec (h (b (P n (Compose f g)))) (h (b (Compose f g))) x
gdistribute Proxy n
_ = forall {k} p a (x :: k). K1 R a x -> Rec p a x
Rec forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k i c (p :: k). c -> K1 i c p
K1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (g :: * -> *) (f :: * -> *) a.
(Distributive g, Functor f) =>
f (g a) -> g (f a)
distribute forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall k i c (p :: k). K1 i c p -> c
unK1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} p a (x :: k). Rec p a x -> K1 R a x
unRec)
  {-# INLINE gdistribute #-}

-- --------------------------------
-- Instances for base types
-- --------------------------------

instance DistributiveB Proxy where
  bdistribute :: forall (f :: * -> *) (g :: k -> *).
Functor f =>
f (Proxy g) -> Proxy (Compose f g)
bdistribute f (Proxy g)
_ = forall {k} (t :: k). Proxy t
Proxy
  {-# INLINE bdistribute #-}

fstF :: Product f g a -> f a
fstF :: forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
Product f g a -> f a
fstF (Pair f a
x g a
_y) = f a
x

sndF :: Product f g a -> g a
sndF :: forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
Product f g a -> g a
sndF (Pair f a
_x g a
y) = g a
y

instance (DistributiveB a, DistributiveB b) => DistributiveB (Product a b) where
  bdistribute :: forall (f :: * -> *) (g :: k -> *).
Functor f =>
f (Product a b g) -> Product a b (Compose f g)
bdistribute f (Product a b g)
xy = forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
f a -> g a -> Product f g a
Pair (forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute forall a b. (a -> b) -> a -> b
$ forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
Product f g a -> f a
fstF forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Product a b g)
xy) (forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute forall a b. (a -> b) -> a -> b
$ forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
Product f g a -> g a
sndF forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Product a b g)
xy)
  {-# INLINE bdistribute #-}

instance (Distributive h, DistributiveB b) => DistributiveB (h `Compose` b) where
  bdistribute :: forall (f :: * -> *) (g :: k -> *).
Functor f =>
f (Compose h b g) -> Compose h b (Compose f g)
bdistribute = forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (g :: * -> *) (f :: * -> *) a.
(Distributive g, Functor f) =>
f (g a) -> g (f a)
distribute forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose
  {-# INLINE bdistribute #-}