{-# 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 :: 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 :: KeyGen -> IO (Key a)
makeKey (KeyGen IORef Int
gen) = do
    Int
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
    Key a -> IO (Key a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Key a -> IO (Key a)) -> Key a -> IO (Key a)
forall a b. (a -> b) -> a -> b
$ Int -> Key a
forall a. Int -> Key a
Key Int
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 :: Key a -> HeterogeneousEnvironment -> Maybe a
lookup (Key Int
k) (HeterogeneousEnvironment IntMap Any
m) = (Any -> a) -> Maybe Any -> Maybe a
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 :: 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 :: 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 :: (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