{-# LANGUAGE ExistentialQuantification, ConstraintKinds, KindSignatures, GADTs, ScopedTypeVariables, Rank2Types #-}

module General.TypeMap(
    Map, empty, singleton, insert, map, lookup, unionWith, toList, size
    ) where

import qualified Data.HashMap.Strict as Map
import Data.Typeable
import Unsafe.Coerce
import Data.Functor
import qualified Prelude
import Prelude hiding (lookup, map)


data F f = forall a . F !(f a)

unF :: F f -> f a
unF :: forall (f :: * -> *) a. F f -> f a
unF F f
x = case F f
x of F f a
x -> forall a b. a -> b
unsafeCoerce f a
x

newtype Map f = Map (Map.HashMap TypeRep (F f))

empty :: Map f
empty :: forall (f :: * -> *). Map f
empty = forall (f :: * -> *). HashMap TypeRep (F f) -> Map f
Map forall k v. HashMap k v
Map.empty

singleton :: Typeable a => f a -> Map f
singleton :: forall a (f :: * -> *). Typeable a => f a -> Map f
singleton f a
x = forall (f :: * -> *). HashMap TypeRep (F f) -> Map f
Map forall a b. (a -> b) -> a -> b
$ forall k v. Hashable k => k -> v -> HashMap k v
Map.singleton (forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep f a
x) (forall (f :: * -> *) a. f a -> F f
F f a
x)

insert :: Typeable a => f a -> Map f -> Map f
insert :: forall a (f :: * -> *). Typeable a => f a -> Map f -> Map f
insert f a
x (Map HashMap TypeRep (F f)
mp) = forall (f :: * -> *). HashMap TypeRep (F f) -> Map f
Map forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert (forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep f a
x) (forall (f :: * -> *) a. f a -> F f
F f a
x) HashMap TypeRep (F f)
mp

lookup :: forall a f . Typeable a => Map f -> Maybe (f a)
lookup :: forall a (f :: * -> *). Typeable a => Map f -> Maybe (f a)
lookup (Map HashMap TypeRep (F f)
mp) = forall (f :: * -> *) a. F f -> f a
unF forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup (forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (forall {k} (t :: k). Proxy t
Proxy :: Proxy a)) HashMap TypeRep (F f)
mp

unionWith :: (forall a . f a -> f a -> f a) -> Map f -> Map f -> Map f
unionWith :: forall (f :: * -> *).
(forall a. f a -> f a -> f a) -> Map f -> Map f -> Map f
unionWith forall a. f a -> f a -> f a
f (Map HashMap TypeRep (F f)
mp1) (Map HashMap TypeRep (F f)
mp2) = forall (f :: * -> *). HashMap TypeRep (F f) -> Map f
Map forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
Map.unionWith (\F f
x1 F f
x2 -> forall (f :: * -> *) a. f a -> F f
F forall a b. (a -> b) -> a -> b
$ forall a. f a -> f a -> f a
f (forall (f :: * -> *) a. F f -> f a
unF F f
x1) (forall (f :: * -> *) a. F f -> f a
unF F f
x2)) HashMap TypeRep (F f)
mp1 HashMap TypeRep (F f)
mp2

map :: (forall a . f1 a -> f2 a) -> Map f1 -> Map f2
map :: forall (f1 :: * -> *) (f2 :: * -> *).
(forall a. f1 a -> f2 a) -> Map f1 -> Map f2
map forall a. f1 a -> f2 a
f (Map HashMap TypeRep (F f1)
mp) = forall (f :: * -> *). HashMap TypeRep (F f) -> Map f
Map forall a b. (a -> b) -> a -> b
$ forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
Map.map (\(F f1 a
a) -> forall (f :: * -> *) a. f a -> F f
F forall a b. (a -> b) -> a -> b
$ forall a. f1 a -> f2 a
f f1 a
a) HashMap TypeRep (F f1)
mp

toList :: (forall a . f a -> b) -> Map f -> [b]
toList :: forall (f :: * -> *) b. (forall a. f a -> b) -> Map f -> [b]
toList forall a. f a -> b
f (Map HashMap TypeRep (F f)
mp) = forall a b. (a -> b) -> [a] -> [b]
Prelude.map (\(F f a
a) -> forall a. f a -> b
f f a
a) forall a b. (a -> b) -> a -> b
$ forall k v. HashMap k v -> [v]
Map.elems HashMap TypeRep (F f)
mp

size :: Map f -> Int
size :: forall (f :: * -> *). Map f -> Int
size (Map HashMap TypeRep (F f)
mp) = forall k v. HashMap k v -> Int
Map.size HashMap TypeRep (F f)
mp