{-# 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 Control.Monad (forM_)
import Data.Functor.Identity
import Data.Functor.Product
import Data.Primitive.SmallArray
import Data.Proxy
import Data.SOP.Classes (fn_2)
import Unsafe.Coerce (unsafeCoerce)

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 :: forall a x. Index a x -> Int
indexToInt (UnsafeIndex Int
ix) = Int
ix

getAtIndex :: Index a x -> Rep f a -> f x
getAtIndex :: forall a x (f :: * -> *). Index a x -> Rep f a -> f x
getAtIndex (UnsafeIndex Int
ix) (Rep SmallArray (f Any)
v) =
    forall a b. a -> b
unsafeCoerce forall a b. (a -> b) -> a -> b
$ forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray (f Any)
v Int
ix

putAtIndex :: Index a x -> f x -> Rep f a -> Rep f a
putAtIndex :: forall a x (f :: * -> *). Index a x -> f x -> Rep f a -> Rep f a
putAtIndex (UnsafeIndex Int
ix) f x
x (Rep SmallArray (f Any)
v) = forall (f :: * -> *) a. SmallArray (f Any) -> Rep f a
Rep forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s (f Any)
v' <- forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray SmallArray (f Any)
v Int
0 (forall a. SmallArray a -> Int
sizeofSmallArray SmallArray (f Any)
v)
    forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s (f Any)
v' Int
ix (forall a b. a -> b
noInlineUnsafeCo f x
x)
    forall (m :: * -> *) a. Monad m => a -> m a
return SmallMutableArray s (f Any)
v'

updateAtIndex ::
     Functor m
  => Index a x
  -> (f x -> m (f x))
  -> Rep f a -> m (Rep f a)
updateAtIndex :: forall (m :: * -> *) a x (f :: * -> *).
Functor m =>
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 -> 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) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f x -> m (f x)
f (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 :: forall a. Generic a => Rep (Index a) a
allIndices = forall (f :: * -> *) a. SmallArray (f Any) -> Rep f a
Rep forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s (Index a Any)
v <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
size forall a. HasCallStack => a
undefined
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
size forall a. Num a => a -> a -> a
- Int
1] forall a b. (a -> b) -> a -> b
$ \Int
i ->
      forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s (Index a Any)
v Int
i (forall a x. Int -> Index a x
UnsafeIndex Int
i)
    forall (m :: * -> *) a. Monad m => a -> m a
return SmallMutableArray s (Index a Any)
v
  where
    size :: Int
    size :: Int
size = forall a. Metadata a -> Int
recordSize (forall a (proxy :: * -> *). Generic a => proxy a -> Metadata a
metadata (forall {k} (t :: k). Proxy t
Proxy @a))

-- | 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 (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
f Rep f a
as = 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' forall a. Generic a => Rep (Index a) a
allIndices
  where
    f' :: Index a x -> g x
    f' :: forall x. Index a x -> g x
f' Index a x
ix = forall x. Index a x -> f x -> g x
f Index a x
ix (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 (f :: * -> *) a. Generic a => (forall x. f x) -> Rep f a
pure forall x. f x
f = forall (f :: * -> *) a. SmallArray (f Any) -> Rep f a
Rep forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
size forall x. f x
f
  where
    size :: Int
    size :: Int
size = forall a. Metadata a -> Int
recordSize (forall a (proxy :: * -> *). Generic a => proxy a -> Metadata a
metadata (forall {k} (t :: k). Proxy t
Proxy @a))

cpure ::
     (Generic a, Constraints a c)
  => Proxy c
  -> (forall x. c x => f x)
  -> Rep f a
cpure :: 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 forall x. c x => f x
f = forall (f :: * -> *) (g :: * -> *) a.
(forall x. f x -> g x) -> Rep f a -> Rep g a
map' (\Dict c x
Dict -> forall x. c x => f x
f) (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 :: forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
Rep (f -.-> g) a -> Rep f a -> Rep g a
ap Rep (f -.-> g) a
fs Rep f a
as = 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' :: forall x. Index a x -> (-.->) f g x -> g x
f' Index a x
ix (-.->) f g x
f = (-.->) f g x
f forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
(-.->) f g a -> f a -> g a
`apFn` 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 a (f :: * -> *) (g :: * -> *).
Generic a =>
(forall x. f x -> g x) -> Rep f a -> Rep g a
map forall x. f x -> g x
f = forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
(forall x. Index a x -> f x -> g x) -> Rep f a -> Rep g a
mapWithIndex (forall a b. a -> b -> a
const 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 (m :: * -> *) a (f :: * -> *) (g :: * -> *).
(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)
f = forall (m :: * -> *) (f :: * -> *) a.
Applicative m =>
Rep (m :.: f) a -> m (Rep f a)
sequenceA forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
(forall x. Index a x -> f x -> g x) -> Rep f a -> Rep g a
mapWithIndex (forall a b. a -> b -> a
const (forall l k (f :: l -> *) (g :: k -> l) (p :: k).
f (g p) -> (:.:) f g p
Comp forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: 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 forall x. c x => f x -> g x
f = forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
Rep (f -.-> g) a -> Rep f a -> Rep g a
ap forall a b. (a -> b) -> a -> b
$ 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 (forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
(f a -> g a) -> (-.->) f g a
Fn 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 :: forall (m :: * -> *) (f :: * -> *) (g :: * -> *)
       (c :: * -> Constraint) 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
p forall x. c x => f x -> m (g x)
f = forall (m :: * -> *) (f :: * -> *) a.
Applicative m =>
Rep (m :.: f) a -> m (Rep f a)
sequenceA forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 (forall l k (f :: l -> *) (g :: k -> l) (p :: k).
f (g p) -> (:.:) f g p
Comp forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 (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)
f Rep f a
a Rep g a
b = forall (m :: * -> *) (f :: * -> *) a.
Applicative m =>
Rep (m :.: f) a -> m (Rep f a)
sequenceA forall a b. (a -> b) -> a -> b
$
    forall (f :: * -> *) a. Generic a => (forall x. f x) -> Rep f a
pure (forall {k} (f :: k -> *) (a :: k) (f' :: k -> *) (f'' :: k -> *).
(f a -> f' a -> f'' a) -> (-.->) f (f' -.-> f'') a
fn_2 forall a b. (a -> b) -> a -> b
$ \f x
x g x
y -> forall l k (f :: l -> *) (g :: k -> l) (p :: k).
f (g p) -> (:.:) f g p
Comp forall a b. (a -> b) -> a -> b
$ forall x. f x -> g x -> m (h x)
f f x
x g x
y) forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
Rep (f -.-> g) a -> Rep f a -> Rep g a
`ap` Rep f a
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 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 -> h x
f Rep f a
a Rep g a
b = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$
    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 -> forall a. a -> Identity a
Identity forall a b. (a -> b) -> a -> b
$ 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 :: forall a (f :: * -> *) (g :: * -> *).
Generic a =>
Rep f a -> Rep g a -> Rep (Product f g) a
zip = 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 {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 :: 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 forall x. c x => f x -> g x -> m (h x)
f Rep f a
a Rep g a
b = forall (m :: * -> *) (f :: * -> *) a.
Applicative m =>
Rep (m :.: f) a -> m (Rep f a)
sequenceA forall a b. (a -> b) -> a -> b
$
    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 (forall {k} (f :: k -> *) (a :: k) (f' :: k -> *) (f'' :: k -> *).
(f a -> f' a -> f'' a) -> (-.->) f (f' -.-> f'') a
fn_2 forall a b. (a -> b) -> a -> b
$ \f x
x g x
y -> forall l k (f :: l -> *) (g :: k -> l) (p :: k).
f (g p) -> (:.:) f g p
Comp forall a b. (a -> b) -> a -> b
$ forall x. c x => f x -> g x -> m (h x)
f f x
x g x
y) forall (f :: * -> *) (g :: * -> *) a.
Generic a =>
Rep (f -.-> g) a -> Rep f a -> Rep g a
`ap` Rep f a
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 :: forall a (c :: * -> Constraint) (f :: * -> *) (g :: * -> *)
       (h :: * -> *).
(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
p forall x. c x => f x -> g x -> h x
f Rep f a
a Rep g a
b = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$
    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 -> forall a. a -> Identity a
Identity (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