{-# 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 = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM IORef Int -> KeyGen
KeyGen forall a b. (a -> b) -> a -> b
$ 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 forall a b. (a -> b) -> a -> b
$ forall a. IntMap a
IM.empty


------------------------------------------------------------------------------
makeKey :: KeyGen -> IO (Key a)
makeKey :: forall a. KeyGen -> IO (Key a)
makeKey (KeyGen IORef Int
gen) = do
    Int
k <- forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef Int
gen forall {b}. (Ord b, Bounded b, Num b) => b -> (b, b)
nextKey
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Int -> Key a
Key Int
k
  where
    nextKey :: b -> (b, b)
nextKey !b
x = if b
x forall a. Ord a => a -> a -> Bool
>= forall a. Bounded a => a
maxBoundforall a. Num a => a -> a -> a
-b
1
                   then forall a. HasCallStack => [Char] -> a
error [Char]
"too many keys generated"
                   else let !x' :: b
x' = b
xforall 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) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. a -> b
unsafeCoerce forall a b. (a -> b) -> a -> b
$ 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 forall a b. (a -> b) -> a -> b
$
                                                forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
k (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 forall a b. (a -> b) -> a -> b
$
                                              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 forall a b. (a -> b) -> a -> b
$
                                                forall a. (a -> a) -> Int -> IntMap a -> IntMap a
IM.adjust forall a b. a -> b
f' Int
k IntMap Any
m
  where
    f' :: a -> c
f' = forall a b. a -> b
unsafeCoerce forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> b
unsafeCoerce