    {-# LANGUAGE BangPatterns #-}
    {-# LANGUAGE RankNTypes   #-}

    ------------------------------------------------------------------------------
    module Data.HeterogeneousEnvironment
      ( KeyGen
      , HeterogeneousEnvironment
      , Key
      , newKeyGen
      , empty
      , makeKey
      , lookup
      , insert
      , delete
      , adjust
      , getKeyId
      ) where

    ------------------------------------------------------------------------------
    import           Control.Monad
    import           Data.IntMap   (IntMap)
    import qualified Data.IntMap   as IM
    import           Data.IORef
    import           GHC.Exts
    import           Prelude hiding (lookup)
    import           Unsafe.Coerce

    ------------------------------------------------------------------------------
    data HeterogeneousEnvironment = HeterogeneousEnvironment (IntMap Any)
    newtype Key a = Key Int
    newtype KeyGen = KeyGen (IORef Int)

    ------------------------------------------------------------------------------
    -- | If you use two different KeyGens to work with the same map, you deserve
    -- what you get.
    newKeyGen :: IO KeyGen
    newKeyGen :: IO KeyGen
    newKeyGen = (IORef Int -> KeyGen) -> IO (IORef Int) -> IO KeyGen
    forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
    liftM IORef Int -> KeyGen
    KeyGen (IO (IORef Int) -> IO KeyGen) -> IO (IORef Int) -> IO KeyGen
    forall a b. (a -> b) -> a -> b
    $ Int -> IO (IORef Int)
    forall a. a -> IO (IORef a)
    newIORef Int
    0

    ------------------------------------------------------------------------------
    getKeyId :: Key a -> Int
    getKeyId :: forall a. Key a -> Int
    getKeyId (Key Int
    x) = Int
    x

    ------------------------------------------------------------------------------
    empty :: HeterogeneousEnvironment
    empty :: HeterogeneousEnvironment
    empty = IntMap Any -> HeterogeneousEnvironment
    HeterogeneousEnvironment (IntMap Any -> HeterogeneousEnvironment)
    -> IntMap Any -> HeterogeneousEnvironment
    forall a b. (a -> b) -> a -> b
    $ IntMap Any
    forall a. IntMap a
    IM.empty

    ------------------------------------------------------------------------------
    makeKey :: KeyGen -> IO (Key a)
    makeKey :: forall a. KeyGen -> IO (Key a)
    makeKey (KeyGen IORef Int
    gen) = do
        k <- IORef Int -> (Int -> (Int, Int)) -> IO Int
    forall a b. IORef a -> (a -> (a, b)) -> IO b
    atomicModifyIORef IORef Int
    gen Int -> (Int, Int)
    forall {b}. (Ord b, Bounded b, Num b) => b -> (b, b)
    nextKey
        return $ Key k
      where
        nextKey :: b -> (b, b)
    nextKey !b
    x = if b
    x b -> b -> Bool
    forall a. Ord a => a -> a -> Bool
    >= b
    forall a. Bounded a => a
    maxBoundb -> b -> b
    forall a. Num a => a -> a -> a
    -b
    1
                       then [Char] -> (b, b)
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "too many keys generated"
                       else let !x' :: b
    x' = b
    xb -> b -> b
    forall a. Num a => a -> a -> a
    +b
    1 in (b
    x',b
    x)

    ------------------------------------------------------------------------------
    lookup :: Key a -> HeterogeneousEnvironment -> Maybe a
    lookup :: forall a. Key a -> HeterogeneousEnvironment -> Maybe a
    lookup (Key Int
    k) (HeterogeneousEnvironment IntMap Any
    m) = (Any -> a) -> Maybe Any -> Maybe a
    forall a b. (a -> b) -> Maybe a -> Maybe b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap Any -> a
    forall a b. a -> b
    unsafeCoerce (Maybe Any -> Maybe a) -> Maybe Any -> Maybe a
    forall a b. (a -> b) -> a -> b
    $ Int -> IntMap Any -> Maybe Any
    forall a. Int -> IntMap a -> Maybe a
    IM.lookup Int
    k IntMap Any
    m

    ------------------------------------------------------------------------------
    insert :: Key a -> a -> HeterogeneousEnvironment -> HeterogeneousEnvironment
    insert :: forall a.
    Key a -> a -> HeterogeneousEnvironment -> HeterogeneousEnvironment
    insert (Key Int
    k) a
    v (HeterogeneousEnvironment IntMap Any
    m) = IntMap Any -> HeterogeneousEnvironment
    HeterogeneousEnvironment (IntMap Any -> HeterogeneousEnvironment)
    -> IntMap Any -> HeterogeneousEnvironment
    forall a b. (a -> b) -> a -> b
    $
                                                    Int -> Any -> IntMap Any -> IntMap Any
    forall a. Int -> a -> IntMap a -> IntMap a
    IM.insert Int
    k (a -> Any
    forall a b. a -> b
    unsafeCoerce a
    v) IntMap Any
    m

    ------------------------------------------------------------------------------
    delete :: Key a -> HeterogeneousEnvironment -> HeterogeneousEnvironment
    delete :: forall a.
    Key a -> HeterogeneousEnvironment -> HeterogeneousEnvironment
    delete (Key Int
    k) (HeterogeneousEnvironment IntMap Any
    m) = IntMap Any -> HeterogeneousEnvironment
    HeterogeneousEnvironment (IntMap Any -> HeterogeneousEnvironment)
    -> IntMap Any -> HeterogeneousEnvironment
    forall a b. (a -> b) -> a -> b
    $
                                                  Int -> IntMap Any -> IntMap Any
    forall a. Int -> IntMap a -> IntMap a
    IM.delete Int
    k IntMap Any
    m

    ------------------------------------------------------------------------------
    adjust :: (a -> a) -> Key a -> HeterogeneousEnvironment -> HeterogeneousEnvironment
    adjust :: forall a.
    (a -> a)
    -> Key a -> HeterogeneousEnvironment -> HeterogeneousEnvironment
    adjust a -> a
    f (Key Int
    k) (HeterogeneousEnvironment IntMap Any
    m) = IntMap Any -> HeterogeneousEnvironment
    HeterogeneousEnvironment (IntMap Any -> HeterogeneousEnvironment)
    -> IntMap Any -> HeterogeneousEnvironment
    forall a b. (a -> b) -> a -> b
    $
                                                    (Any -> Any) -> Int -> IntMap Any -> IntMap Any
    forall a. (a -> a) -> Int -> IntMap a -> IntMap a
    IM.adjust Any -> Any
    forall a b. a -> b
    f' Int
    k IntMap Any
    m
      where
        f' :: a -> c
    f' = a -> c
    forall a b. a -> b
    unsafeCoerce (a -> c) -> (a -> a) -> a -> c
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . a -> a
    f (a -> a) -> (a -> a) -> a -> a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . a -> a
    forall a b. a -> b
    unsafeCoerce