{-# 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 = f (b g) -> b (Compose f g)
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' :: f (b Identity) -> b f
bdistribute' = (forall a. Compose f Identity a -> f a)
-> b (Compose f Identity) -> b f
forall k (b :: (k -> *) -> *) (f :: k -> *) (g :: k -> *).
FunctorB b =>
(forall (a :: k). f a -> g a) -> b f -> b g
bmap ((Identity a -> a) -> f (Identity a) -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Identity a -> a
forall a. Identity a -> a
runIdentity (f (Identity a) -> f a)
-> (Compose f Identity a -> f (Identity a))
-> Compose f Identity a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Compose f Identity a -> f (Identity a)
forall k1 (f :: k1 -> *) k2 (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose) (b (Compose f Identity) -> b f)
-> (f (b Identity) -> b (Compose f Identity))
-> f (b Identity)
-> b f
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (b Identity) -> b (Compose f Identity)
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 a. f (g a) -> f a) -> f (b g) -> b f
bcotraverse forall a. f (g a) -> f a
h = (forall a. Compose f g a -> f a) -> b (Compose f g) -> b f
forall k (b :: (k -> *) -> *) (f :: k -> *) (g :: k -> *).
FunctorB b =>
(forall (a :: k). f a -> g a) -> b f -> b g
bmap (f (g a) -> f a
forall a. f (g a) -> f a
h (f (g a) -> f a)
-> (Compose f g a -> f (g a)) -> Compose f g a -> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Compose f g a -> f (g a)
forall k1 (f :: k1 -> *) k2 (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose) (b (Compose f g) -> b f)
-> (f (b g) -> b (Compose f g)) -> f (b g) -> b f
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (b g) -> b (Compose f g)
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 :: (a -> b Identity) -> b ((->) a)
bdecompose = (a -> b Identity) -> b ((->) a)
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 :: b ((->) a) -> a -> b Identity
brecompose b ((->) a)
bfs = \a
a -> (forall a. (a -> a) -> Identity a) -> b ((->) a) -> b Identity
forall k (b :: (k -> *) -> *) (f :: k -> *) (g :: k -> *).
FunctorB b =>
(forall (a :: k). f a -> g a) -> b f -> b g
bmap (a -> Identity a
forall a. a -> Identity a
Identity (a -> Identity a) -> ((a -> a) -> a) -> (a -> a) -> Identity a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a -> a) -> a -> a
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 :: f (b g) -> b (Compose f g)
gbdistributeDefault
  = Proxy 0 -> RepP 0 (b (Compose f g)) Any -> b (Compose f g)
forall (n :: Nat) a x. GenericP n a => Proxy n -> RepP n a x -> a
toP (Proxy 0
forall k (t :: k). Proxy t
Proxy @0) (Zip
   (Rep (FilterIndex 0 (Indexed b 1) (Param 0 (Compose f g))))
   (Rep (b (Compose f g)))
   Any
 -> b (Compose f g))
-> (f (b g)
    -> Zip
         (Rep (FilterIndex 0 (Indexed b 1) (Param 0 (Compose f g))))
         (Rep (b (Compose f g)))
         Any)
-> f (b g)
-> b (Compose f g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy 0
-> f (Zip
        (Rep (FilterIndex 0 (Indexed b 1) (Param 0 g))) (Rep (b g)) Any)
-> Zip
     (Rep (FilterIndex 0 (Indexed b 1) (Param 0 (Compose f g))))
     (Rep (b (Compose f g)))
     Any
forall k (n :: Nat) (f :: * -> *) (repbg :: k -> *)
       (repbfg :: k -> *) (x :: k).
GDistributive n f repbg repbfg =>
Proxy n -> f (repbg x) -> repbfg x
gdistribute (Proxy 0
forall k (t :: k). Proxy t
Proxy @0) (f (Zip
      (Rep (FilterIndex 0 (Indexed b 1) (Param 0 g))) (Rep (b g)) Any)
 -> Zip
      (Rep (FilterIndex 0 (Indexed b 1) (Param 0 (Compose f g))))
      (Rep (b (Compose f g)))
      Any)
-> (f (b g)
    -> f (Zip
            (Rep (FilterIndex 0 (Indexed b 1) (Param 0 g))) (Rep (b g)) Any))
-> f (b g)
-> Zip
     (Rep (FilterIndex 0 (Indexed b 1) (Param 0 (Compose f g))))
     (Rep (b (Compose f g)))
     Any
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b g
 -> Zip
      (Rep (FilterIndex 0 (Indexed b 1) (Param 0 g))) (Rep (b g)) Any)
-> f (b g)
-> f (Zip
        (Rep (FilterIndex 0 (Indexed b 1) (Param 0 g))) (Rep (b g)) Any)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Proxy 0 -> b g -> RepP 0 (b g) Any
forall (n :: Nat) a x. GenericP n a => Proxy n -> a -> RepP n a x
fromP (Proxy 0
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 :: 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
_ = K1 R (b (Compose f g)) x
-> Rec (b' (P 0 (Compose f g))) (b (Compose f g)) x
forall k p a (x :: k). K1 R a x -> Rec p a x
Rec (K1 R (b (Compose f g)) x
 -> Rec (b' (P 0 (Compose f g))) (b (Compose f g)) x)
-> (f (Rec (b' (P 0 g)) (b g) x) -> K1 R (b (Compose f g)) x)
-> f (Rec (b' (P 0 g)) (b g) x)
-> Rec (b' (P 0 (Compose f g))) (b (Compose f g)) x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b (Compose f g) -> K1 R (b (Compose f g)) x
forall k i c (p :: k). c -> K1 i c p
K1 (b (Compose f g) -> K1 R (b (Compose f g)) x)
-> (f (Rec (b' (P 0 g)) (b g) x) -> b (Compose f g))
-> f (Rec (b' (P 0 g)) (b g) x)
-> K1 R (b (Compose f g)) x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (b g) -> b (Compose f g)
forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute (f (b g) -> b (Compose f g))
-> (f (Rec (b' (P 0 g)) (b g) x) -> f (b g))
-> f (Rec (b' (P 0 g)) (b g) x)
-> b (Compose f g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Rec (b' (P 0 g)) (b g) x -> b g)
-> f (Rec (b' (P 0 g)) (b g) x) -> f (b g)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (K1 R (b g) x -> b g
forall i c k (p :: k). K1 i c p -> c
unK1 (K1 R (b g) x -> b g)
-> (Rec (b' (P 0 g)) (b g) x -> K1 R (b g) x)
-> Rec (b' (P 0 g)) (b g) x
-> b g
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rec (b' (P 0 g)) (b g) x -> K1 R (b g) x
forall p a k (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 :: 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
_ = K1 R (h (b (Compose f g))) x
-> Rec (h (b (P n (Compose f g)))) (h (b (Compose f g))) x
forall k p a (x :: k). K1 R a x -> Rec p a x
Rec (K1 R (h (b (Compose f g))) x
 -> Rec (h (b (P n (Compose f g)))) (h (b (Compose f g))) x)
-> (f (Rec (h (b (P n g))) (h (b g)) x)
    -> K1 R (h (b (Compose f g))) x)
-> 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
forall b c a. (b -> c) -> (a -> b) -> a -> c
. h (b (Compose f g)) -> K1 R (h (b (Compose f g))) x
forall k i c (p :: k). c -> K1 i c p
K1 (h (b (Compose f g)) -> K1 R (h (b (Compose f g))) x)
-> (f (Rec (h (b (P n g))) (h (b g)) x) -> h (b (Compose f g)))
-> f (Rec (h (b (P n g))) (h (b g)) x)
-> K1 R (h (b (Compose f g))) x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (f (b g) -> b (Compose f g)) -> h (f (b g)) -> h (b (Compose f g))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap f (b g) -> b (Compose f g)
forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute (h (f (b g)) -> h (b (Compose f g)))
-> (f (Rec (h (b (P n g))) (h (b g)) x) -> h (f (b g)))
-> f (Rec (h (b (P n g))) (h (b g)) x)
-> h (b (Compose f g))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (h (b g)) -> h (f (b g))
forall (g :: * -> *) (f :: * -> *) a.
(Distributive g, Functor f) =>
f (g a) -> g (f a)
distribute (f (h (b g)) -> h (f (b g)))
-> (f (Rec (h (b (P n g))) (h (b g)) x) -> f (h (b g)))
-> f (Rec (h (b (P n g))) (h (b g)) x)
-> h (f (b g))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Rec (h (b (P n g))) (h (b g)) x -> h (b g))
-> f (Rec (h (b (P n g))) (h (b g)) x) -> f (h (b g))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (K1 R (h (b g)) x -> h (b g)
forall i c k (p :: k). K1 i c p -> c
unK1 (K1 R (h (b g)) x -> h (b g))
-> (Rec (h (b (P n g))) (h (b g)) x -> K1 R (h (b g)) x)
-> Rec (h (b (P n g))) (h (b g)) x
-> h (b g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rec (h (b (P n g))) (h (b g)) x -> K1 R (h (b g)) x
forall p a k (x :: k). Rec p a x -> K1 R a x
unRec)
  {-# INLINE gdistribute #-}

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

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

fstF :: Product f g a -> f a
fstF :: 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 :: 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 :: f (Product a b g) -> Product a b (Compose f g)
bdistribute f (Product a b g)
xy = a (Compose f g) -> b (Compose f g) -> Product a b (Compose f g)
forall k (f :: k -> *) (g :: k -> *) (a :: k).
f a -> g a -> Product f g a
Pair (f (a g) -> a (Compose f g)
forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute (f (a g) -> a (Compose f g)) -> f (a g) -> a (Compose f g)
forall a b. (a -> b) -> a -> b
$ Product a b g -> a g
forall k (f :: k -> *) (g :: k -> *) (a :: k). Product f g a -> f a
fstF (Product a b g -> a g) -> f (Product a b g) -> f (a g)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Product a b g)
xy) (f (b g) -> b (Compose f g)
forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute (f (b g) -> b (Compose f g)) -> f (b g) -> b (Compose f g)
forall a b. (a -> b) -> a -> b
$ Product a b g -> b g
forall k (f :: k -> *) (g :: k -> *) (a :: k). Product f g a -> g a
sndF (Product a b g -> b g) -> f (Product a b g) -> f (b g)
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 :: f (Compose h b g) -> Compose h b (Compose f g)
bdistribute = h (b (Compose f g)) -> Compose h b (Compose f g)
forall k k1 (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (h (b (Compose f g)) -> Compose h b (Compose f g))
-> (f (Compose h b g) -> h (b (Compose f g)))
-> f (Compose h b g)
-> Compose h b (Compose f g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (f (b g) -> b (Compose f g)) -> h (f (b g)) -> h (b (Compose f g))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap f (b g) -> b (Compose f g)
forall k (b :: (k -> *) -> *) (f :: * -> *) (g :: k -> *).
(DistributiveB b, Functor f) =>
f (b g) -> b (Compose f g)
bdistribute (h (f (b g)) -> h (b (Compose f g)))
-> (f (Compose h b g) -> h (f (b g)))
-> f (Compose h b g)
-> h (b (Compose f g))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (h (b g)) -> h (f (b g))
forall (g :: * -> *) (f :: * -> *) a.
(Distributive g, Functor f) =>
f (g a) -> g (f a)
distribute (f (h (b g)) -> h (f (b g)))
-> (f (Compose h b g) -> f (h (b g)))
-> f (Compose h b g)
-> h (f (b g))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Compose h b g -> h (b g)) -> f (Compose h b g) -> f (h (b g))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Compose h b g -> h (b g)
forall k1 (f :: k1 -> *) k2 (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose
  {-# INLINE bdistribute #-}