module Data.DefaultMap
where
import Prelude hiding (lookup)
import Control.Applicative ((<$>), (<*>))
import Test.QuickCheck (Arbitrary(..))
import Control.DeepSeq (NFData(..), deepseq)
import qualified Data.HashSet as S
import Data.Hashable
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as Map
import Data.Serialize
import GHC.Generics
data DefaultMap k v =
DefMap { defDefault :: v
, defMap :: HashMap k v
} deriving (Read, Show, Eq, Generic)
instance (Eq k, Hashable k, Serialize k, Serialize v) => Serialize (DefaultMap k v) where
put (DefMap d theMap) = put d >> put (Map.toList theMap)
get = DefMap <$> get <*> (Map.fromList <$> get)
instance (NFData k, NFData v, Hashable k) => NFData (DefaultMap k v) where
rnf (DefMap d m) = d `deepseq` m `deepseq` ()
empty :: v -> DefaultMap k v
empty def = DefMap { defDefault = def
, defMap = Map.empty }
lookup :: (Eq k, Hashable k) => k -> DefaultMap k v -> v
lookup k m = Map.lookupDefault (defDefault m) k (defMap m)
fromList :: (Eq k, Hashable k) => v -> [(k, v)] -> DefaultMap k v
fromList def entries = DefMap { defDefault = def
, defMap = Map.fromList entries }
keys :: DefaultMap k a -> [k]
keys m = Map.keys (defMap m)
elems :: DefaultMap k a -> [a]
elems = Map.elems . defMap
map :: (a -> a) -> DefaultMap k a -> DefaultMap k a
map f m = m { defMap = f <$> defMap m }
foldl :: (a -> b -> a) -> a -> DefaultMap k b -> a
foldl fn acc m = Map.foldl' fn acc (defMap m)
unionWith :: forall v k . (Eq k, Hashable k)
=> (v -> v -> v)
-> v
-> DefaultMap k v
-> DefaultMap k v
-> DefaultMap k v
unionWith f newDef a b =
let allKeys = S.unions [ S.fromList $ Map.keys $ defMap a
, S.fromList $ Map.keys $ defMap b
]
mergeKey :: k -> Map.HashMap k v -> Map.HashMap k v
mergeKey k m = Map.insert k (f (lookup k a) (lookup k b)) m
in DefMap { defDefault = newDef
, defMap = S.foldr mergeKey Map.empty allKeys
}
instance (Arbitrary k, Arbitrary v, Hashable k, Eq k) => Arbitrary (DefaultMap k v) where
arbitrary = do
def <- arbitrary
entries <- arbitrary
return $ fromList def entries