1 {-# LANGUAGE BangPatterns #-}
    2 {-# LANGUAGE CPP #-}
    3 {-# LANGUAGE MagicHash #-}
    4 {-# LANGUAGE RankNTypes #-}
    5 {-# LANGUAGE TemplateHaskell #-}
    6 
    7 module Data.Concurrent.HashMap
    8   ( HashMap
    9   , new
   10   , new'
   11   , null
   12   , insert
   13   , delete
   14   , lookup
   15   , update
   16   , fromList
   17   , toList
   18   , hashString
   19   , hashBS
   20   , hashInt
   21   , nextHighestPowerOf2 ) where
   22 
   23 ------------------------------------------------------------------------------
   24 
   25 import           Control.Concurrent.MVar
   26 import           Control.Monad
   27 import           Data.Bits
   28 import qualified Data.ByteString as B
   29 import qualified Data.Digest.Murmur32 as Murmur
   30 import qualified Data.Digest.Murmur64 as Murmur
   31 import           Data.IntMap (IntMap)
   32 import qualified Data.IntMap as IM
   33 import           Data.Maybe
   34 import qualified Data.Vector as V
   35 import           Data.Vector (Vector)
   36 import           GHC.Conc (numCapabilities)
   37 import           Prelude hiding (lookup, null)
   38 import qualified Prelude
   39 
   40 #if __GLASGOW_HASKELL__ >= 503
   41 import GHC.Exts ( Word(..), Int(..), shiftRL# )
   42 #else
   43 import Data.Word
   44 #endif
   45 
   46 import           Data.Concurrent.HashMap.Internal
   47 
   48 
   49 hashString :: String -> Word
   50 hashString = $(whichHash [| Murmur.asWord32 . Murmur.hash32 |]
   51                          [| Murmur.asWord64 . Murmur.hash64 |])
   52 {-# INLINE hashString #-}
   53 
   54 
   55 hashInt :: Int -> Word
   56 hashInt = $(whichHash [| Murmur.asWord32 . Murmur.hash32 |]
   57                       [| Murmur.asWord64 . Murmur.hash64 |])
   58 {-# INLINE hashInt #-}
   59 
   60 
   61 hashBS :: B.ByteString -> Word
   62 hashBS =
   63     $(let h32 = [| \s -> s `seq`
   64                          Murmur.asWord32 $
   65                          B.foldl' (\h c -> h `seq` c `seq`
   66                                            Murmur.hash32AddInt (fromEnum c) h)
   67                                   (Murmur.hash32 ([] :: [Int]))
   68                                   s
   69                 |]
   70           h64 = [| \s -> s `seq`
   71                          Murmur.asWord64 $
   72                          B.foldl' (\h c -> h `seq` c `seq`
   73                                            Murmur.hash64AddInt (fromEnum c) h)
   74                                   (Murmur.hash64 ([] :: [Int]))
   75                                   s
   76                 |]
   77       in whichHash h32 h64)
   78 {-# INLINE hashBS #-}
   79 
   80 
   81 data HashMap k v = HM {
   82       _hash         :: !(k -> Word)
   83     , _hashToBucket :: !(Word -> Word)
   84     , _maps         :: !(Vector (MVar (Submap k v)))
   85 }
   86 
   87 
   88 
   89 null :: HashMap k v -> IO Bool
   90 null ht = liftM V.and $ V.mapM f $ _maps ht
   91 
   92   where
   93     f mv = withMVar mv (return . IM.null)
   94 
   95 
   96 new' :: Eq k =>
   97         Int            -- ^ number of locks to use
   98      -> (k -> Word)    -- ^ hash function
   99      -> IO (HashMap k v)
  100 new' numLocks hashFunc = do
  101     vector <- V.replicateM (fromEnum n) (newMVar IM.empty)
  102     return $! HM hf bh vector
  103 
  104   where
  105     hf !x = hashFunc x
  106     bh !x = x .&. (n-1)
  107     !n    = nextHighestPowerOf2 $ toEnum numLocks
  108 
  109 
  110 new :: Eq k =>
  111        (k -> Word)      -- ^ hash function
  112     -> IO (HashMap k v)
  113 new = new' defaultNumberOfLocks
  114 
  115 
  116 insert :: k -> v -> HashMap k v -> IO ()
  117 insert key value ht =
  118     modifyMVar_ submap $ \m ->
  119         return $! insSubmap hashcode key value m
  120 
  121   where
  122     hashcode = _hash ht key
  123     bucket   = _hashToBucket ht hashcode
  124     submap   = V.unsafeIndex (_maps ht) (fromEnum bucket)
  125 
  126 
  127 delete :: (Eq k) => k -> HashMap k v -> IO ()
  128 delete key ht =
  129     modifyMVar_ submap $ \m ->
  130         return $! delSubmap hashcode key m
  131   where
  132     hashcode = _hash ht key
  133     bucket   = _hashToBucket ht hashcode
  134     submap   = V.unsafeIndex (_maps ht) (fromEnum bucket)
  135 
  136 
  137 lookup :: (Eq k) => k -> HashMap k v -> IO (Maybe v)
  138 lookup key ht =
  139     withMVar submap $ \m ->
  140         return $! lookupSubmap hashcode key m
  141   where
  142     hashcode = _hash ht key
  143     bucket   = _hashToBucket ht hashcode
  144     submap   = V.unsafeIndex (_maps ht) (fromEnum bucket)
  145 
  146 
  147 update :: (Eq k) => k -> v -> HashMap k v -> IO Bool
  148 update key value ht =
  149     modifyMVar submap $ \m ->
  150         return $! updateSubmap hashcode key value m
  151   where
  152     hashcode = _hash ht key
  153     bucket   = _hashToBucket ht hashcode
  154     submap   = V.unsafeIndex (_maps ht) (fromEnum bucket)
  155 
  156 
  157 toList :: HashMap k v -> IO [(k,v)]
  158 toList ht = liftM (concat . V.toList) $ V.mapM f $ _maps ht
  159   where
  160     f m = withMVar m $ \sm -> return $ concat $ IM.elems sm
  161 
  162 
  163 fromList :: (Eq k) => (k -> Word) -> [(k,v)] -> IO (HashMap k v)
  164 fromList hf xs = do
  165     ht <- new hf
  166     mapM_ (\(k,v) -> insert k v ht) xs
  167     return $! ht
  168 
  169 
  170 ------------------------------------------------------------------------------
  171 -- helper functions
  172 ------------------------------------------------------------------------------
  173 
  174 -- nicked this technique from Data.IntMap
  175 
  176 shiftRL :: Word -> Int -> Word
  177 #if __GLASGOW_HASKELL__
  178 {--------------------------------------------------------------------
  179   GHC: use unboxing to get @shiftRL@ inlined.
  180 --------------------------------------------------------------------}
  181 shiftRL (W# x) (I# i)
  182   = W# (shiftRL# x i)
  183 #else
  184 shiftRL x i   = shiftR x i
  185 #endif
  186 
  187 
  188 type Submap k v = IntMap [(k,v)]
  189 
  190 
  191 nextHighestPowerOf2 :: Word -> Word
  192 nextHighestPowerOf2 w = highestBitMask (w-1) + 1
  193 
  194 
  195 highestBitMask :: Word -> Word
  196 highestBitMask !x0 = case (x0 .|. shiftRL x0 1) of
  197                       x1 -> case (x1 .|. shiftRL x1 2) of
  198                        x2 -> case (x2 .|. shiftRL x2 4) of
  199                         x3 -> case (x3 .|. shiftRL x3 8) of
  200                          x4 -> case (x4 .|. shiftRL x4 16) of
  201                           x5 -> x5 .|. shiftRL x5 32
  202 
  203 
  204 
  205 insSubmap :: Word -> k -> v -> Submap k v -> Submap k v
  206 insSubmap hashcode key value m = let !x = f m in x
  207   where
  208     f = IM.insertWith (++) (fromIntegral hashcode) [(key,value)]
  209 
  210 
  211 delSubmap :: (Eq k) => Word -> k -> Submap k v -> Submap k v
  212 delSubmap hashcode key m =
  213     let !z = IM.update f (fromIntegral hashcode) m in z
  214 
  215   where
  216     f l = let l' = del l in if Prelude.null l' then Nothing else Just l'
  217 
  218     del = filter ((/= key) . fst)
  219 
  220 
  221 lookupSubmap :: (Eq k) => Word -> k -> Submap k v -> Maybe v
  222 lookupSubmap hashcode key m = maybe Nothing (Prelude.lookup key) mbBucket
  223   where
  224     mbBucket = IM.lookup (fromIntegral hashcode) m
  225 
  226 
  227 updateSubmap :: (Eq k) => Word -> k -> v -> Submap k v -> (Submap k v, Bool)
  228 updateSubmap hashcode key value m = (m'', b)
  229   where
  230     oldV = lookupSubmap hashcode key m
  231     m'   = maybe m (const $ delSubmap hashcode key m) oldV
  232     m''  = insSubmap hashcode key value m'
  233     b    = isJust oldV
  234 
  235 
  236 defaultNumberOfLocks :: Int
  237 defaultNumberOfLocks = 8 * numCapabilities