{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples       #-}
module Data.Word64Map where

import Data.Bits
import qualified Data.IntMap as IntMap
         ;import Data.IntMap (IntMap)
import Data.Typeable
import Data.Word

-- | Since 'Int' may be 32 or 64 bits, this function is provided as a
-- convenience to test if an integral type, such as 'Data.Word.Word64', can be
-- safely transformed into an 'Int' for use with 'IntMap'.
--
-- Returns 'True' if the proxied type can be losslessly converted to 'Int' using
-- 'fromIntegral'.
fitsInInt :: forall proxy word. (Bounded word, Integral word) => proxy word -> Bool
fitsInInt proxy = (original == casted)
 where
    original = div maxBound 2 :: word
    casted   = fromIntegral (fromIntegral original :: Int) :: word

newtype Word64Map a = Word64Map (IntMap (IntMap a))

empty :: Word64Map a
empty = Word64Map IntMap.empty

-- Warning: This function assumes an 'Int' is either 64 or 32 bits.
keyFrom64 :: Word64 -> (# Int,Int #)
keyFrom64 w8 =
    if fitsInInt (Proxy :: Proxy Word64)
        then (# fromIntegral w8              , 0               #)
        else (# fromIntegral (w8 `shiftR` 32), fromIntegral w8 #)
{-# INLINE keyFrom64 #-}

lookup :: Word64 -> Word64Map b -> Maybe b
lookup w8 (Word64Map m) | (# hi,lo #) <- keyFrom64 w8 = do
    m' <- IntMap.lookup hi m
    IntMap.lookup lo m'
{-# INLINE lookup #-}

insert :: Word64 -> b -> Word64Map b -> Word64Map b
insert w8 b (Word64Map m) | (# hi,lo #) <- keyFrom64 w8
    = Word64Map $ IntMap.alter (Just . maybe (IntMap.singleton lo b)
                                             (IntMap.insert lo b))
                               hi
                               m
{-# INLINE insert #-}

delete :: Word64 -> Word64Map b -> Word64Map b
delete w8 (Word64Map m) | (# hi,lo #) <- keyFrom64 w8
    = Word64Map $ IntMap.alter (maybe Nothing
                                      (\m' -> case IntMap.delete lo m' of
                                                m'' | IntMap.null m'' -> Nothing
                                                m''                   -> Just m''))
                               hi
                               m
{-# INLINE delete #-}