{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

-- | Operations on the generic representation
--
-- We also re-export some non-derive functions to clarify where they belong
-- in this list.
--
-- This module is intended for qualified import.
--
-- > import qualified Data.Record.Generic.Rep as Rep
--
-- TODO: Could we provide instances for the @generics-sop@ type classes?
-- Might lessen the pain of switching between the two or using both?
module Data.Record.Generic.Rep (
    Rep(..) -- TODO: Make opaque?
    -- * "Functor"
  , map
  , mapM
  , cmap
  , cmapM
    -- * Zipping
  , zip
  , zipWith
  , zipWithM
  , czipWith
  , czipWithM
    -- * "Foldable"
  , collapse
    -- * "Traversable"
  , sequenceA
    -- * "Applicable"
  , pure
  , cpure
  , ap
    -- * Array-like interface
  , Index -- opaque
  , indexToInt
  , getAtIndex
  , putAtIndex
  , updateAtIndex
  , allIndices
  , mapWithIndex
  ) where

import Prelude hiding (
    map
  , mapM
  , pure
  , sequenceA
  , zip
  , zipWith
  )

import Data.Proxy
import Data.Functor.Identity
import Data.Functor.Product
import Data.SOP.Classes (fn_2)
import Unsafe.Coerce (unsafeCoerce)

import qualified Data.Vector as V

import Data.Record.Generic
import Data.Record.Generic.Rep.Internal

--
-- NOTE: In order to avoid circular definitions, this module is strictly defined
-- in order: every function only depends on the functions defined before it.
--

{-------------------------------------------------------------------------------
  Array-like interface
-------------------------------------------------------------------------------}

newtype Index a x = UnsafeIndex Int

indexToInt :: Index a x -> Int
indexToInt :: Index a x -> Int
indexToInt (UnsafeIndex Int
ix) = Int
ix

getAtIndex :: Index a x -> Rep f a -> f x
getAtIndex :: Index a x -> Rep f a -> f x
getAtIndex (UnsafeIndex Int
ix) (Rep Vector (f Any)
v) =
    f Any -> f x
forall a b. a -> b
unsafeCoerce (f Any -> f x) -> f Any -> f x
forall a b. (a -> b) -> a -> b
$ Vector (f Any) -> Int -> f Any
forall a. Vector a -> Int -> a
V.unsafeIndex Vector (f Any)
v Int
ix

putAtIndex :: Index a x -> f x -> Rep f a -> Rep f a
putAtIndex :: Index a x -> f x -> Rep f a -> Rep f a
putAtIndex (UnsafeIndex Int
ix) f x
x (Rep Vector (f Any)
v) = Vector (f Any) -> Rep f a
forall (f :: * -> *) a. Vector (f Any) -> Rep f a
Rep (Vector (f Any) -> Rep f a) -> Vector (f Any) -> Rep f a
forall a b. (a -> b) -> a -> b
$
    Vector (f Any) -> [(Int, f Any)] -> Vector (f Any)
forall a. Vector a -> [(Int, a)] -> Vector a
V.unsafeUpd Vector (f Any)
v [(Int
ix, f x -> f Any
forall a b. a -> b
unsafeCoerce f x
x)]

updateAtIndex ::
     Functor m
  => Index a x
  -> (f x -> m (f x))
  -> Rep f a -> m (Rep f a)
updateAtIndex :: Index a x -> (f x -> m (f x)) -> Rep f a -> m (Rep f a)
updateAtIndex Index a x
ix f x -> m (f x)
f Rep f a
a = (\f x
x -> Index a x -> f x -> Rep f a -> Rep f a
forall a x (f :: * -> *). Index a x -> f x -> Rep f a -> Rep f a
putAtIndex Index a x
ix f x
x Rep f a
a) (f x -> Rep f a) -> m (f x) -> m (Rep f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f x -> m (f x)
f (Index a x -> Rep f a -> f x
forall a x (f :: * -> *). Index a x -> Rep f a -> f x
getAtIndex Index a x
ix Rep f a
a)

allIndices :: forall a. Generic a => Rep (Index a) a
allIndices :: Rep (Index a) a
allIndices = Vector (Index a Any) -> Rep (Index a) a
forall (f :: * -> *) a. Vector (f Any) -> Rep f a
Rep (Vector (Index a Any) -> Rep (Index a) a)
-> Vector (Index a Any) -> Rep (Index a) a
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Index a Any) -> Vector (Index a Any)
forall a. Int -> (Int -> a) -> Vector a
V.generate (Metadata a -> Int
forall a. Metadata a -> Int
recordSize (Proxy a -> Metadata a
forall a (proxy :: * -> *). Generic a => proxy a -> Metadata a
metadata (Proxy a
forall k (t :: k). Proxy t
Proxy @a))) Int -> Index a Any
forall a x. Int -> Index a x
UnsafeIndex

-- | Map with index
--
-- This is an important building block in this module.
-- Crucially, @mapWithIndex f a@ is lazy in @a@, reading elements from @a@
-- only if and when @f@ demands them.
mapWithIndex ::
     forall f g a. Generic a
  => (forall x. Index a x -> f x -> g x)
  -> Rep f a -> Rep g a
mapWithIndex :: (forall x. Index a x -> f x -> g x) -> Rep f a -> Rep g a
mapWithIndex forall x. Index a x -> f x -> g x
f Rep f a
as = (forall x. Index a x -> g x) -> Rep (Index a) a -> Rep g a
forall (f :: * -> *) (g :: * -> *) a.
(forall x. f x -> g x) -> Rep f a -> Rep g a
map' forall x. Index a x -> g x
f' Rep (Index a) a
forall a. Generic a => Rep (Index a) a
allIndices
  where
    f' :: Index a x -> g x
    f' :: Index a x -> g x
f' Index a x
ix = Index a x -> f x -> g x
forall x. Index a x -> f x -> g x
f Index a x
ix (Index a x -> Rep f a -> f x
forall a x (f :: * -> *). Index a x -> Rep f a -> f x
getAtIndex Index a x
ix Rep f a
as)

{-------------------------------------------------------------------------------
  "Applicative"
-------------------------------------------------------------------------------}

pure :: forall f a. Generic a => (forall x. f x) -> Rep f a
pure :: (forall x. f x) -> Rep f a
pure forall x. f x
f = Vector (f Any) -> Rep f a
forall (f :: * -> *) a. Vector (f Any) -> Rep f a
Rep (Int -> f Any -> Vector (f Any)
forall a. Int -> a -> Vector a
V.replicate (Metadata a -> Int
forall a. Metadata a -> Int
recordSize (Proxy a -> Metadata a
forall a (proxy :: * -> *). Generic a => proxy a -> Metadata a
metadata (Proxy a
forall k (t :: k). Proxy t
Proxy @a))) f Any
forall x. f x
f)

cpure ::
     (Generic a, Constraints a c)
  => Proxy c
  -> (forall x. c x => f x)
  -> Rep f a
cpure :: Proxy c -> (forall x. c x => f x) -> Rep f a
cpure Proxy c
p forall x. c x => f x
f = (forall x. Dict c x -> f x) -> Rep (Dict c) a -> Rep f a
forall (f :: * -> *) (g :: * -> *) a.
(forall x. f x -> g x) -> Rep f a -> Rep g a
map' (\Dict c x
Dict -> f x
forall x. c x => f x
f) (Proxy c -> Rep (Dict c) a
forall a (c :: * -> Constraint).
(Generic a, Constraints a c) =>
Proxy c -> Rep (Dict c) a
dict Proxy c
p)

-- | Higher-order version of @<*>@
--
-- Lazy in the second argument.
ap :: forall f g a. Generic a => Rep (f -.-> g) a -> Rep f a -> Rep g a
ap :: Rep (f -.-> g) a -> Rep f a -> Rep g a
ap Rep (f -.-> g) a
fs Rep f a
as = (forall x. Index a x -> (-.->) f g x -> g x)
-> Rep (f -.-> g) a -> Rep g a
forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
(forall x. Index a x -> f x -> g x) -> Rep f a -> Rep g a
mapWithIndex forall x. Index a x -> (-.->) f g x -> g x
f' Rep (f -.-> g) a
fs
  where
    f' :: Index a x -> (-.->) f g x -> g x
    f' :: Index a x -> (-.->) f g x -> g x
f' Index a x
ix (-.->) f g x
f = (-.->) f g x
f (-.->) f g x -> f x -> g x
forall k (f :: k -> *) (g :: k -> *) (a :: k).
(-.->) f g a -> f a -> g a
`apFn` Index a x -> Rep f a -> f x
forall a x (f :: * -> *). Index a x -> Rep f a -> f x
getAtIndex Index a x
ix Rep f a
as

{-------------------------------------------------------------------------------
  "Functor"
-------------------------------------------------------------------------------}

map :: Generic a => (forall x. f x -> g x) -> Rep f a -> Rep g a
map :: (forall x. f x -> g x) -> Rep f a -> Rep g a
map forall x. f x -> g x
f = (forall x. Index a x -> f x -> g x) -> Rep f a -> Rep g a
forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
(forall x. Index a x -> f x -> g x) -> Rep f a -> Rep g a
mapWithIndex ((f x -> g x) -> Index a x -> f x -> g x
forall a b. a -> b -> a
const f x -> g x
forall x. f x -> g x
f)

mapM ::
     (Applicative m, Generic a)
  => (forall x. f x -> m (g x))
  -> Rep f a -> m (Rep g a)
mapM :: (forall x. f x -> m (g x)) -> Rep f a -> m (Rep g a)
mapM forall x. f x -> m (g x)
f = Rep (m :.: g) a -> m (Rep g a)
forall (m :: * -> *) (f :: * -> *) a.
Applicative m =>
Rep (m :.: f) a -> m (Rep f a)
sequenceA (Rep (m :.: g) a -> m (Rep g a))
-> (Rep f a -> Rep (m :.: g) a) -> Rep f a -> m (Rep g a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall x. Index a x -> f x -> (:.:) m g x)
-> Rep f a -> Rep (m :.: g) a
forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
(forall x. Index a x -> f x -> g x) -> Rep f a -> Rep g a
mapWithIndex ((f x -> (:.:) m g x) -> Index a x -> f x -> (:.:) m g x
forall a b. a -> b -> a
const (m (g x) -> (:.:) m g x
forall l k (f :: l -> *) (g :: k -> l) (p :: k).
f (g p) -> (:.:) f g p
Comp (m (g x) -> (:.:) m g x) -> (f x -> m (g x)) -> f x -> (:.:) m g x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f x -> m (g x)
forall x. f x -> m (g x)
f))

cmap ::
     (Generic a, Constraints a c)
  => Proxy c
  -> (forall x. c x => f x -> g x)
  -> Rep f a -> Rep g a
cmap :: Proxy c -> (forall x. c x => f x -> g x) -> Rep f a -> Rep g a
cmap Proxy c
p forall x. c x => f x -> g x
f = Rep (f -.-> g) a -> Rep f a -> Rep g a
forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
Rep (f -.-> g) a -> Rep f a -> Rep g a
ap (Rep (f -.-> g) a -> Rep f a -> Rep g a)
-> Rep (f -.-> g) a -> Rep f a -> Rep g a
forall a b. (a -> b) -> a -> b
$ Proxy c -> (forall x. c x => (-.->) f g x) -> Rep (f -.-> g) a
forall a (c :: * -> Constraint) (f :: * -> *).
(Generic a, Constraints a c) =>
Proxy c -> (forall x. c x => f x) -> Rep f a
cpure Proxy c
p ((f x -> g x) -> (-.->) f g x
forall k (f :: k -> *) (g :: k -> *) (a :: k).
(f a -> g a) -> (-.->) f g a
Fn f x -> g x
forall x. c x => f x -> g x
f)

cmapM ::
     forall m f g c a. (Generic a, Applicative m, Constraints a c)
  => Proxy c
  -> (forall x. c x => f x -> m (g x))
  -> Rep f a -> m (Rep g a)
cmapM :: Proxy c
-> (forall x. c x => f x -> m (g x)) -> Rep f a -> m (Rep g a)
cmapM Proxy c
p forall x. c x => f x -> m (g x)
f = Rep (m :.: g) a -> m (Rep g a)
forall (m :: * -> *) (f :: * -> *) a.
Applicative m =>
Rep (m :.: f) a -> m (Rep f a)
sequenceA (Rep (m :.: g) a -> m (Rep g a))
-> (Rep f a -> Rep (m :.: g) a) -> Rep f a -> m (Rep g a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy c
-> (forall x. c x => f x -> (:.:) m g x)
-> Rep f a
-> Rep (m :.: g) a
forall a (c :: * -> Constraint) (f :: * -> *) (g :: * -> *).
(Generic a, Constraints a c) =>
Proxy c -> (forall x. c x => f x -> g x) -> Rep f a -> Rep g a
cmap Proxy c
p (m (g x) -> (:.:) m g x
forall l k (f :: l -> *) (g :: k -> l) (p :: k).
f (g p) -> (:.:) f g p
Comp (m (g x) -> (:.:) m g x) -> (f x -> m (g x)) -> f x -> (:.:) m g x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f x -> m (g x)
forall x. c x => f x -> m (g x)
f)

{-------------------------------------------------------------------------------
  Zipping
-------------------------------------------------------------------------------}

zipWithM ::
     forall m f g h a. (Generic a, Applicative m)
  => (forall x. f x -> g x -> m (h x))
  -> Rep f a -> Rep g a -> m (Rep h a)
zipWithM :: (forall x. f x -> g x -> m (h x))
-> Rep f a -> Rep g a -> m (Rep h a)
zipWithM forall x. f x -> g x -> m (h x)
f Rep f a
a Rep g a
b = Rep (m :.: h) a -> m (Rep h a)
forall (m :: * -> *) (f :: * -> *) a.
Applicative m =>
Rep (m :.: f) a -> m (Rep f a)
sequenceA (Rep (m :.: h) a -> m (Rep h a)) -> Rep (m :.: h) a -> m (Rep h a)
forall a b. (a -> b) -> a -> b
$
    (forall x. (-.->) f (g -.-> (m :.: h)) x)
-> Rep (f -.-> (g -.-> (m :.: h))) a
forall (f :: * -> *) a. Generic a => (forall x. f x) -> Rep f a
pure ((f x -> g x -> (:.:) m h x) -> (-.->) f (g -.-> (m :.: h)) x
forall k (f :: k -> *) (a :: k) (f' :: k -> *) (f'' :: k -> *).
(f a -> f' a -> f'' a) -> (-.->) f (f' -.-> f'') a
fn_2 ((f x -> g x -> (:.:) m h x) -> (-.->) f (g -.-> (m :.: h)) x)
-> (f x -> g x -> (:.:) m h x) -> (-.->) f (g -.-> (m :.: h)) x
forall a b. (a -> b) -> a -> b
$ \f x
x g x
y -> m (h x) -> (:.:) m h x
forall l k (f :: l -> *) (g :: k -> l) (p :: k).
f (g p) -> (:.:) f g p
Comp (m (h x) -> (:.:) m h x) -> m (h x) -> (:.:) m h x
forall a b. (a -> b) -> a -> b
$ f x -> g x -> m (h x)
forall x. f x -> g x -> m (h x)
f f x
x g x
y) Rep (f -.-> (g -.-> (m :.: h))) a
-> Rep f a -> Rep (g -.-> (m :.: h)) a
forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
Rep (f -.-> g) a -> Rep f a -> Rep g a
`ap` Rep f a
a Rep (g -.-> (m :.: h)) a -> Rep g a -> Rep (m :.: h) a
forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
Rep (f -.-> g) a -> Rep f a -> Rep g a
`ap` Rep g a
b

zipWith ::
     Generic a
  => (forall x. f x -> g x -> h x)
  -> Rep f a -> Rep g a -> Rep h a
zipWith :: (forall x. f x -> g x -> h x) -> Rep f a -> Rep g a -> Rep h a
zipWith forall x. f x -> g x -> h x
f Rep f a
a Rep g a
b = Identity (Rep h a) -> Rep h a
forall a. Identity a -> a
runIdentity (Identity (Rep h a) -> Rep h a) -> Identity (Rep h a) -> Rep h a
forall a b. (a -> b) -> a -> b
$
    (forall x. f x -> g x -> Identity (h x))
-> Rep f a -> Rep g a -> Identity (Rep h a)
forall (m :: * -> *) (f :: * -> *) (g :: * -> *) (h :: * -> *) a.
(Generic a, Applicative m) =>
(forall x. f x -> g x -> m (h x))
-> Rep f a -> Rep g a -> m (Rep h a)
zipWithM (\f x
x g x
y -> h x -> Identity (h x)
forall a. a -> Identity a
Identity (h x -> Identity (h x)) -> h x -> Identity (h x)
forall a b. (a -> b) -> a -> b
$ f x -> g x -> h x
forall x. f x -> g x -> h x
f f x
x g x
y) Rep f a
a Rep g a
b

zip :: Generic a => Rep f a -> Rep g a -> Rep (Product f g) a
zip :: Rep f a -> Rep g a -> Rep (Product f g) a
zip = (forall x. f x -> g x -> Product f g x)
-> Rep f a -> Rep g a -> Rep (Product f g) a
forall a (f :: * -> *) (g :: * -> *) (h :: * -> *).
Generic a =>
(forall x. f x -> g x -> h x) -> Rep f a -> Rep g a -> Rep h a
zipWith forall x. f x -> g x -> Product f g x
forall k (f :: k -> *) (g :: k -> *) (a :: k).
f a -> g a -> Product f g a
Pair

czipWithM ::
     forall m f g h c a. (Generic a, Applicative m, Constraints a c)
  => Proxy c
  -> (forall x. c x => f x -> g x -> m (h x))
  -> Rep f a -> Rep g a -> m (Rep h a)
czipWithM :: Proxy c
-> (forall x. c x => f x -> g x -> m (h x))
-> Rep f a
-> Rep g a
-> m (Rep h a)
czipWithM Proxy c
p forall x. c x => f x -> g x -> m (h x)
f Rep f a
a Rep g a
b = Rep (m :.: h) a -> m (Rep h a)
forall (m :: * -> *) (f :: * -> *) a.
Applicative m =>
Rep (m :.: f) a -> m (Rep f a)
sequenceA (Rep (m :.: h) a -> m (Rep h a)) -> Rep (m :.: h) a -> m (Rep h a)
forall a b. (a -> b) -> a -> b
$
    Proxy c
-> (forall x. c x => (-.->) f (g -.-> (m :.: h)) x)
-> Rep (f -.-> (g -.-> (m :.: h))) a
forall a (c :: * -> Constraint) (f :: * -> *).
(Generic a, Constraints a c) =>
Proxy c -> (forall x. c x => f x) -> Rep f a
cpure Proxy c
p ((f x -> g x -> (:.:) m h x) -> (-.->) f (g -.-> (m :.: h)) x
forall k (f :: k -> *) (a :: k) (f' :: k -> *) (f'' :: k -> *).
(f a -> f' a -> f'' a) -> (-.->) f (f' -.-> f'') a
fn_2 ((f x -> g x -> (:.:) m h x) -> (-.->) f (g -.-> (m :.: h)) x)
-> (f x -> g x -> (:.:) m h x) -> (-.->) f (g -.-> (m :.: h)) x
forall a b. (a -> b) -> a -> b
$ \f x
x g x
y -> m (h x) -> (:.:) m h x
forall l k (f :: l -> *) (g :: k -> l) (p :: k).
f (g p) -> (:.:) f g p
Comp (m (h x) -> (:.:) m h x) -> m (h x) -> (:.:) m h x
forall a b. (a -> b) -> a -> b
$ f x -> g x -> m (h x)
forall x. c x => f x -> g x -> m (h x)
f f x
x g x
y) Rep (f -.-> (g -.-> (m :.: h))) a
-> Rep f a -> Rep (g -.-> (m :.: h)) a
forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
Rep (f -.-> g) a -> Rep f a -> Rep g a
`ap` Rep f a
a Rep (g -.-> (m :.: h)) a -> Rep g a -> Rep (m :.: h) a
forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
Rep (f -.-> g) a -> Rep f a -> Rep g a
`ap` Rep g a
b

czipWith ::
     (Generic a, Constraints a c)
  => Proxy c
  -> (forall x. c x => f x -> g x -> h x)
  -> Rep f a -> Rep g a -> Rep h a
czipWith :: Proxy c
-> (forall x. c x => f x -> g x -> h x)
-> Rep f a
-> Rep g a
-> Rep h a
czipWith Proxy c
p forall x. c x => f x -> g x -> h x
f Rep f a
a Rep g a
b = Identity (Rep h a) -> Rep h a
forall a. Identity a -> a
runIdentity (Identity (Rep h a) -> Rep h a) -> Identity (Rep h a) -> Rep h a
forall a b. (a -> b) -> a -> b
$
    Proxy c
-> (forall x. c x => f x -> g x -> Identity (h x))
-> Rep f a
-> Rep g a
-> Identity (Rep h a)
forall (m :: * -> *) (f :: * -> *) (g :: * -> *) (h :: * -> *)
       (c :: * -> Constraint) a.
(Generic a, Applicative m, Constraints a c) =>
Proxy c
-> (forall x. c x => f x -> g x -> m (h x))
-> Rep f a
-> Rep g a
-> m (Rep h a)
czipWithM Proxy c
p (\f x
x g x
y -> h x -> Identity (h x)
forall a. a -> Identity a
Identity (f x -> g x -> h x
forall x. c x => f x -> g x -> h x
f f x
x g x
y)) Rep f a
a Rep g a
b