{-# LANGUAGE RecordWildCards #-}
module Data.DAWG.HashMap
( Hash (..)
, HashMap (..)
, empty
, lookup
, insertUnsafe
, lookupUnsafe
, deleteUnsafe
) where
import Prelude hiding (lookup)
import Control.Applicative ((<$>), (<*>))
import Data.Binary (Binary, Get, put, get)
import qualified Data.Map as M
import qualified Data.IntMap as I
fromJust :: Maybe a -> a
fromJust (Just x) = x
fromJust Nothing = error "fromJust: Nothing"
{-# INLINE fromJust #-}
class Ord a => Hash a where
hash :: a -> Int
data Value a b
= Single !a !b
| Multi !(M.Map a b)
deriving (Show, Eq, Ord)
instance (Ord a, Binary a, Binary b) => Binary (Value a b) where
put (Single x y) = put (1 :: Int) >> put x >> put y
put (Multi m) = put (2 :: Int) >> put m
get = do
x <- get :: Get Int
case x of
1 -> Single <$> get <*> get
_ -> Multi <$> get
find :: Ord a => a -> Value a b -> Maybe b
find x (Single x' y) = if x == x'
then Just y
else Nothing
find x (Multi m) = M.lookup x m
findUnsafe :: Ord a => a -> Value a b -> Maybe b
findUnsafe _ (Single _ y) = Just y
findUnsafe x (Multi m) = M.lookup x m
trySingle :: Ord a => M.Map a b -> Value a b
trySingle m = if M.size m == 1
then (uncurry Single) (M.findMin m)
else Multi m
embed :: Ord a => a -> b -> Value a b -> Value a b
embed x y (Single x' y') = Multi $ M.fromList [(x, y), (x', y')]
embed x y (Multi m) = Multi $ M.insert x y m
ejectUnsafe :: Ord a => a -> Value a b -> Maybe (Value a b)
ejectUnsafe _ (Single _ _) = Nothing
ejectUnsafe x (Multi m) = (Just . trySingle) (M.delete x m)
data HashMap a b = HashMap
{ size :: {-# UNPACK #-} !Int
, hashMap :: !(I.IntMap (Value a b)) }
deriving (Show, Eq, Ord)
instance (Ord a, Binary a, Binary b) => Binary (HashMap a b) where
put HashMap{..} = put size >> put hashMap
get = HashMap <$> get <*> get
empty :: HashMap a b
empty = HashMap 0 I.empty
lookup :: Hash a => a -> HashMap a b -> Maybe b
lookup x (HashMap _ m) = I.lookup (hash x) m >>= find x
lookupUnsafe :: Hash a => a -> HashMap a b -> b
lookupUnsafe x (HashMap _ m) = fromJust (I.lookup (hash x) m >>= findUnsafe x)
insertUnsafe :: Hash a => a -> b -> HashMap a b -> HashMap a b
insertUnsafe x y (HashMap n m) =
let i = hash x
f (Just v) = embed x y v
f Nothing = Single x y
in HashMap (n + 1) $ I.alter (Just . f) i m
deleteUnsafe :: Hash a => a -> HashMap a b -> HashMap a b
deleteUnsafe x (HashMap n m) =
HashMap (n - 1) $ I.update (ejectUnsafe x) (hash x) m