{-# 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 :: F f -> f a
unF F f
x = case F f
x of F f a
x -> f a -> f a
forall a b. a -> b
unsafeCoerce f a
x

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

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

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

insert :: Typeable a => f a -> Map f -> Map f
insert :: f a -> Map f -> Map f
insert f a
x (Map HashMap TypeRep (F f)
mp) = HashMap TypeRep (F f) -> Map f
forall (f :: * -> *). HashMap TypeRep (F f) -> Map f
Map (HashMap TypeRep (F f) -> Map f) -> HashMap TypeRep (F f) -> Map f
forall a b. (a -> b) -> a -> b
$ TypeRep -> F f -> HashMap TypeRep (F f) -> HashMap TypeRep (F f)
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
Map.insert (f a -> TypeRep
forall k (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep f a
x) (f a -> F f
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 :: Map f -> Maybe (f a)
lookup (Map HashMap TypeRep (F f)
mp) = F f -> f a
forall (f :: * -> *) a. F f -> f a
unF (F f -> f a) -> Maybe (F f) -> Maybe (f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeRep -> HashMap TypeRep (F f) -> Maybe (F f)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.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)) HashMap TypeRep (F f)
mp

unionWith :: (forall a . f a -> f a -> f a) -> Map f -> Map f -> Map f
unionWith :: (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) = HashMap TypeRep (F f) -> Map f
forall (f :: * -> *). HashMap TypeRep (F f) -> Map f
Map (HashMap TypeRep (F f) -> Map f) -> HashMap TypeRep (F f) -> Map f
forall a b. (a -> b) -> a -> b
$ (F f -> F f -> F f)
-> HashMap TypeRep (F f)
-> HashMap TypeRep (F f)
-> HashMap TypeRep (F f)
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 -> f Any -> F f
forall (f :: * -> *) a. f a -> F f
F (f Any -> F f) -> f Any -> F f
forall a b. (a -> b) -> a -> b
$ f Any -> f Any -> f Any
forall a. f a -> f a -> f a
f (F f -> f Any
forall (f :: * -> *) a. F f -> f a
unF F f
x1) (F f -> f Any
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 a. f1 a -> f2 a) -> Map f1 -> Map f2
map forall a. f1 a -> f2 a
f (Map HashMap TypeRep (F f1)
mp) = HashMap TypeRep (F f2) -> Map f2
forall (f :: * -> *). HashMap TypeRep (F f) -> Map f
Map (HashMap TypeRep (F f2) -> Map f2)
-> HashMap TypeRep (F f2) -> Map f2
forall a b. (a -> b) -> a -> b
$ (F f1 -> F f2) -> HashMap TypeRep (F f1) -> HashMap TypeRep (F f2)
forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
Map.map (\(F f1 a
a) -> f2 a -> F f2
forall (f :: * -> *) a. f a -> F f
F (f2 a -> F f2) -> f2 a -> F f2
forall a b. (a -> b) -> a -> b
$ f1 a -> f2 a
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 a. f a -> b) -> Map f -> [b]
toList forall a. f a -> b
f (Map HashMap TypeRep (F f)
mp) = (F f -> b) -> [F f] -> [b]
forall a b. (a -> b) -> [a] -> [b]
Prelude.map (\(F f a
a) -> f a -> b
forall a. f a -> b
f f a
a) ([F f] -> [b]) -> [F f] -> [b]
forall a b. (a -> b) -> a -> b
$ HashMap TypeRep (F f) -> [F f]
forall k v. HashMap k v -> [v]
Map.elems HashMap TypeRep (F f)
mp

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