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

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

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


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

unF :: F f -> f a
unF x = case x of F x -> unsafeCoerce x

newtype Map (f :: * -> *) = Map (Map.HashMap TypeRep (F f))

empty :: Map f
empty = Map Map.empty

singleton :: Typeable a => f a -> Map f
singleton x = Map $ Map.singleton (typeRep x) (F x)

insert :: Typeable a => f a -> Map f -> Map f
insert x (Map mp) = Map $ Map.insert (typeRep x) (F x) mp

lookup :: forall a f . Typeable a => Map f -> Maybe (f a)
lookup (Map mp) = unF <$> Map.lookup (typeRep (Proxy :: Proxy a)) mp

unionWith :: (forall a . f a -> f a -> f a) -> Map f -> Map f -> Map f
unionWith f (Map mp1) (Map mp2) = Map $ Map.unionWith (\x1 x2 -> F $ f (unF x1) (unF x2)) mp1 mp2

map :: (forall a . f1 a -> f2 a) -> Map f1 -> Map f2
map f (Map mp) = Map $ Map.map (\(F a) -> F $ f a) mp