{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples       #-}
{-# LANGUAGE CPP                 #-}
-- |  This is a wrapper around 'Data.IntMap', from the containers package, but with a
--    guaranteed bitwidth of 64 bits.
--
--    On 32 bit platforms, this is currently accomplished simply by composing two IntMaps.
--
--    On obvious 64 bit platforms(platform name as shown by System.Info.arch ends in 64), CPP
--    is used and it is simply a newtype around IntMap.
--
--    If the CPP is not defined, the Word64Map will be a composition of two IntMaps, but it should
--    work anyway provided Int is either 32 or 64 bits. Nevertheless, for highest guaranteed efficiency,
--    report your platform so it can be detected and the CPP defined accordingly.
--
module Data.Word64Map
        ( Word64Map
        , fitsInInt
        , Data.Word64Map.lookup
        , insert
        , delete
        , size
        , empty
        ) where

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

-- | Returns 'True' if the proxied type can be losslessly converted to 'Int' using
-- 'fromIntegral'.
--
-- 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'.
--
-- It should be optimized away at runtime.
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

#if KNOWN64
newtype Word64Map a = Word64Map (IntMap a)
#else
newtype Word64Map a = Word64Map (IntMap (IntMap a))
#endif


#if KNOWN64
size :: Word64Map a -> Int
size (Word64Map mp) = IntMap.size mp
{-# INLINE size #-}
empty :: Word64Map a
empty = Word64Map (IntMap.empty)
{-# INLINE empty #-}
lookup :: Word64 -> Word64Map b -> Maybe b
lookup key (Word64Map mp) = IntMap.lookup (fromIntegral key) mp
{-# INLINE lookup #-}
insert :: Word64 -> b -> Word64Map b -> Word64Map b
insert key val (Word64Map mp) = Word64Map (IntMap.insert (fromIntegral key) val mp)
{-# INLINE insert #-}
delete :: Word64 -> Word64Map b -> Word64Map b
delete key (Word64Map m) = Word64Map (IntMap.delete (fromIntegral key) m)
{-# INLINE delete #-}
#else
size :: Word64Map a -> Int
size (Word64Map m) = getSum $ foldMap (\n -> Sum (IntMap.size n)) m

empty :: Word64Map a
empty = Word64Map IntMap.empty

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

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

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

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