Ticket #4874: HashMap.hs

{-# LANGUAGE BangPatterns, CPP #-}
module Data.HashMap
    (
      HashMap
    , lookup
    ) where

import Data.Bits ((.&.))
import Data.Hashable (Hashable(hash))
import qualified Data.FullList as FL
import Data.Word (Word)
import Prelude hiding (lookup)

data HashMap k v
    = Nil
    | Tip {-# UNPACK #-} !Hash
          {-# UNPACK #-} !(FL.FullList k v)
    | Bin {-# UNPACK #-} !Prefix
          {-# UNPACK #-} !Mask
          !(HashMap k v)
          !(HashMap k v)

type Prefix = Int
type Mask   = Int
type Hash   = Int

lookup :: (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
lookup k0 t = go h0 k0 t
  where
    h0 = hash k0
    go !h !k (Bin _ m l r)
      | zero h m  = go h k l
      | otherwise = go h k r
    go h k (Tip h' l)
      | h == h'   = FL.lookup k l
      | otherwise = Nothing
    go _ _ Nil    = Nothing
#if __GLASGOW_HASKELL__ >= 700
{-# INLINABLE lookup #-}
#endif

zero :: Hash -> Mask -> Bool
zero i m = (fromIntegral i :: Word) .&. (fromIntegral m :: Word) == 0
{-# INLINE zero #-}