{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
#if __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
#endif
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Distributive
-- Copyright   :  (C) 2011-2016 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
----------------------------------------------------------------------------
module Data.Distributive.Generic
  ( GDistributive(..)
  , genericCollect
  , genericDistribute
  ) where

import Data.Distributive
import GHC.Generics
import Data.Coerce

-- | 'collect' derived from a 'Generic1' type
--
-- This can be used to easily produce a 'Distributive' instance for a
-- type with a 'Generic1' instance,
--
-- > data V2 a = V2 a a deriving (Show, Functor, Generic1)
-- > instance Distributive V2' where collect = genericCollect
genericCollect :: (Functor f, Generic1 g, GDistributive (Rep1 g))
               => (a -> g b) -> f a -> g (f b)
genericCollect :: forall (f :: * -> *) (g :: * -> *) a b.
(Functor f, Generic1 g, GDistributive (Rep1 g)) =>
(a -> g b) -> f a -> g (f b)
genericCollect a -> g b
f = forall k (f :: k -> *) (a :: k). Generic1 f => Rep1 f a -> f a
to1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (g :: * -> *) (f :: * -> *) a b.
(GDistributive g, Functor f) =>
(a -> g b) -> f a -> g (f b)
gcollect (forall k (f :: k -> *) (a :: k). Generic1 f => f a -> Rep1 f a
from1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> g b
f)

-- | 'distribute' derived from a 'Generic1' type
--
-- It's often more efficient to use 'genericCollect' instead.
genericDistribute  :: (Functor f, Generic1 g, GDistributive (Rep1 g)) => f (g a) -> g (f a)
genericDistribute :: forall (f :: * -> *) (g :: * -> *) a.
(Functor f, Generic1 g, GDistributive (Rep1 g)) =>
f (g a) -> g (f a)
genericDistribute = forall k (f :: k -> *) (a :: k). Generic1 f => Rep1 f a -> f a
to1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (g :: * -> *) (f :: * -> *) b.
(GDistributive g, Functor f) =>
f (g b) -> g (f b)
gdistribute 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 (f :: k -> *) (a :: k). Generic1 f => f a -> Rep1 f a
from1


-- Can't distribute over,
--   * sums (:+:)
--   * K1
--   * V1
class GDistributive g where
  gcollect :: Functor f => (a -> g b) -> f a -> g (f b)

gdistribute :: (GDistributive g, Functor f) => f (g b) -> g (f b)
gdistribute :: forall (g :: * -> *) (f :: * -> *) b.
(GDistributive g, Functor f) =>
f (g b) -> g (f b)
gdistribute = forall (g :: * -> *) (f :: * -> *) a b.
(GDistributive g, Functor f) =>
(a -> g b) -> f a -> g (f b)
gcollect forall a. a -> a
id
{-# INLINE gdistribute #-}

instance GDistributive U1 where
  gcollect :: forall (f :: * -> *) a b.
Functor f =>
(a -> U1 b) -> f a -> U1 (f b)
gcollect a -> U1 b
_ f a
_ = forall k (p :: k). U1 p
U1
  {-# INLINE gcollect #-}

instance (GDistributive a, GDistributive b) => GDistributive (a :*: b) where
  -- It might be tempting to fuse `gcollect fstP (fmap f x)` into
  -- `gcollect (fstP . f) x`, but this would lead to a loss of sharing.
  gcollect :: forall (f :: * -> *) a b.
Functor f =>
(a -> (:*:) a b b) -> f a -> (:*:) a b (f b)
gcollect a -> (:*:) a b b
f f a
x = forall (g :: * -> *) (f :: * -> *) a b.
(GDistributive g, Functor f) =>
(a -> g b) -> f a -> g (f b)
gcollect forall {f :: * -> *} {g :: * -> *} {p}. (:*:) f g p -> f p
fstP f ((:*:) a b b)
x' forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*: forall (g :: * -> *) (f :: * -> *) a b.
(GDistributive g, Functor f) =>
(a -> g b) -> f a -> g (f b)
gcollect forall {f :: * -> *} {g :: * -> *} {p}. (:*:) f g p -> g p
sndP f ((:*:) a b b)
x' where
    x' :: f ((:*:) a b b)
x' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> (:*:) a b b
f f a
x
    fstP :: (:*:) f g p -> f p
fstP (f p
l :*: g p
_) = f p
l
    sndP :: (:*:) f g p -> g p
sndP (f p
_ :*: g p
r) = g p
r
  {-# INLINE gcollect #-}

instance (Distributive a, GDistributive b) => GDistributive (a :.: b) where
  gcollect :: forall (f :: * -> *) a b.
Functor f =>
(a -> (:.:) a b b) -> f a -> (:.:) a b (f b)
gcollect a -> (:.:) a b b
f = forall k2 k1 (f :: k2 -> *) (g :: k1 -> k2) (p :: k1).
f (g p) -> (:.:) f g p
Comp1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (g :: * -> *) (f :: * -> *) b.
(GDistributive g, Functor f) =>
f (g b) -> g (f b)
gdistribute forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (g :: * -> *) (f :: * -> *) a b.
(Distributive g, Functor f) =>
(a -> g b) -> f a -> g (f b)
collect (coerce :: forall a b. Coercible a b => a -> b
coerce a -> (:.:) a b b
f)
  {-# INLINE gcollect #-}

instance GDistributive Par1 where
  gcollect :: forall (f :: * -> *) a b.
Functor f =>
(a -> Par1 b) -> f a -> Par1 (f b)
gcollect = coerce :: forall a b. Coercible a b => a -> b
coerce (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap :: (a -> b) -> f a -> f b)
    :: forall f a b . Functor f => (a -> Par1 b) -> f a -> Par1 (f b)
  {-# INLINE gcollect #-}

instance Distributive f => GDistributive (Rec1 f) where
  gcollect :: forall (f :: * -> *) a b.
Functor f =>
(a -> Rec1 f b) -> f a -> Rec1 f (f b)
gcollect = coerce :: forall a b. Coercible a b => a -> b
coerce (forall (g :: * -> *) (f :: * -> *) a b.
(Distributive g, Functor f) =>
(a -> g b) -> f a -> g (f b)
collect :: (a -> f b) -> g a -> f (g b))
    :: forall g a b . Functor g
    => (a -> Rec1 f b) -> g a -> Rec1 f (g b)
  {-# INLINE gcollect #-}

instance GDistributive f => GDistributive (M1 i c f) where
  gcollect :: forall (f :: * -> *) a b.
Functor f =>
(a -> M1 i c f b) -> f a -> M1 i c f (f b)
gcollect = coerce :: forall a b. Coercible a b => a -> b
coerce (forall (g :: * -> *) (f :: * -> *) a b.
(GDistributive g, Functor f) =>
(a -> g b) -> f a -> g (f b)
gcollect :: (a -> f b) -> g a -> f (g b))
    :: forall g a b . Functor g
    => (a -> M1 i c f b) -> g a -> M1 i c f (g b)
  {-# INLINE gcollect #-}