{-# LANGUAGE PolyKinds #-}

{- |
Copyright:  (c) 2017-2020 Kowainik
SPDX-License-Identifier: MPL-2.0
Maintainer: Kowainik <xrom.xkov@gmail.com>

TypeRepMap implementation based on @containers@ 'Map'.
-}

module Data.TypeRep.CMap
       ( TypeRepMap (..)
       , empty
       , insert
       , keys
       , lookup
       , size
       ) where

import Prelude hiding (lookup)

import Control.DeepSeq
import Data.Kind (Type)
import Data.Proxy (Proxy (..))
import Data.Typeable (TypeRep, Typeable, typeRep)
import GHC.Base (Any)
import Unsafe.Coerce (unsafeCoerce)

import qualified Data.Map.Lazy as LMap


-- | Map-like data structure with types served as the keys.
newtype TypeRepMap (f :: k -> Type) = TypeRepMap
    { forall k (f :: k -> *). TypeRepMap f -> Map TypeRep Any
unMap :: LMap.Map TypeRep Any
    }

instance NFData (TypeRepMap f) where
  rnf :: TypeRepMap f -> ()
rnf TypeRepMap f
x = [TypeRep] -> ()
forall a. NFData a => a -> ()
rnf (TypeRepMap f -> [TypeRep]
forall {k} (f :: k -> *). TypeRepMap f -> [TypeRep]
keys TypeRepMap f
x) () -> () -> ()
forall a b. a -> b -> b
`seq` ()

-- | Empty structure.
empty :: TypeRepMap f
empty :: forall {k} (f :: k -> *). TypeRepMap f
empty = Map TypeRep Any -> TypeRepMap f
forall k (f :: k -> *). Map TypeRep Any -> TypeRepMap f
TypeRepMap Map TypeRep Any
forall a. Monoid a => a
mempty

-- | Inserts the value with its type as a key.
insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
insert :: forall {k} (a :: k) (f :: k -> *).
Typeable a =>
f a -> TypeRepMap f -> TypeRepMap f
insert f a
val = Map TypeRep Any -> TypeRepMap f
forall k (f :: k -> *). Map TypeRep Any -> TypeRepMap f
TypeRepMap (Map TypeRep Any -> TypeRepMap f)
-> (TypeRepMap f -> Map TypeRep Any)
-> TypeRepMap f
-> TypeRepMap f
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRep -> Any -> Map TypeRep Any -> Map TypeRep Any
forall k a. Ord k => k -> a -> Map k a -> Map k a
LMap.insert (Proxy a -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)) (f a -> Any
forall a b. a -> b
unsafeCoerce f a
val) (Map TypeRep Any -> Map TypeRep Any)
-> (TypeRepMap f -> Map TypeRep Any)
-> TypeRepMap f
-> Map TypeRep Any
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRepMap f -> Map TypeRep Any
forall k (f :: k -> *). TypeRepMap f -> Map TypeRep Any
unMap

-- | Looks up the value at the type.
-- >>> let x = lookup $ insert (11 :: Int) empty
-- >>> x :: Maybe Int
-- Just 11
-- >>> x :: Maybe ()
-- Nothing
lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
lookup :: forall {k} (a :: k) (f :: k -> *).
Typeable a =>
TypeRepMap f -> Maybe (f a)
lookup = (Any -> f a) -> Maybe Any -> Maybe (f a)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Any -> f a
forall a b. a -> b
unsafeCoerce (Maybe Any -> Maybe (f a))
-> (TypeRepMap f -> Maybe Any) -> TypeRepMap f -> Maybe (f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRep -> Map TypeRep Any -> Maybe Any
forall k a. Ord k => k -> Map k a -> Maybe a
LMap.lookup (Proxy a -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)) (Map TypeRep Any -> Maybe Any)
-> (TypeRepMap f -> Map TypeRep Any) -> TypeRepMap f -> Maybe Any
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRepMap f -> Map TypeRep Any
forall k (f :: k -> *). TypeRepMap f -> Map TypeRep Any
unMap

size :: TypeRepMap f -> Int
size :: forall {k} (f :: k -> *). TypeRepMap f -> Int
size = Map TypeRep Any -> Int
forall k a. Map k a -> Int
LMap.size (Map TypeRep Any -> Int)
-> (TypeRepMap f -> Map TypeRep Any) -> TypeRepMap f -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRepMap f -> Map TypeRep Any
forall k (f :: k -> *). TypeRepMap f -> Map TypeRep Any
unMap

keys :: TypeRepMap f -> [TypeRep]
keys :: forall {k} (f :: k -> *). TypeRepMap f -> [TypeRep]
keys = Map TypeRep Any -> [TypeRep]
forall k a. Map k a -> [k]
LMap.keys (Map TypeRep Any -> [TypeRep])
-> (TypeRepMap f -> Map TypeRep Any) -> TypeRepMap f -> [TypeRep]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRepMap f -> Map TypeRep Any
forall k (f :: k -> *). TypeRepMap f -> Map TypeRep Any
unMap