    {-# LANGUAGE BangPatterns, CPP #-}

    #if __GLASGOW_HASKELL__ >= 702
    {-# LANGUAGE Trustworthy #-}
    #endif

    ------------------------------------------------------------------------
    -- |
    -- Module      :  Data.HashMap.Strict
    -- Copyright   :  2010-2012 Johan Tibell
    -- License     :  BSD-style
    -- Maintainer  :  johan.tibell@gmail.com
    -- Stability   :  provisional
    -- Portability :  portable
    --
    -- A map from /hashable/ keys to values.  A map cannot contain
    -- duplicate keys; each key can map to at most one value.  A 'HashMap'
    -- makes no guarantees as to the order of its elements.
    --
    -- This map is strict in both the keys and the values; keys and values
    -- are evaluated to /weak head normal form/ before they are added to
    -- the map.  Exception: the provided instances are the same as for the
    -- lazy version of this module.
    --
    -- The implementation is based on /hash array mapped tries/.  A
    -- 'HashMap' is often faster than other tree-based set types,
    -- especially when key comparison is expensive, as in the case of
    -- strings.
    --
    -- Many operations have a average-case complexity of /O(log n)/.  The
    -- implementation uses a large base (i.e. 16) so in practice these
    -- operations are constant time.
    module Data.HashMap.Strict
        (
          HashMap

          -- * Construction
        , empty
        , singleton

          -- * Basic interface
        , HM.null
        , size
        , HM.member
        , HM.lookup
        , lookupDefault
        , (!)
        , insert
        , insertWith
        , delete
        , adjust

          -- * Combine
          -- ** Union
        , union
        , unionWith
        , unions

          -- * Transformations
        , map
        , traverseWithKey

          -- * Difference and intersection
        , difference
        , intersection
        , intersectionWith

          -- * Folds
        , foldl'
        , foldlWithKey'
        , HM.foldr
        , foldrWithKey

          -- * Filter
        , HM.filter
        , filterWithKey

          -- * Conversions
        , keys
        , elems

          -- ** Lists
        , toList
        , fromList
        , fromListWith
        ) where

    import Data.Bits ((.&.), (.|.))
    import qualified Data.List as L
    import Data.Hashable (Hashable)
    import Prelude hiding (map)

    import qualified Data.HashMap.Array as A
    import qualified Data.HashMap.Base as HM
    import Data.HashMap.Base hiding (
        adjust, fromList, fromListWith, insert, insertWith, intersectionWith,
        lookupDefault, map, singleton, unionWith)
    import Data.HashMap.Unsafe (runST)

    ------------------------------------------------------------------------
    -- * Construction

    -- | /O(1)/ Construct a map with a single element.
    singleton :: (Hashable k) => k -> v -> HashMap k v
    singleton k !v = HM.singleton k v

    ------------------------------------------------------------------------
    -- * Basic interface

    -- | /O(log n)/ Return the value to which the specified key is mapped,
    -- or the default value if this map contains no mapping for the key.
    lookupDefault :: (Eq k, Hashable k)
                  => v          -- ^ Default value to return.
                  -> k -> HashMap k v -> v
    lookupDefault !def k t = HM.lookupDefault def k t
    {-# INLINABLE lookupDefault #-}

    -- | /O(log n)/ Associate the specified value with the specified
    -- key in this map.  If this map previously contained a mapping for
    -- the key, the old value is replaced.
    insert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v
    insert k !v = HM.insert k v
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE insert #-}
    #endif

    -- | /O(log n)/ Associate the value with the key in this map.  If
    -- this map previously contained a mapping for the key, the old value
    -- is replaced by the result of applying the given function to the new
    -- and old value.  Example:
    --
    -- > insertWith f k v map
    -- >   where f new old = new + old
    insertWith :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v
               -> HashMap k v
    insertWith f k0 !v0 m0 = go h0 k0 v0 0 m0
      where
        h0 = hash k0
        go !h !k x !_ Empty = Leaf h (L k x)
        go h k x s (Leaf hy l@(L ky y))
            | hy == h = if ky == k
                        then let !v' = f x y in Leaf h (L k v')
                        else collision h l (L k x)
            | otherwise = runST (two s h k x hy ky y)
        go h k x s (BitmapIndexed b ary)
            | b .&. m == 0 =
                let ary' = A.insert ary i $! Leaf h (L k x)
                in bitmapIndexedOrFull (b .|. m) ary'
            | otherwise =
                let st   = A.index ary i
                    st'  = go h k x (s+bitsPerSubkey) st
                    ary' = A.update ary i $! st'
                in BitmapIndexed b ary'
          where m = mask h s
                i = sparseIndex b m
        go h k x s (Full ary) =
            let st   = A.index ary i
                st'  = go h k x (s+bitsPerSubkey) st
                ary' = update16 ary i $! st'
            in Full ary'
          where i = index h s
        go h k x s t@(Collision hy v)
            | h == hy   = Collision h (updateOrSnocWith f k x v)
            | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE insertWith #-}
    #endif

    -- | In-place update version of insertWith
    unsafeInsertWith :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v
                     -> HashMap k v
    unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
      where
        h0 = hash k0
        go !h !k x !_ Empty = return $! Leaf h (L k x)
        go h k x s (Leaf hy l@(L ky y))
            | hy == h = if ky == k
                        then let !v' = f x y in return $! Leaf h (L k v')
                        else return $! collision h l (L k x)
            | otherwise = two s h k x hy ky y
        go h k x s t@(BitmapIndexed b ary)
            | b .&. m == 0 = do
                ary' <- A.insertM ary i $! Leaf h (L k x)
                return $! bitmapIndexedOrFull (b .|. m) ary'
            | otherwise = do
                st <- A.indexM ary i
                st' <- go h k x (s+bitsPerSubkey) st
                A.unsafeUpdateM ary i st'
                return t
          where m = mask h s
                i = sparseIndex b m
        go h k x s t@(Full ary) = do
            st <- A.indexM ary i
            st' <- go h k x (s+bitsPerSubkey) st
            A.unsafeUpdateM ary i st'
            return t
          where i = index h s
        go h k x s t@(Collision hy v)
            | h == hy   = return $! Collision h (updateOrSnocWith f k x v)
            | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE unsafeInsertWith #-}
    #endif

    -- | /O(log n)/ Adjust the value tied to a given key in this map only
    -- if it is present. Otherwise, leave the map alone.
    adjust :: (Eq k, Hashable k) => (v -> v) -> k -> HashMap k v -> HashMap k v
    adjust f k0 = go h0 k0 0
      where
        h0 = hash k0
        go !_ !_ !_ Empty = Empty
        go h k _ t@(Leaf hy (L ky y))
            | hy == h && ky == k = let !v' = f y in Leaf h (L k v')
            | otherwise          = t
        go h k s t@(BitmapIndexed b ary)
            | b .&. m == 0 = t
            | otherwise = let st   = A.index ary i
                              st'  = go h k (s+bitsPerSubkey) st
                              ary' = A.update ary i $! st'
                          in BitmapIndexed b ary'
          where m = mask h s
                i = sparseIndex b m
        go h k s (Full ary) =
            let i    = index h s
                st   = A.index ary i
                st'  = go h k (s+bitsPerSubkey) st
                ary' = update16 ary i $! st'
            in Full ary'
        go h k _ t@(Collision hy v)
            | h == hy   = Collision h (updateWith f k v)
            | otherwise = t
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE adjust #-}
    #endif

    ------------------------------------------------------------------------
    -- * Combine

    -- | /O(n+m)/ The union of two maps.  If a key occurs in both maps,
    -- the provided function (first argument) will be used to compute the result.
    unionWith :: (Eq k, Hashable k) => (v -> v -> v) -> HashMap k v -> HashMap k v
              -> HashMap k v
    unionWith f = go 0
      where
        -- empty vs. anything
        go !_ t1 Empty = t1
        go _ Empty t2 = t2
        -- leaf vs. leaf
        go s t1@(Leaf h1 l1@(L k1 v1)) t2@(Leaf h2 l2@(L k2 v2))
            | h1 == h2  = if k1 == k2
                          then Leaf h1 . L k1 $! f v1 v2
                          else collision h1 l1 l2
            | otherwise = goDifferentHash s h1 h2 t1 t2
        go s t1@(Leaf h1 (L k1 v1)) t2@(Collision h2 ls2)
            | h1 == h2  = Collision h1 (updateOrSnocWith f k1 v1 ls2)
            | otherwise = goDifferentHash s h1 h2 t1 t2
        go s t1@(Collision h1 ls1) t2@(Leaf h2 (L k2 v2))
            | h1 == h2  = Collision h1 (updateOrSnocWith (flip f) k2 v2 ls1)
            | otherwise = goDifferentHash s h1 h2 t1 t2
        go s t1@(Collision h1 ls1) t2@(Collision h2 ls2)
            | h1 == h2  = Collision h1 (updateOrConcatWith f ls1 ls2)
            | otherwise = goDifferentHash s h1 h2 t1 t2
        -- branch vs. branch
        go s (BitmapIndexed b1 ary1) (BitmapIndexed b2 ary2) =
            let b'   = b1 .|. b2
                ary' = unionArrayBy (go (s+bitsPerSubkey)) b1 b2 ary1 ary2
            in bitmapIndexedOrFull b' ary'
        go s (BitmapIndexed b1 ary1) (Full ary2) =
            let ary' = unionArrayBy (go (s+bitsPerSubkey)) b1 fullNodeMask ary1 ary2
            in Full ary'
        go s (Full ary1) (BitmapIndexed b2 ary2) =
            let ary' = unionArrayBy (go (s+bitsPerSubkey)) fullNodeMask b2 ary1 ary2
            in Full ary'
        go s (Full ary1) (Full ary2) =
            let ary' = unionArrayBy (go (s+bitsPerSubkey)) fullNodeMask fullNodeMask
                       ary1 ary2
            in Full ary'
        -- leaf vs. branch
        go s (BitmapIndexed b1 ary1) t2
            | b1 .&. m2 == 0 = let ary' = A.insert ary1 i t2
                                   b'   = b1 .|. m2
                               in bitmapIndexedOrFull b' ary'
            | otherwise      = let ary' = A.updateWith' ary1 i $ \st1 ->
                                       go (s+bitsPerSubkey) st1 t2
                               in BitmapIndexed b1 ary'
            where
              h2 = leafHashCode t2
              m2 = mask h2 s
              i = sparseIndex b1 m2
        go s t1 (BitmapIndexed b2 ary2)
            | b2 .&. m1 == 0 = let ary' = A.insert ary2 i $! t1
                                   b'   = b2 .|. m1
                               in bitmapIndexedOrFull b' ary'
            | otherwise      = let ary' = A.updateWith' ary2 i $ \st2 ->
                                       go (s+bitsPerSubkey) t1 st2
                               in BitmapIndexed b2 ary'
          where
            h1 = leafHashCode t1
            m1 = mask h1 s
            i = sparseIndex b2 m1
        go s (Full ary1) t2 =
            let h2   = leafHashCode t2
                i    = index h2 s
                ary' = update16With' ary1 i $ \st1 -> go (s+bitsPerSubkey) st1 t2
            in Full ary'
        go s t1 (Full ary2) =
            let h1   = leafHashCode t1
                i    = index h1 s
                ary' = update16With' ary2 i $ \st2 -> go (s+bitsPerSubkey) t1 st2
            in Full ary'

        leafHashCode (Leaf h _) = h
        leafHashCode (Collision h _) = h
        leafHashCode _ = error "leafHashCode"

        goDifferentHash s h1 h2 t1 t2
            | m1 == m2  = BitmapIndexed m1 (A.singleton $! go (s+bitsPerSubkey) t1 t2)
            | m1 <  m2  = BitmapIndexed (m1 .|. m2) (A.pair t1 t2)
            | otherwise = BitmapIndexed (m1 .|. m2) (A.pair t2 t1)
          where
            m1 = mask h1 s
            m2 = mask h2 s
    {-# INLINE unionWith #-}

    ------------------------------------------------------------------------
    -- * Transformations

    -- | /O(n)/ Transform this map by applying a function to every value.
    mapWithKey :: (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2
    mapWithKey f = go
      where
        go Empty                 = Empty
        go (Leaf h (L k v))      = let !v' = f k v in Leaf h $ L k v'
        go (BitmapIndexed b ary) = BitmapIndexed b $ A.map' go ary
        go (Full ary)            = Full $ A.map' go ary
        go (Collision h ary)     =
            Collision h $ A.map' (\ (L k v) -> let !v' = f k v in L k v') ary
    {-# INLINE mapWithKey #-}

    map :: (v1 -> v2) -> HashMap k v1 -> HashMap k v2
    map f = mapWithKey (const f)
    {-# INLINE map #-}

    -- TODO: Should we add a strict traverseWithKey?

    ------------------------------------------------------------------------
    -- * Difference and intersection

    -- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps
    -- the provided function is used to combine the values from the two
    -- maps.
    intersectionWith :: (Eq k, Hashable k) => (v1 -> v2 -> v3) -> HashMap k v1
                     -> HashMap k v2 -> HashMap k v3
    intersectionWith f a b = foldlWithKey' go empty a
      where
        go m k v = case HM.lookup k b of
                     Just w -> insert k (f v w) m
                     _      -> m
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE intersectionWith #-}
    #endif

    ------------------------------------------------------------------------
    -- ** Lists

    -- | /O(n*log n)/ Construct a map with the supplied mappings.  If the
    -- list contains duplicate mappings, the later mappings take
    -- precedence.
    fromList :: (Eq k, Hashable k) => [(k, v)] -> HashMap k v
    fromList = L.foldl' (\ m (k, v) -> HM.unsafeInsert k v m) empty
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE fromList #-}
    #endif

    -- | /O(n*log n)/ Construct a map from a list of elements.  Uses
    -- the provided function to merge duplicate entries.
    fromListWith :: (Eq k, Hashable k) => (v -> v -> v) -> [(k, v)] -> HashMap k v
    fromListWith f = L.foldl' (\ m (k, v) -> unsafeInsertWith f k v m) empty
    {-# INLINE fromListWith #-}

    ------------------------------------------------------------------------
    -- Array operations

    updateWith :: Eq k => (v -> v) -> k -> A.Array (Leaf k v) -> A.Array (Leaf k v)
    updateWith f k0 ary0 = go k0 ary0 0 (A.length ary0)
      where
        go !k !ary !i !n
            | i >= n    = ary
            | otherwise = case A.index ary i of
                (L kx y) | k == kx   -> let !v' = f y in A.update ary i (L k v')
                         | otherwise -> go k ary (i+1) n
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE updateWith #-}
    #endif

    updateOrSnocWith :: Eq k => (v -> v -> v) -> k -> v -> A.Array (Leaf k v)
                     -> A.Array (Leaf k v)
    updateOrSnocWith f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0)
      where
        go !k v !ary !i !n
            | i >= n = A.run $ do
                -- Not found, append to the end.
                mary <- A.new_ (n + 1)
                A.copy ary 0 mary 0 n
                A.write mary n (L k v)
                return mary
            | otherwise = case A.index ary i of
                (L kx y) | k == kx   -> let !v' = f v y in A.update ary i (L k v')
                         | otherwise -> go k v ary (i+1) n
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE updateOrSnocWith #-}
    #endif