content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE Trustworthy #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE MultiParamTypeClasses, UndecidableInstances #-}
  -----------------------------------------------------------------------------
  -- |
  -- Module      :  Data.Extensible.Product
  -- Copyright   :  (c) Fumiaki Kinoshita 2018
  -- License     :  BSD3
  --
  -- Maintainer  :  Fumiaki Kinoshita <fumiexcel@gmail.com>
  --
  ------------------------------------------------------------------------
  module Data.Extensible.Product (
    -- * Basic operations
    (:&)
    , nil
    , (<:)
    , (<!)
    , (=<:)
    , hlength
    , type (++)
    , happend
    , hmap
    , hmapWithIndex
    , hzipWith
    , hzipWith3
    , hfoldMap
    , hfoldMapWithIndex
    , hfoldrWithIndex
    , hfoldlWithIndex
    , htraverse
    , htraverseWithIndex
    , hsequence
    -- * Constrained fold
    , hmapWithIndexFor
    , hfoldMapFor
    , hfoldMapWithIndexFor
    , hfoldrWithIndexFor
    , hfoldlWithIndexFor
    -- * Constraind fold without proxies
    , hfoldMapWith
    , hfoldMapWithIndexWith
    , hfoldrWithIndexWith
    , hfoldlWithIndexWith
    , hmapWithIndexWith
    -- * Evaluating
    , hforce
    -- * Update
    , haccumMap
    , haccum
    , hpartition
    -- * Lookup
    , hlookup
    , hindex
    -- * Generation
    , Generate(..)
    , hgenerate
    , htabulate
    , hrepeat
    , hcollect
    , hdistribute
    , fromHList
    , toHList
    , Forall(..)
    , hgenerateFor
    , htabulateFor
    , hrepeatFor
    , hgenerateWith
    , htabulateWith
    , hrepeatWith) where
  
  import Data.Extensible.Internal.Rig (review)
  import Data.Extensible.Struct
  import Data.Extensible.Sum
  import Data.Extensible.Class
  import Data.Extensible.Wrapper
  import Data.Functor.Compose
  import Data.Proxy
  import qualified Type.Membership.HList as HList
  
  -- | O(n) Prepend an element onto a product.
  -- Expressions like @a <: b <: c <: nil@ are transformed to a single 'fromHList'.
  (<:) :: h x -> xs :& h -> (x ': xs) :& h
  (<:) x = fromHList . HList.HCons x . toHList
  {-# INLINE (<:) #-}
  infixr 0 <:
  
  (=<:) :: Wrapper h => Repr h x -> xs :& h -> (x ': xs) :& h
  (=<:) = (<:) . review _Wrapper
  {-# INLINE (=<:) #-}
  infixr 0 =<:
  
  -- | Strict version of ('<:').
  (<!) :: h x -> xs :& h -> (x ': xs) :& h
  (<!) x = fromHList . (HList.HCons $! x) . toHList
  {-# INLINE (<!) #-}
  infixr 0 <!
  
  -- | An empty product.
  nil :: '[] :& h
  nil = hfrozen $ new $ error "Impossible"
  {-# NOINLINE nil #-}
  {-# RULES "toHList/nil" toHList nil = HList.HNil #-}
  
  -- | Convert 'HList.HList' into a product.
  fromHList :: HList.HList h xs -> xs :& h
  fromHList xs = hfrozen (newFromHList xs)
  {-# INLINE fromHList #-}
  
  -- | Flipped 'hlookup'
  hindex :: xs :& h -> Membership xs x ->  h x
  hindex = flip hlookup
  {-# INLINE hindex #-}
  
  -- | Map a function to every element of a product.
  hmapWithIndex :: (forall x. Membership xs x -> g x -> h x) -> xs :& g -> xs :& h
  hmapWithIndex t p = hfrozen (newFrom p t)
  {-# INLINE hmapWithIndex #-}
  
  -- | Map a function to every element of a product.
  hmapWithIndexFor :: Forall c xs
    => proxy c
    -> (forall x. c x => Membership xs x -> g x -> h x)
    -> xs :& g -> xs :& h
  hmapWithIndexFor c t p = hfrozen $ newFor c $ \i -> t i $ hlookup i p
  {-# INLINE hmapWithIndexFor #-}
  
  hmapWithIndexWith :: forall c xs g h. Forall c xs
    => (forall x. c x => Membership xs x -> g x -> h x)
    -> xs :& g -> xs :& h
  hmapWithIndexWith = hmapWithIndexFor (Proxy @c)
  
  -- | Transform every element in a product, preserving the order.
  --
  -- @
  -- 'hmap' 'id' ≡ 'id'
  -- 'hmap' (f . g) ≡ 'hmap' f . 'hmap' g
  -- @
  hmap :: (forall x. g x -> h x) -> xs :& g -> xs :& h
  hmap f = hmapWithIndex (const f)
  {-# INLINE hmap #-}
  
  -- | 'zipWith' for heterogeneous product
  hzipWith :: (forall x. f x -> g x -> h x) -> xs :& f -> xs :& g -> xs :& h
  hzipWith t xs = hmapWithIndex (\i -> t (hlookup i xs))
  {-# INLINE hzipWith #-}
  
  -- | 'zipWith3' for heterogeneous product
  hzipWith3 :: (forall x. f x -> g x -> h x -> i x) -> xs :& f -> xs :& g -> xs :& h -> xs :& i
  hzipWith3 t xs ys = hmapWithIndex (\i -> t (hlookup i xs) (hlookup i ys))
  {-# INLINE hzipWith3 #-}
  
  -- | Map elements to a monoid and combine the results.
  --
  -- @'hfoldMap' f . 'hmap' g ≡ 'hfoldMap' (f . g)@
  hfoldMap :: Monoid a => (forall x. h x -> a) -> xs :& h -> a
  hfoldMap f = hfoldMapWithIndex (const f)
  {-# INLINE hfoldMap #-}
  
  -- | 'hfoldMap' with the membership of elements.
  hfoldMapWithIndex :: Monoid a
    => (forall x. Membership xs x -> g x -> a) -> xs :& g -> a
  hfoldMapWithIndex f = hfoldrWithIndex (\i -> mappend . f i) mempty
  {-# INLINE hfoldMapWithIndex #-}
  
  -- | Perform a strict left fold over the elements.
  hfoldlWithIndex :: (forall x. Membership xs x -> r -> h x -> r) -> r -> xs :& h -> r
  hfoldlWithIndex f r xs = hfoldrWithIndex (\i x c a -> c $! f i a x) id xs r
  {-# INLINE hfoldlWithIndex #-}
  
  -- | 'hfoldrWithIndex' with a constraint for each element.
  hfoldrWithIndexFor :: forall c xs h r proxy. (Forall c xs) => proxy c
    -> (forall x. c x => Membership xs x -> h x -> r -> r) -> r -> xs :& h -> r
  hfoldrWithIndexFor p f r xs = henumerateFor p (Proxy :: Proxy xs) (\i -> f i (hlookup i xs)) r
  {-# INLINE hfoldrWithIndexFor #-}
  
  hfoldrWithIndexWith :: forall c xs h r. (Forall c xs)
    => (forall x. c x => Membership xs x -> h x -> r -> r) -> r -> xs :& h -> r
  hfoldrWithIndexWith f r xs = henumerateFor (Proxy @c) (Proxy @xs) (\i -> f i (hlookup i xs)) r
  {-# INLINE hfoldrWithIndexWith #-}
  
  -- | Constrained 'hfoldlWithIndex'
  hfoldlWithIndexFor :: (Forall c xs) => proxy c
    -> (forall x. c x => Membership xs x -> r -> h x -> r) -> r -> xs :& h -> r
  hfoldlWithIndexFor p f r xs = hfoldrWithIndexFor p (\i x c a -> c $! f i a x) id xs r
  {-# INLINE hfoldlWithIndexFor #-}
  
  -- | Constrained 'hfoldlWithIndex'
  hfoldlWithIndexWith :: forall c xs h r. (Forall c xs)
    => (forall x. c x => Membership xs x -> r -> h x -> r) -> r -> xs :& h -> r
  hfoldlWithIndexWith f r xs = hfoldrWithIndexWith @c (\i x c a -> c $! f i a x) id xs r
  {-# INLINE hfoldlWithIndexWith #-}
  
  -- | 'hfoldMapWithIndex' with a constraint for each element.
  hfoldMapWithIndexFor :: (Forall c xs, Monoid a) => proxy c
    -> (forall x. c x => Membership xs x -> h x -> a) -> xs :& h -> a
  hfoldMapWithIndexFor p f = hfoldrWithIndexFor p (\i -> mappend . f i) mempty
  {-# INLINE hfoldMapWithIndexFor #-}
  
  -- | 'hfoldMapWithIndex' with a constraint for each element.
  hfoldMapWithIndexWith :: forall c xs h a. (Forall c xs, Monoid a)
    => (forall x. c x => Membership xs x -> h x -> a) -> xs :& h -> a
  hfoldMapWithIndexWith f = hfoldrWithIndexWith @c (\i -> mappend . f i) mempty
  {-# INLINE hfoldMapWithIndexWith #-}
  
  -- | Constrained 'hfoldMap'
  hfoldMapFor :: (Forall c xs, Monoid a) => proxy c
    -> (forall x. c x => h x -> a) -> xs :& h -> a
  hfoldMapFor p f = hfoldMapWithIndexFor p (const f)
  {-# INLINE hfoldMapFor #-}
  
  -- | Constrained 'hfoldMap'
  hfoldMapWith :: forall c xs h a. (Forall c xs, Monoid a)
    => (forall x. c x => h x -> a) -> xs :& h -> a
  hfoldMapWith f = hfoldMapWithIndexFor (Proxy @c) (const f)
  {-# INLINE hfoldMapWith #-}
  
  -- | Traverse all elements and combine the result sequentially.
  -- @
  -- htraverse (fmap f . g) ≡ fmap (hmap f) . htraverse g
  -- htraverse pure ≡ pure
  -- htraverse (Compose . fmap g . f) ≡ Compose . fmap (htraverse g) . htraverse f
  -- @
  htraverse :: Applicative f => (forall x. g x -> f (h x)) -> xs :& g -> f (xs :& h)
  htraverse f = fmap fromHList . HList.htraverse f . toHList
  {-# INLINE htraverse #-}
  
  -- | 'sequence' analog for extensible products
  hsequence :: Applicative f => xs :& Compose f h -> f (xs :& h)
  hsequence = htraverse getCompose
  {-# INLINE hsequence #-}
  
  -- | The dual of 'htraverse'
  hcollect :: (Functor f, Generate xs) => (a -> xs :& h) -> f a -> xs :& Compose f h
  hcollect f m = htabulate $ \i -> Compose $ fmap (hlookup i . f) m
  {-# INLINABLE hcollect #-}
  
  -- | The dual of 'hsequence'
  hdistribute :: (Functor f, Generate xs) => f (xs :& h) -> xs :& Compose f h
  hdistribute = hcollect id
  {-# INLINE hdistribute #-}
  
  -- | 'htraverse' with 'Membership's.
  htraverseWithIndex :: Applicative f
    => (forall x. Membership xs x -> g x -> f (h x)) -> xs :& g -> f (xs :& h)
  htraverseWithIndex f = fmap fromHList . HList.htraverseWithIndex f . toHList
  {-# INLINE htraverseWithIndex #-}
  
  -- | A product filled with the specified value.
  hrepeat :: Generate xs => (forall x. h x) -> xs :& h
  hrepeat x = hfrozen $ newRepeat x
  {-# INLINE hrepeat #-}
  
  -- | Construct a product using a function which takes a 'Membership'.
  --
  -- @
  -- 'hmap' f ('htabulate' g) ≡ 'htabulate' (f . g)
  -- 'htabulate' ('hindex' m) ≡ m
  -- 'hindex' ('htabulate' k) ≡ k
  -- @
  htabulate :: Generate xs => (forall x. Membership xs x -> h x) -> xs :& h
  htabulate f = hfrozen $ new f
  {-# INLINE htabulate #-}
  
  -- | 'Applicative' version of 'htabulate'.
  hgenerate :: (Generate xs, Applicative f)
    => (forall x. Membership xs x -> f (h x)) -> f (xs :& h)
  hgenerate f = fmap fromHList $ hgenerateList f
  {-# INLINE hgenerate #-}
  
  -- | Pure version of 'hgenerateFor'.
  htabulateFor :: Forall c xs => proxy c -> (forall x. c x => Membership xs x -> h x) -> xs :& h
  htabulateFor p f = hfrozen $ newFor p f
  {-# INLINE htabulateFor #-}
  
  -- | Pure version of 'hgenerateFor'.
  htabulateWith :: forall c xs h. Forall c xs => (forall x. c x => Membership xs x -> h x) -> xs :& h
  htabulateWith f = hfrozen $ newFor (Proxy @c) f
  {-# INLINE htabulateWith #-}
  
  -- | A product filled with the specified value.
  hrepeatFor :: Forall c xs => proxy c -> (forall x. c x => h x) -> xs :& h
  hrepeatFor p f = htabulateFor p (const f)
  {-# INLINE hrepeatFor #-}
  
  -- | A product filled with the specified value.
  hrepeatWith :: forall c xs h. Forall c xs => (forall x. c x => h x) -> xs :& h
  hrepeatWith f = htabulateFor (Proxy @c) (const f)
  {-# INLINE hrepeatWith #-}
  
  -- | 'Applicative' version of 'htabulateFor'.
  hgenerateFor :: (Forall c xs, Applicative f)
    => proxy c -> (forall x. c x => Membership xs x -> f (h x)) -> f (xs :& h)
  hgenerateFor p f = fmap fromHList $ hgenerateListFor p f
  {-# INLINE hgenerateFor #-}
  
  -- | 'Applicative' version of 'htabulateFor'.
  hgenerateWith :: forall c xs f h. (Forall c xs, Applicative f)
    => (forall x. c x => Membership xs x -> f (h x)) -> f (xs :& h)
  hgenerateWith f = fmap fromHList $ hgenerateListFor (Proxy @c) f
  {-# INLINE hgenerateWith #-}
  
  -- | Accumulate sums on a product.
  haccumMap :: Foldable f
    => (a -> xs :/ g)
    -> (forall x. Membership xs x -> g x -> h x -> h x)
    -> xs :& h -> f a -> xs :& h
  haccumMap f g p0 xs = hmodify
    (\s -> mapM_ (\x -> case f x of EmbedAt i v -> get s i >>= set s i . g i v) xs)
    p0
  {-# INLINE haccumMap #-}
  
  -- | @haccum = 'haccumMap' 'id'@
  haccum :: Foldable f
    => (forall x. Membership xs x -> g x -> h x -> h x)
    -> xs :& h -> f (xs :/ g) -> xs :& h
  haccum = haccumMap id
  {-# INLINE haccum #-}
  
  -- | Group sums by type.
  hpartition :: (Foldable f, Generate xs) => (a -> xs :/ h) -> f a -> xs :& Compose [] h
  hpartition f = haccumMap f (\_ x (Compose xs) -> Compose (x:xs)) $ hrepeat $ Compose []
  {-# INLINE hpartition #-}
  
  -- | Evaluate every element in a product.
  hforce :: xs :& h -> xs :& h
  hforce p = hfoldrWithIndex (const seq) p p
  {-# INLINE hforce #-}
  