{-# LANGUAGE CPP #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.HashMap
-- Copyright   :  (c) Milan Straka 2010
-- License     :  BSD-style
-- Maintainer  :  fox@ucw.cz
-- Stability   :  provisional
-- Portability :  portable
--
-- Persistent 'HashMap', which is defined as
--
-- @
--   data 'HashMap' k v = 'Data.IntMap.IntMap' ('Data.Map.Map' k v)
-- @
--
-- is an 'Data.IntMap.IntMap' indexed by hash values of keys,
-- containing a map @'Data.Map.Map' k v@ with keys of the same hash values.
--
-- The interface of a 'HashMap' is a suitable subset of 'Data.IntMap.IntMap'.
--
-- The complexity of operations is determined by the complexities of
-- 'Data.IntMap.IntMap' and 'Data.Map.Map' operations. See the sources of
-- 'HashMap' to see which operations from @containers@ package are used.
-----------------------------------------------------------------------------

module Data.HashMap ( HashMap

                    -- * Operators
                    , (!), (\\)

                    -- * Query
                    , null
                    , size
                    , member
                    , notMember
                    , lookup
                    , findWithDefault

                    -- * Construction
                    , empty
                    , singleton

                    -- ** Insertion
                    , insert
                    , insertWith, insertWithKey, insertLookupWithKey

                    -- ** Delete\/Update
                    , delete
                    , adjust
                    , adjustWithKey
                    , update
                    , updateWithKey
                    , updateLookupWithKey
                    , alter

                    -- * Combine

                    -- ** Union
                    , union
                    , unionWith
                    , unionWithKey
                    , unions
                    , unionsWith

                    -- ** Difference
                    , difference
                    , differenceWith
                    , differenceWithKey

                    -- ** Intersection
                    , intersection
                    , intersectionWith
                    , intersectionWithKey

                    -- * Traversal
                    -- ** Map
                    , map
                    , mapWithKey
                    , mapAccum
                    , mapAccumWithKey

                    -- ** Fold
                    , fold
                    , foldWithKey

                    -- * Conversion
                    , elems
                    , keys
                    , keysSet
                    , assocs

                    -- ** Lists
                    , toList
                    , fromList
                    , fromListWith
                    , fromListWithKey

                    -- * Filter
                    , filter
                    , filterWithKey
                    , partition
                    , partitionWithKey

                    , mapMaybe
                    , mapMaybeWithKey
                    , mapEither
                    , mapEitherWithKey

                    -- * Submap
                    , isSubmapOf, isSubmapOfBy
                    , isProperSubmapOf, isProperSubmapOfBy
                    ) where

import Prelude hiding (lookup,map,filter,null)

import Control.Applicative (Applicative(pure,(<*>)),(<$>))
import Control.Monad ( liftM )
import Data.Hashable
import Data.Foldable (Foldable(foldMap))
import Data.List (foldl')
import Data.Maybe (fromMaybe)
import Data.Monoid (Monoid(..))
import Data.Traversable (Traversable(traverse))
import Data.Typeable

#if __GLASGOW_HASKELL__
import Text.Read
import Data.Data (Data(..), mkNoRepType)
#endif

import qualified Data.IntMap as I
import qualified Data.Map as M
import qualified Data.Set as S


{--------------------------------------------------------------------
  Operators
--------------------------------------------------------------------}

-- | Find the value at a key.
-- Calls 'error' when the element can not be found.
(!) :: (Hashable k, Ord k) => HashMap k a -> k -> a
m ! k = case lookup k m of
          Nothing -> error "HashMap.(!): key not an element of the map"
          Just v -> v

-- | Same as 'difference'.
(\\) :: Ord k => HashMap k a -> HashMap k b -> HashMap k a
m1 \\ m2 = difference m1 m2


{--------------------------------------------------------------------
  Types
--------------------------------------------------------------------}

-- | The abstract type of a @HashMap@. Its interface is a suitable
-- subset of 'Data.IntMap.IntMap'.
newtype HashMap k v = HashMap (I.IntMap (M.Map k v)) deriving (Eq, Ord)

instance Functor (HashMap k) where
  fmap = map

instance Ord k => Monoid (HashMap k a) where
  mempty  = empty
  mappend = union
  mconcat = unions

instance Foldable (HashMap k) where
  foldMap f (HashMap m) = foldMap (foldMap f) m

instance Traversable (HashMap k) where
  traverse f (HashMap m) = pure HashMap <*> traverse (traverse f) m

instance (Show k, Show a) => Show (HashMap k a) where
  showsPrec d m   = showParen (d > 10) $
    showString "fromList " . shows (toList m)

instance (Read k, Hashable k, Ord k, Read a) => Read (HashMap k a) where
#ifdef __GLASGOW_HASKELL__
  readPrec = parens $ prec 10 $ do
    Ident "fromList" <- lexP
    xs <- readPrec
    return (fromList xs)

  readListPrec = readListPrecDefault
#else
  readsPrec p = readParen (p > 10) $ \ r -> do
    ("fromList",s) <- lex r
    (xs,t) <- reads s
    return (fromList xs,t)
#endif

#include "Typeable.h"
INSTANCE_TYPEABLE2(HashMap,hashMapTc,"HashMap")

#if __GLASGOW_HASKELL__


{--------------------------------------------------------------------
  A Data instance
--------------------------------------------------------------------}

-- This instance preserves data abstraction at the cost of inefficiency.
-- We omit reflection services for the sake of data abstraction.

instance (Data k, Hashable k, Ord k, Data a) => Data (HashMap k a) where
  gfoldl f z m = z fromList `f` (toList m)
  toConstr _   = error "toConstr"
  gunfold _ _  = error "gunfold"
  dataTypeOf _ = mkNoRepType "Data.HashMap.HashMap"
  dataCast1 f  = gcast1 f
#endif


{--------------------------------------------------------------------
  Query
--------------------------------------------------------------------}
-- | Is the map empty?
null :: HashMap k a -> Bool
null (HashMap m) = I.null m

-- | Number of elements in the map.
size :: HashMap k a -> Int
size (HashMap m) = I.fold ((+) . M.size) 0 m

-- | Is the key a member of the map?
member :: (Hashable k, Ord k) => k -> HashMap k a -> Bool
member k m = case lookup k m of
               Nothing -> False
               Just _  -> True

-- | Is the key not a member of the map?
notMember :: (Hashable k, Ord k) => k -> HashMap k a -> Bool
notMember k m = not $ member k m

-- | Lookup the value at a key in the map.
lookup :: (Hashable k, Ord k) => k -> HashMap k a -> Maybe a
lookup k (HashMap m) = I.lookup (hash k) m >>= M.lookup k

-- | The expression @('findWithDefault' def k map)@ returns the value at key
-- @k@ or returns @def@ when the key is not an element of the map.
findWithDefault :: (Hashable k, Ord k) => a -> k -> HashMap k a -> a
findWithDefault def k m = case lookup k m of
                            Nothing -> def
                            Just x  -> x


{--------------------------------------------------------------------
  Construction
--------------------------------------------------------------------}
-- | The empty map.
empty :: HashMap k a
empty = HashMap I.empty

-- | A map of one element.
singleton :: Hashable k => k -> a -> HashMap k a
singleton k x = HashMap $
  I.singleton (hash k) $ M.singleton k x


{--------------------------------------------------------------------
  Insert
--------------------------------------------------------------------}
-- | Insert a new key\/value pair in the map.  If the key is already present in
-- the map, the associated value is replaced with the supplied value, i.e.
-- 'insert' is equivalent to @'insertWith' 'const'@.
insert :: (Hashable k, Ord k)
       => k -> a -> HashMap k a -> HashMap k a
insert k x (HashMap m) = HashMap $
  I.insertWith (\_ -> M.insert k x) (hash k) (M.singleton k x) m

-- | Insert with a combining function.  @'insertWith' f key value mp@ will
-- insert the pair (key, value) into @mp@ if key does not exist in the map. If
-- the key does exist, the function will insert @f new_value old_value@.
insertWith :: (Hashable k, Ord k)
           => (a -> a -> a) -> k -> a -> HashMap k a -> HashMap k a
insertWith f k x (HashMap m) = HashMap $
  I.insertWith (\_ -> M.insertWith f k x) (hash k) (M.singleton k x) m

-- | Insert with a combining function.  @'insertWithKey' f key value mp@ will
-- insert the pair (key, value) into @mp@ if key does not exist in the map. If
-- the key does exist, the function will insert @f key new_value old_value@.
insertWithKey :: (Hashable k, Ord k)
              => (k -> a -> a -> a) -> k -> a -> HashMap k a -> HashMap k a
insertWithKey f k x (HashMap m) = HashMap $
  I.insertWith (\_ -> M.insertWithKey f k x) (hash k) (M.singleton k x) m

-- | The expression (@'insertLookupWithKey' f k x map@) is a pair where the
-- first element is equal to (@'lookup' k map@) and the second element equal to
-- (@'insertWithKey' f k x map@).
insertLookupWithKey :: (Hashable k, Ord k)
                    => (k -> a -> a -> a) -> k -> a -> HashMap k a -> (Maybe a, HashMap k a)
insertLookupWithKey f k x (HashMap m) =
  case I.insertLookupWithKey (\_ _ -> M.insertWithKey f k x) (hash k) (M.singleton k x) m of
    (found, insert) -> (found >>= M.lookup k, HashMap insert)


{--------------------------------------------------------------------
  Deletion
--------------------------------------------------------------------}

nonempty :: M.Map k a -> Maybe (M.Map k a)
nonempty m | M.null m  = Nothing
           | otherwise = Just m

-- | Delete a key and its value from the map. When the key is not
-- a member of the map, the original map is returned.
delete :: (Hashable k, Ord k)
       => k -> HashMap k a -> HashMap k a
delete k (HashMap m) = HashMap $
  I.update (nonempty . M.delete k) (hash k) m

-- | Adjust a value at a specific key. When the key is not a member of the map,
-- the original map is returned.
adjust :: (Hashable k, Ord k)
       => (a -> a) -> k -> HashMap k a -> HashMap k a
adjust f k (HashMap m) = HashMap $
  I.adjust (M.adjust f k) (hash k) m

-- | Adjust a value at a specific key. When the key is not a member of the map,
-- the original map is returned.
adjustWithKey :: (Hashable k, Ord k)
              => (k -> a -> a) -> k -> HashMap k a -> HashMap k a
adjustWithKey f k (HashMap m) = HashMap $
  I.adjust (M.adjustWithKey f k) (hash k) m

-- | The expression (@'update' f k map@) updates the value @x@ at @k@ (if it is
-- in the map). If (@f x@) is 'Nothing', the element is deleted. If it is
-- (@'Just' y@), the key @k@ is bound to the new value @y@.
update :: (Hashable k, Ord k)
       => (a -> Maybe a) -> k -> HashMap k a -> HashMap k a
update f k (HashMap m) = HashMap $
  I.update (nonempty . M.update f k) (hash k) m

-- | The expression (@'update' f k map@) updates the value @x@ at @k@ (if it is
-- in the map). If (@f k x@) is 'Nothing', the element is deleted. If it is
-- (@'Just' y@), the key @k@ is bound to the new value @y@.
updateWithKey :: (Hashable k, Ord k)
              => (k -> a -> Maybe a) -> k -> HashMap k a -> HashMap k a
updateWithKey f k (HashMap m) = HashMap $
  I.update (nonempty . M.updateWithKey f k) (hash k) m

-- | Lookup and update.  The function returns original value, if it is updated.
-- This is different behavior than 'Data.Map.updateLookupWithKey'.  Returns the
-- original key value if the map entry is deleted.
updateLookupWithKey :: (Hashable k, Ord k)
                    => (k -> a -> Maybe a) -> k -> HashMap k a -> (Maybe a, HashMap k a)
updateLookupWithKey f k (HashMap m) =
  case I.updateLookupWithKey (\_ -> nonempty . M.updateWithKey f k) (hash k) m of
    (found, update) -> (found >>= M.lookup k, HashMap update)

-- | The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence
-- thereof.  'alter' can be used to insert, delete, or update a value in an
-- 'HashMap'.
alter :: (Hashable k, Ord k)
      => (Maybe a -> Maybe a) -> k -> HashMap k a -> HashMap k a
alter f k (HashMap m) = HashMap $
  I.alter (nonempty . M.alter f k . fromMaybe M.empty) (hash k) m


{--------------------------------------------------------------------
  Union
--------------------------------------------------------------------}
-- | The union of a list of maps.
unions :: Ord k => [HashMap k a] -> HashMap k a
unions xs = foldl' union empty xs

-- | The union of a list of maps, with a combining operation.
unionsWith :: Ord k => (a->a->a) -> [HashMap k a] -> HashMap k a
unionsWith f xs = foldl' (unionWith f) empty xs

-- | The (left-biased) union of two maps.
-- It prefers the first map when duplicate keys are encountered,
-- i.e. (@'union' == 'unionWith' 'const'@).
union :: Ord k => HashMap k a -> HashMap k a -> HashMap k a
union (HashMap m1) (HashMap m2) = HashMap $
  I.unionWith M.union m1 m2

-- | The union with a combining function.
unionWith :: Ord k => (a -> a -> a) -> HashMap k a -> HashMap k a -> HashMap k a
unionWith f (HashMap m1) (HashMap m2) = HashMap $
  I.unionWith (M.unionWith f) m1 m2

-- | The union with a combining function.
unionWithKey :: Ord k => (k -> a -> a -> a) -> HashMap k a -> HashMap k a -> HashMap k a
unionWithKey f (HashMap m1) (HashMap m2) = HashMap $
  I.unionWith (M.unionWithKey f) m1 m2


{--------------------------------------------------------------------
  Difference
--------------------------------------------------------------------}
-- | Difference between two maps (based on keys).
difference :: Ord k => HashMap k a -> HashMap k b -> HashMap k a
difference (HashMap m1) (HashMap m2) = HashMap $
  I.differenceWith (\n1 n2 -> nonempty $ M.difference n1 n2) m1 m2

-- | Difference with a combining function.
differenceWith :: Ord k => (a -> b -> Maybe a) -> HashMap k a -> HashMap k b -> HashMap k a
differenceWith f (HashMap m1) (HashMap m2) = HashMap $
  I.differenceWith (\n1 n2 -> nonempty $ M.differenceWith f n1 n2) m1 m2

-- | Difference with a combining function. When two equal keys are
-- encountered, the combining function is applied to the key and both values.
-- If it returns 'Nothing', the element is discarded (proper set difference).
-- If it returns (@'Just' y@), the element is updated with a new value @y@. 
differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> HashMap k a -> HashMap k b -> HashMap k a
differenceWithKey f (HashMap m1) (HashMap m2) = HashMap $
  I.differenceWith (\n1 n2 -> nonempty $ M.differenceWithKey f n1 n2) m1 m2


{--------------------------------------------------------------------
  Intersection
--------------------------------------------------------------------}
delete_empty :: I.IntMap (M.Map k a) -> I.IntMap (M.Map k a)
delete_empty = I.filter (not . M.null)

-- | The (left-biased) intersection of two maps (based on keys).
intersection :: Ord k => HashMap k a -> HashMap k b -> HashMap k a
intersection (HashMap m1) (HashMap m2) = HashMap $ delete_empty $
  I.intersectionWith M.intersection m1 m2

-- | The intersection with a combining function.
intersectionWith :: Ord k => (a -> b -> c) -> HashMap k a -> HashMap k b -> HashMap k c
intersectionWith f (HashMap m1) (HashMap m2) = HashMap $ delete_empty $
  I.intersectionWith (M.intersectionWith f) m1 m2

-- | The intersection with a combining function.
intersectionWithKey :: Ord k => (k -> a -> b -> c) -> HashMap k a -> HashMap k b -> HashMap k c
intersectionWithKey f (HashMap m1) (HashMap m2) = HashMap $ delete_empty $
  I.intersectionWith (M.intersectionWithKey f) m1 m2


{--------------------------------------------------------------------
  Submap
--------------------------------------------------------------------}
-- | Is this a proper submap? (ie. a submap but not equal).
isProperSubmapOf :: (Ord k, Eq a) => HashMap k a -> HashMap k a -> Bool
isProperSubmapOf m1 m2 = isSubmapOf m1 m2 && size m1 < size m2

-- | Is this a proper submap? (ie. a submap but not equal).  The expression
-- (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when @m1@ and @m2@ are not
-- equal, all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
-- applied to their respective values.
isProperSubmapOfBy :: Ord k => (a -> b -> Bool) -> HashMap k a -> HashMap k b -> Bool
isProperSubmapOfBy f m1 m2 = isSubmapOfBy f m1 m2 && size m1 < size m2

-- | Is this a submap?
isSubmapOf :: (Ord k, Eq a) => HashMap k a -> HashMap k a -> Bool
isSubmapOf (HashMap m1) (HashMap m2) =
  I.isSubmapOfBy (M.isSubmapOf) m1 m2

-- | The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if all keys in
-- @m1@ are in @m2@, and when @f@ returns 'True' when applied to their
-- respective values.
isSubmapOfBy :: Ord k => (a -> b -> Bool) -> HashMap k a -> HashMap k b -> Bool
isSubmapOfBy f (HashMap m1) (HashMap m2) =
  I.isSubmapOfBy (M.isSubmapOfBy f) m1 m2


{--------------------------------------------------------------------
  Mapping
--------------------------------------------------------------------}
-- | Map a function over all values in the map.
map :: (a -> b) -> HashMap k a -> HashMap k b
map f (HashMap m) = HashMap $
  I.map (M.map f) m

-- | Map a function over all values in the map.
mapWithKey :: (k -> a -> b) -> HashMap k a -> HashMap k b
mapWithKey f (HashMap m) = HashMap $
  I.map (M.mapWithKey f) m

-- | The function @'mapAccum'@ threads an accumulating argument through the map
-- in unspecified order of keys.
mapAccum :: (a -> b -> (a,c)) -> a -> HashMap k b -> (a,HashMap k c)
mapAccum f a (HashMap m) =
  case I.mapAccum (M.mapAccum f) a m of
    (acc, m) -> (acc, HashMap m)

-- | The function @'mapAccumWithKey'@ threads an accumulating argument through
-- the map in unspecified order of keys.
mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> HashMap k b -> (a,HashMap k c)
mapAccumWithKey f a (HashMap m) =
  case I.mapAccum (M.mapAccumWithKey f) a m of
    (acc, m) -> (acc, HashMap m)


{--------------------------------------------------------------------
  Filter
--------------------------------------------------------------------}
-- | Filter all values that satisfy some predicate.
filter :: Ord k => (a -> Bool) -> HashMap k a -> HashMap k a
filter p (HashMap m) = HashMap $
  I.mapMaybe (nonempty . M.filter p) m

-- | Filter all keys\/values that satisfy some predicate.
filterWithKey :: Ord k => (k -> a -> Bool) -> HashMap k a -> HashMap k a
filterWithKey p (HashMap m) = HashMap $
  I.mapMaybe (nonempty . M.filterWithKey p) m

-- | Partition the map according to some predicate. The first map contains all
-- elements that satisfy the predicate, the second all elements that fail the
-- predicate. See also 'split'.
partition :: Ord k => (a -> Bool) -> HashMap k a -> (HashMap k a, HashMap k a)
partition p m = (mapMaybe (maybe_true p) m, mapMaybe (maybe_false p) m)

-- | Partition the map according to some predicate. The first map contains all
-- elements that satisfy the predicate, the second all elements that fail the
-- predicate.
partitionWithKey :: Ord k => (k -> a -> Bool) -> HashMap k a -> (HashMap k a, HashMap k a)
partitionWithKey p m = (mapMaybeWithKey (\k -> maybe_true  (p k)) m
                       ,mapMaybeWithKey (\k -> maybe_false (p k)) m)

-- | Map values and collect the 'Just' results.
mapMaybe :: Ord k => (a -> Maybe b) -> HashMap k a -> HashMap k b
mapMaybe f (HashMap m) = HashMap $
  I.mapMaybe (nonempty . M.mapMaybe f) m

-- | Map keys\/values and collect the 'Just' results.
mapMaybeWithKey :: Ord k => (k -> a -> Maybe b) -> HashMap k a -> HashMap k b
mapMaybeWithKey f (HashMap m) = HashMap $
  I.mapMaybe (nonempty . M.mapMaybeWithKey f) m

-- | Map values and separate the 'Left' and 'Right' results.
mapEither :: Ord k => (a -> Either b c) -> HashMap k a -> (HashMap k b, HashMap k c)
mapEither f m = (mapMaybe (maybe_left . f) m, mapMaybe (maybe_right . f) m)

-- | Map keys\/values and separate the 'Left' and 'Right' results.
mapEitherWithKey :: Ord k => (k -> a -> Either b c) -> HashMap k a -> (HashMap k b, HashMap k c)
mapEitherWithKey f m = (mapMaybeWithKey (\k a -> maybe_left  (f k a)) m
                       ,mapMaybeWithKey (\k a -> maybe_right (f k a)) m)

-- Helper functions for this section
maybe_left (Left a) = Just a
maybe_left (Right _) = Nothing

maybe_right (Right a) = Just a
maybe_right (Left _) = Nothing

maybe_true  p a = if p a then Just a else Nothing

maybe_false p a = if p a then Nothing else Just a

{--------------------------------------------------------------------
  Fold
--------------------------------------------------------------------}
-- | Fold the values in the map, such that @'fold' f z == 'Prelude.foldr'
-- f z . 'elems'@.
fold :: (a -> b -> b) -> b -> HashMap k a -> b
fold f z (HashMap m) = I.fold (flip $ M.fold f) z m

-- | Fold the keys and values in the map, such that @'foldWithKey' f z ==
-- 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
foldWithKey :: (k -> a -> b -> b) -> b -> HashMap k a -> b
foldWithKey f z (HashMap m) = I.fold (flip $ M.foldWithKey f) z m


{--------------------------------------------------------------------
  List variations
--------------------------------------------------------------------}
-- | Return all elements of the map in arbitrary order of their keys.
elems :: HashMap k a -> [a]
elems (HashMap m) = I.fold ((++) . M.elems) [] m

-- | Return all keys of the map in arbitrary order.
keys  :: HashMap k a -> [k]
keys (HashMap m) = I.fold ((++) . M.keys) [] m

-- | The set of all keys of the map.
keysSet :: Ord k => HashMap k a -> S.Set k
keysSet (HashMap m) = I.fold (S.union . M.keysSet) S.empty m

-- | Return all key\/value pairs in the map in arbitrary key order.
assocs :: HashMap k a -> [(k,a)]
assocs = toList


{--------------------------------------------------------------------
  Lists
--------------------------------------------------------------------}
-- | Convert the map to a list of key\/value pairs.
toList :: HashMap k a -> [(k,a)]
toList (HashMap m) =
  I.fold ((++) . M.toList) [] m

-- | Create a map from a list of key\/value pairs.
fromList :: (Hashable k, Ord k)
         => [(k,a)] -> HashMap k a
fromList xs = foldl' (\m (k, v) -> insert k v m) empty xs

-- | Create a map from a list of key\/value pairs with a combining function.
fromListWith :: (Hashable k, Ord k) => (a -> a -> a) -> [(k,a)] -> HashMap k a
fromListWith f xs = foldl' (\m (k, v) -> insertWith f k v m) empty xs

-- | Build a map from a list of key\/value pairs with a combining function.
fromListWithKey :: (Hashable k, Ord k) => (k -> a -> a -> a) -> [(k,a)] -> HashMap k a
fromListWithKey f xs = foldl' (\m (k, v) -> insertWithKey f k v m) empty xs