{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE MagicHash              #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeOperators          #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Comp.Dict
-- Copyright   :  (c) 2020 James Koppel
-- License     :  BSD3
--
-- The goal of this package is to define the `All` typeclass, where, if @fs :: [k]@ is a type-level
-- list of types or type constructors, then @All c fs@ holds if @c f@ holds for each @f@ in @fs@.
--
--------------------------------------------------------------------------------

module Data.Comp.Dict
       ( Dict (..)
       , All
       , withDict
       , (\\)
       , dictFor
       ) where

import GHC.Exts
import Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Unsafe.Coerce as U

import Data.Comp.Elem

-- | Existential combinator. See `dictFor` function for a usage.
--   Like the definition of @E@ in @Data.Comp.Multi.HFunctor@, but PolyKinded.
data E (f :: k1 -> *) where
  E :: f e -> E f

-- | A reified `Constraint`. Similar in spirit to @Dict@ from ekmett's @constraints@ package,
--   but that definition has a kind of @Constraint -> */
data Dict (c :: k -> Constraint) (a :: k) where
  Dict :: c a => Dict c a

-- | Puts the typeclass instance of a `Dict` back into the context.
withDict :: Dict c a -> (c a => r) -> r
withDict :: Dict c a -> (c a => r) -> r
withDict Dict x :: c a => r
x = r
c a => r
x

infixl 1 \\

-- | Inline version of `withDict`
(\\) :: (c a => r) -> Dict c a -> r
\\ :: (c a => r) -> Dict c a -> r
(\\) x :: c a => r
x Dict = r
c a => r
x

-- |
-- An instance of @All c fs@ holds if @c f@ holds for all @f@ in @fs@.
--
-- Example: @All `HFunctor` '[Add, Mul]@ holds if there are `HFunctor` instances for signatures
-- @Add@ and @Mul@.
--
-- The primary way to consume an `All` instance is with the `caseCxt` function. E.g.:
--
-- @
-- class Pretty f where
--   pretty :: f e l -> String
--
-- instance (All Pretty fs) => Pretty (Sum fs) where
--   pretty x = caseCxt (Proxy @Pretty) pretty x
-- @
class All (c :: k -> Constraint) (fs :: [k]) where
  -- | Primitive which returns list of dictionaries for each @f@ in @fs@
  dicts :: Proxy# fs -> [E (Dict c)]

instance All c '[] where
  {-# INLINE dicts #-}
  dicts :: Proxy# '[] -> [E (Dict c)]
dicts _ = []

instance (All c fs, c f) => All c (f ': fs) where
  {-# INLINE dicts #-}
  dicts :: Proxy# (f : fs) -> [E (Dict c)]
dicts p :: Proxy# (f : fs)
p = Dict c f -> E (Dict c)
forall k1 (f :: k1 -> *) (e :: k1). f e -> E f
E (Dict c f
forall k (c :: k -> Constraint) (a :: k). c a => Dict c a
Dict :: Dict c f) E (Dict c) -> [E (Dict c)] -> [E (Dict c)]
forall a. a -> [a] -> [a]
: (Proxy# fs -> [E (Dict c)]
forall k (c :: k -> Constraint) (fs :: [k]).
All c fs =>
Proxy# fs -> [E (Dict c)]
dicts (Proxy# (f : fs) -> Proxy# fs
forall k k (b :: k) (a :: k). Proxy# a -> Proxy# b
reproxy @fs Proxy# (f : fs)
p))

reproxy :: forall b a. Proxy# a -> Proxy# b
reproxy :: Proxy# a -> Proxy# b
reproxy _ = Proxy# b
forall k (a :: k). Proxy# a
proxy#

-- | If @All c fs@ holds, and @e@ is type-level proof that @f@ is in @fs@,
--   then @dictFor e@ is the instance of `c` for `f`
--
--   This uses `unsafeCoerce` under the hood, but is safe if @e@ is constructed by the public API.
dictFor :: forall c f fs. (All c fs) => Elem f fs -> Dict c f
dictFor :: Elem f fs -> Dict c f
dictFor (Elem v :: Int
v) =
  let ds :: Vector (E (Dict c))
ds = [E (Dict c)] -> Vector (E (Dict c))
forall a. [a] -> Vector a
V.fromList (Proxy# fs -> [E (Dict c)]
forall k (c :: k -> Constraint) (fs :: [k]).
All c fs =>
Proxy# fs -> [E (Dict c)]
dicts (Proxy# fs
forall k (a :: k). Proxy# a
proxy# :: Proxy# fs)) :: Vector (E (Dict c))
  in case Vector (E (Dict c))
ds Vector (E (Dict c)) -> Int -> E (Dict c)
forall a. Vector a -> Int -> a
V.! Int
v of
       E d :: Dict c e
d -> Dict c e -> Dict c f
forall a b. a -> b
U.unsafeCoerce Dict c e
d
{-# INLINE dictFor #-}