    {-# LANGUAGE BangPatterns, CPP, DeriveDataTypeable, MagicHash #-}
    {-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-}

    module Data.HashMap.Base
        (
          HashMap(..)
        , Leaf(..)

          -- * Construction
        , empty
        , singleton

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

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

          -- * Transformations
        , map
        , traverseWithKey

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

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

          -- * Filter
        , filter
        , filterWithKey

          -- * Conversions
        , keys
        , elems

          -- ** Lists
        , toList
        , fromList
        , fromListWith

          -- Internals used by the strict version
        , Bitmap
        , bitmapIndexedOrFull
        , collision
        , hash
        , mask
        , index
        , bitsPerSubkey
        , fullNodeMask
        , sparseIndex
        , two
        , unionArrayBy
        , update16
        , update16M
        , update16With'
        , updateOrConcatWith
        ) where

    import Control.Applicative ((<$>), Applicative(pure))
    import Control.DeepSeq (NFData(rnf))
    import Control.Monad.ST (ST)
    import Data.Bits ((.&.), (.|.), complement)
    import qualified Data.Foldable as Foldable
    import qualified Data.List as L
    import Data.Monoid (Monoid(mempty, mappend))
    import Data.Traversable (Traversable(..))
    import Data.Word (Word)
    import Prelude hiding (filter, foldr, lookup, map, null, pred)

    import qualified Data.HashMap.Array as A
    import qualified Data.Hashable as H
    import Data.Hashable (Hashable)
    import Data.HashMap.PopCount (popCount)
    import Data.HashMap.Unsafe (runST)
    import Data.HashMap.UnsafeShift (unsafeShiftL, unsafeShiftR)
    import Data.Typeable (Typeable)

    #if defined(__GLASGOW_HASKELL__)
    import Data.Data hiding (Typeable)
    import GHC.Exts ((==#), build, reallyUnsafePtrEquality#)
    #endif

    ------------------------------------------------------------------------

    -- | Convenience function.  Compute a hash value for the given value.
    hash :: H.Hashable a => a -> Hash
    hash = fromIntegral . H.hash

    data Leaf k v = L !k v
      deriving (Eq)

    instance (NFData k, NFData v) => NFData (Leaf k v) where
        rnf (L k v) = rnf k `seq` rnf v

    -- Invariant: The length of the 1st argument to 'Full' is
    -- 2^bitsPerSubkey

    -- | A map from keys to values.  A map cannot contain duplicate keys;
    -- each key can map to at most one value.
    data HashMap k v
        = Empty
        | BitmapIndexed !Bitmap !(A.Array (HashMap k v))
        | Leaf !Hash !(Leaf k v)
        | Full !(A.Array (HashMap k v))
        | Collision !Hash !(A.Array (Leaf k v))
          deriving (Typeable)

    instance (NFData k, NFData v) => NFData (HashMap k v) where
        rnf Empty                 = ()
        rnf (BitmapIndexed _ ary) = rnf ary
        rnf (Leaf _ l)            = rnf l
        rnf (Full ary)            = rnf ary
        rnf (Collision _ ary)     = rnf ary

    instance Functor (HashMap k) where
        fmap = map

    instance Foldable.Foldable (HashMap k) where
        foldr f = foldrWithKey (const f)

    instance (Eq k, Hashable k) => Monoid (HashMap k v) where
      mempty = empty
      {-# INLINE mempty #-}
      mappend = union
      {-# INLINE mappend #-}

    #if __GLASGOW_HASKELL__
    instance (Data k, Data v, Eq k, Hashable k) => Data (HashMap k v) where
        gfoldl f z m   = z fromList `f` toList m
        toConstr _     = fromListConstr
        gunfold k z c  = case constrIndex c of
            1 -> k (z fromList)
            _ -> error "gunfold"
        dataTypeOf _   = hashMapDataType
        dataCast2 f    = gcast2 f

    fromListConstr :: Constr
    fromListConstr = mkConstr hashMapDataType "fromList" [] Prefix

    hashMapDataType :: DataType
    hashMapDataType = mkDataType "Data.HashMap.Base.HashMap" [fromListConstr]
    #endif

    type Hash   = Word
    type Bitmap = Word
    type Shift  = Int

    instance (Show k, Show v) => Show (HashMap k v) where
        show m = "fromList " ++ show (toList m)

    instance Traversable (HashMap k) where
        traverse f = traverseWithKey (const f)

    instance (Eq k, Eq v) => Eq (HashMap k v) where
        (==) = equal

    equal :: (Eq k, Eq v) => HashMap k v -> HashMap k v -> Bool
    equal t1 t2 = go (toList' t1 []) (toList' t2 [])
      where
        -- If the two trees are the same, then their lists of 'Leaf's and
        -- 'Collision's read from left to right should be the same (modulo the
        -- order of elements in 'Collision').

        go (Leaf k1 l1 : tl1) (Leaf k2 l2 : tl2)
          | k1 == k2 && l1 == l2
          = go tl1 tl2
        go (Collision k1 ary1 : tl1) (Collision k2 ary2 : tl2)
          | k1 == k2 && A.length ary1 == A.length ary2 &&
            L.null (A.toList ary1 L.\\ A.toList ary2)
          = go tl1 tl2
        go [] [] = True
        go _  _  = False

        toList' (BitmapIndexed _ ary) a = A.foldr toList' a ary
        toList' (Full ary)            a = A.foldr toList' a ary
        toList' l@(Leaf _ _)          a = l : a
        toList' c@(Collision _ _)     a = c : a
        toList' Empty                 a = a

    -- Helper function to detect 'Leaf's and 'Collision's.
    isLeafOrCollision :: HashMap k v -> Bool
    isLeafOrCollision (Leaf _ _)      = True
    isLeafOrCollision (Collision _ _) = True
    isLeafOrCollision _               = False

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

    -- | /O(1)/ Construct an empty map.
    empty :: HashMap k v
    empty = Empty

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

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

    -- | /O(1)/ Return 'True' if this map is empty, 'False' otherwise.
    null :: HashMap k v -> Bool
    null Empty = True
    null _   = False

    -- | /O(n)/ Return the number of key-value mappings in this map.
    size :: HashMap k v -> Int
    size t = go t 0
      where
        go Empty                !n = n
        go (Leaf _ _)            n = n + 1
        go (BitmapIndexed _ ary) n = A.foldl' (flip go) n ary
        go (Full ary)            n = A.foldl' (flip go) n ary
        go (Collision _ ary)     n = n + A.length ary

    -- | /O(log n)/ Return 'True' if the specified key is present in the
    -- map, 'False' otherwise.
    member :: (Eq k, Hashable k) => k -> HashMap k a -> Bool
    member k m = case lookup k m of
        Nothing -> False
        Just _  -> True
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINEABLE member #-}
    #endif

    -- | /O(log n)/ Return the value to which the specified key is mapped,
    -- or 'Nothing' if this map contains no mapping for the key.
    lookup :: (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
    lookup k0 = go h0 k0 0
      where
        h0 = hash k0
        go !_ !_ !_ Empty = Nothing
        go h k _ (Leaf hx (L kx x))
            | h == hx && k == kx = Just x  -- TODO: Split test in two
            | otherwise          = Nothing
        go h k s (BitmapIndexed b v)
            | b .&. m == 0 = Nothing
            | otherwise    = go h k (s+bitsPerSubkey) (A.index v (sparseIndex b m))
          where m = mask h s
        go h k s (Full v) = go h k (s+bitsPerSubkey) (A.index v (index h s))
        go h k _ (Collision hx v)
            | h == hx   = lookupInArray k v
            | otherwise = Nothing
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE lookup #-}
    #endif

    -- | /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 = case lookup k t of
        Just v -> v
        _      -> def
    {-# INLINABLE lookupDefault #-}

    -- | /O(log n)/ Return the value to which the specified key is mapped.
    -- Calls 'error' if this map contains no mapping for the key.
    (!) :: (Eq k, Hashable k) => HashMap k v -> k -> v
    (!) m k = case lookup k m of
        Just v  -> v
        Nothing -> error "Data.HashMap.Base.(!): key not found"
    {-# INLINABLE (!) #-}

    infixl 9 !

    -- | Create a 'Collision' value with two 'Leaf' values.
    collision :: Hash -> Leaf k v -> Leaf k v -> HashMap k v
    collision h e1 e2 =
        let v = A.run $ do mary <- A.new 2 e1
                           A.write mary 1 e2
                           return mary
        in Collision h v
    {-# INLINE collision #-}

    -- | Create a 'BitmapIndexed' or 'Full' node.
    bitmapIndexedOrFull :: Bitmap -> A.Array (HashMap k v) -> HashMap k v
    bitmapIndexedOrFull b ary
        | b == fullNodeMask = Full ary
        | otherwise         = BitmapIndexed b ary
    {-# INLINE bitmapIndexedOrFull #-}

    -- | /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 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 t@(Leaf hy l@(L ky y))
            | hy == h = if ky == k
                        then if x `ptrEq` y
                             then t
                             else Leaf h (L k x)
                        else collision h l (L k x)
            | otherwise = runST (two s h k x hy ky y)
        go h k x s t@(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
                in if st' `ptrEq` st
                   then t
                   else BitmapIndexed b (A.update ary i st')
          where m = mask h s
                i = sparseIndex b m
        go h k x s t@(Full ary) =
            let !st  = A.index ary i
                !st' = go h k x (s+bitsPerSubkey) st
            in if st' `ptrEq` st
                then t
                else Full (update16 ary i st')
          where i = index h s
        go h k x s t@(Collision hy v)
            | h == hy   = Collision h (updateOrSnocWith const k x v)
            | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE insert #-}
    #endif

    -- | In-place update version of insert
    unsafeInsert :: (Eq k, Hashable k) => k -> v -> HashMap k v -> HashMap k v
    unsafeInsert 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 t@(Leaf hy l@(L ky y))
            | hy == h = if ky == k
                        then if x `ptrEq` y
                             then return t
                             else return $! Leaf h (L k x)
                        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 const k x v)
            | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE unsafeInsert #-}
    #endif

    -- | Create a map from two key-value pairs which hashes don't collide.
    two :: Shift -> Hash -> k -> v -> Hash -> k -> v -> ST s (HashMap k v)
    two = go
      where
        go s h1 k1 v1 h2 k2 v2
            | bp1 == bp2 = do
                st <- go (s+bitsPerSubkey) h1 k1 v1 h2 k2 v2
                ary <- A.singletonM st
                return $! BitmapIndexed bp1 ary
            | otherwise  = do
                mary <- A.new 2 $ Leaf h1 (L k1 v1)
                A.write mary idx2 $ Leaf h2 (L k2 v2)
                ary <- A.unsafeFreeze mary
                return $! BitmapIndexed (bp1 .|. bp2) ary
          where
            bp1  = mask h1 s
            bp2  = mask h2 s
            idx2 | index h1 s < index h2 s = 1
                 | otherwise               = 0
    {-# INLINE two #-}

    -- | /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 Leaf h (L k (f x y))
                        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 return $! Leaf h (L k (f x y))
                        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)/ Remove the mapping for the specified key from this map
    -- if present.
    delete :: (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
    delete k0 m0 = go h0 k0 0 m0
      where
        h0 = hash k0
        go !_ !_ !_ Empty = Empty
        go h k _ t@(Leaf hy (L ky _))
            | hy == h && ky == k = Empty
            | 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
                in if st' `ptrEq` st
                    then t
                    else case st' of
                    Empty | A.length ary == 1 -> Empty
                          | A.length ary == 2 ->
                              case (i, A.index ary 0, A.index ary 1) of
                              (0, _, l) | isLeafOrCollision l -> l
                              (1, l, _) | isLeafOrCollision l -> l
                              _                               -> bIndexed
                          | otherwise -> bIndexed
                        where
                          bIndexed = BitmapIndexed (b .&. complement m) (A.delete ary i)
                    l | isLeafOrCollision l && A.length ary == 1 -> l
                    _ -> BitmapIndexed b (A.update ary i st')
          where m = mask h s
                i = sparseIndex b m
        go h k s t@(Full ary) =
            let !st   = A.index ary i
                !st' = go h k (s+bitsPerSubkey) st
            in if st' `ptrEq` st
                then t
                else case st' of
                Empty ->
                    let ary' = A.delete ary i
                        bm   = fullNodeMask .&. complement (1 `unsafeShiftL` i)
                    in BitmapIndexed bm ary'
                _ -> Full (A.update ary i st')
          where i = index h s
        go h k _ t@(Collision hy v)
            | h == hy = case indexOf k v of
                Just i
                    | A.length v == 2 ->
                        if i == 0
                        then Leaf h (A.index v 1)
                        else Leaf h (A.index v 0)
                    | otherwise -> Collision h (A.delete v i)
                Nothing -> t
            | otherwise = t
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE delete #-}
    #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 = Leaf h (L k (f y))
            | 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
    -- mapping from the first will be the mapping in the result.
    union :: (Eq k, Hashable k) => HashMap k v -> HashMap k v -> HashMap k v
    union = unionWith const
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE union #-}
    #endif

    -- | /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 #-}

    -- | Strict in the result of @f@.
    unionArrayBy :: (a -> a -> a) -> Bitmap -> Bitmap -> A.Array a -> A.Array a
                 -> A.Array a
    unionArrayBy f b1 b2 ary1 ary2 = A.run $ do
        let b' = b1 .|. b2
        mary <- A.new_ (popCount b')
        -- iterate over nonzero bits of b1 .|. b2
        -- it would be nice if we could shift m by more than 1 each time
        let ba = b1 .&. b2
            go !i !i1 !i2 !m
                | m > b'        = return ()
                | b' .&. m == 0 = go i i1 i2 (m `unsafeShiftL` 1)
                | ba .&. m /= 0 = do
                    A.write mary i $! f (A.index ary1 i1) (A.index ary2 i2)
                    go (i+1) (i1+1) (i2+1) (m `unsafeShiftL` 1)
                | b1 .&. m /= 0 = do
                    A.write mary i =<< A.indexM ary1 i1
                    go (i+1) (i1+1) (i2  ) (m `unsafeShiftL` 1)
                | otherwise     = do
                    A.write mary i =<< A.indexM ary2 i2
                    go (i+1) (i1  ) (i2+1) (m `unsafeShiftL` 1)
        go 0 0 0 (b' .&. negate b') -- XXX: b' must be non-zero
        return mary
        -- TODO: For the case where b1 .&. b2 == b1, i.e. when one is a
        -- subset of the other, we could use a slightly simpler algorithm,
        -- where we copy one array, and then update.
    {-# INLINE unionArrayBy #-}

    -- TODO: Figure out the time complexity of 'unions'.

    -- | Construct a set containing all elements from a list of sets.
    unions :: (Eq k, Hashable k) => [HashMap k v] -> HashMap k v
    unions = L.foldl' union empty
    {-# INLINE unions #-}

    ------------------------------------------------------------------------
    -- * 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)) = Leaf h $ L k (f 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) -> L k (f k v)) ary
    {-# INLINE mapWithKey #-}

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

    -- TODO: We should be able to use mutation to create the new
    -- 'HashMap'.

    -- | /O(n)/ Transform this map by accumulating an Applicative result
    -- from every value.
    traverseWithKey :: Applicative f => (k -> v1 -> f v2) -> HashMap k v1
                    -> f (HashMap k v2)
    traverseWithKey f = go
      where
        go Empty                 = pure Empty
        go (Leaf h (L k v))      = Leaf h . L k <$> f k v
        go (BitmapIndexed b ary) = BitmapIndexed b <$> A.traverse go ary
        go (Full ary)            = Full <$> A.traverse go ary
        go (Collision h ary)     =
            Collision h <$> A.traverse (\ (L k v) -> L k <$> f k v) ary
    {-# INLINE traverseWithKey #-}

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

    -- | /O(n*log m)/ Difference of two maps. Return elements of the first map
    -- not existing in the second.
    difference :: (Eq k, Hashable k) => HashMap k v -> HashMap k w -> HashMap k v
    difference a b = foldlWithKey' go empty a
      where
        go m k v = case lookup k b of
                     Nothing -> insert k v m
                     _       -> m
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE difference #-}
    #endif

    -- | /O(n*log m)/ Intersection of two maps. Return elements of the first
    -- map for keys existing in the second.
    intersection :: (Eq k, Hashable k) => HashMap k v -> HashMap k w -> HashMap k v
    intersection a b = foldlWithKey' go empty a
      where
        go m k v = case lookup k b of
                     Just _ -> insert k v m
                     _      -> m
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE intersection #-}
    #endif

    -- | /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 lookup k b of
                     Just w -> insert k (f v w) m
                     _      -> m
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE intersectionWith #-}
    #endif

    ------------------------------------------------------------------------
    -- * Folds

    -- | /O(n)/ Reduce this map by applying a binary operator to all
    -- elements, using the given starting value (typically the
    -- left-identity of the operator).  Each application of the operator
    -- is evaluated before before using the result in the next
    -- application.  This function is strict in the starting value.
    foldl' :: (a -> v -> a) -> a -> HashMap k v -> a
    foldl' f = foldlWithKey' (\ z _ v -> f z v)
    {-# INLINE foldl' #-}

    -- | /O(n)/ Reduce this map by applying a binary operator to all
    -- elements, using the given starting value (typically the
    -- left-identity of the operator).  Each application of the operator
    -- is evaluated before before using the result in the next
    -- application.  This function is strict in the starting value.
    foldlWithKey' :: (a -> k -> v -> a) -> a -> HashMap k v -> a
    foldlWithKey' f = go
      where
        go !z Empty                = z
        go z (Leaf _ (L k v))      = f z k v
        go z (BitmapIndexed _ ary) = A.foldl' go z ary
        go z (Full ary)            = A.foldl' go z ary
        go z (Collision _ ary)     = A.foldl' (\ z' (L k v) -> f z' k v) z ary
    {-# INLINE foldlWithKey' #-}

    -- | /O(n)/ Reduce this map by applying a binary operator to all
    -- elements, using the given starting value (typically the
    -- right-identity of the operator).
    foldr :: (v -> a -> a) -> a -> HashMap k v -> a
    foldr f = foldrWithKey (const f)
    {-# INLINE foldr #-}

    -- | /O(n)/ Reduce this map by applying a binary operator to all
    -- elements, using the given starting value (typically the
    -- right-identity of the operator).
    foldrWithKey :: (k -> v -> a -> a) -> a -> HashMap k v -> a
    foldrWithKey f = go
      where
        go z Empty                 = z
        go z (Leaf _ (L k v))      = f k v z
        go z (BitmapIndexed _ ary) = A.foldr (flip go) z ary
        go z (Full ary)            = A.foldr (flip go) z ary
        go z (Collision _ ary)     = A.foldr (\ (L k v) z' -> f k v z') z ary
    {-# INLINE foldrWithKey #-}

    ------------------------------------------------------------------------
    -- * Filter

    -- | Create a new array of the @n@ first elements of @mary@.
    trim :: A.MArray s a -> Int -> ST s (A.Array a)
    trim mary n = do
        mary2 <- A.new_ n
        A.copyM mary 0 mary2 0 n
        A.unsafeFreeze mary2
    {-# INLINE trim #-}

    -- | /O(n)/ Filter this map by retaining only elements satisfying a
    -- predicate.
    filterWithKey :: (k -> v -> Bool) -> HashMap k v -> HashMap k v
    filterWithKey pred = go
      where
        go Empty = Empty
        go t@(Leaf _ (L k v))
            | pred k v  = t
            | otherwise = Empty
        go (BitmapIndexed b ary) = filterA ary b
        go (Full ary) = filterA ary fullNodeMask
        go (Collision h ary) = filterC ary h

        filterA ary0 b0 =
            let !n = A.length ary0
            in runST $ do
                mary <- A.new_ n
                step ary0 mary b0 0 0 1 n
          where
            step !ary !mary !b i !j !bi n
                | i >= n = case j of
                    0 -> return Empty
                    1 -> do
                        ch <- A.read mary 0
                        case ch of
                          t | isLeafOrCollision t -> return t
                          _                       -> BitmapIndexed b <$> trim mary 1
                    _ -> do
                        ary2 <- trim mary j
                        return $! if j == maxChildren
                                  then Full ary2
                                  else BitmapIndexed b ary2
                | bi .&. b == 0 = step ary mary b i j (bi `unsafeShiftL` 1) n
                | otherwise = case go (A.index ary i) of
                    Empty -> step ary mary (b .&. complement bi) (i+1) j
                             (bi `unsafeShiftL` 1) n
                    t     -> do A.write mary j t
                                step ary mary b (i+1) (j+1) (bi `unsafeShiftL` 1) n

        filterC ary0 h =
            let !n = A.length ary0
            in runST $ do
                mary <- A.new_ n
                step ary0 mary 0 0 n
          where
            step !ary !mary i !j n
                | i >= n    = case j of
                    0 -> return Empty
                    1 -> do l <- A.read mary 0
                            return $! Leaf h l
                    _ | i == j -> do ary2 <- A.unsafeFreeze mary
                                     return $! Collision h ary2
                      | otherwise -> do ary2 <- trim mary j
                                        return $! Collision h ary2
                | pred k v  = A.write mary j el >> step ary mary (i+1) (j+1) n
                | otherwise = step ary mary (i+1) j n
              where el@(L k v) = A.index ary i
    {-# INLINE filterWithKey #-}

    -- | /O(n)/ Filter this map by retaining only elements which values
    -- satisfy a predicate.
    filter :: (v -> Bool) -> HashMap k v -> HashMap k v
    filter p = filterWithKey (\_ v -> p v)
    {-# INLINE filter #-}

    ------------------------------------------------------------------------
    -- * Conversions

    -- TODO: Improve fusion rules by modelled them after the Prelude ones
    -- on lists.

    -- | /O(n)/ Return a list of this map's keys.  The list is produced
    -- lazily.
    keys :: HashMap k v -> [k]
    keys = L.map fst . toList
    {-# INLINE keys #-}

    -- | /O(n)/ Return a list of this map's values.  The list is produced
    -- lazily.
    elems :: HashMap k v -> [v]
    elems = L.map snd . toList
    {-# INLINE elems #-}

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

    -- | /O(n)/ Return a list of this map's elements.  The list is
    -- produced lazily.
    toList :: HashMap k v -> [(k, v)]
    #if defined(__GLASGOW_HASKELL__)
    toList t = build (\ c z -> foldrWithKey (curry c) z t)
    #else
    toList = foldrWithKey (\ k v xs -> (k, v) : xs) []
    #endif
    {-# INLINE toList #-}

    -- | /O(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) -> 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
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINE fromListWith #-}
    #endif

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

    -- | /O(n)/ Lookup the value associated with the given key in this
    -- array.  Returns 'Nothing' if the key wasn't found.
    lookupInArray :: Eq k => k -> A.Array (Leaf k v) -> Maybe v
    lookupInArray k0 ary0 = go k0 ary0 0 (A.length ary0)
      where
        go !k !ary !i !n
            | i >= n    = Nothing
            | otherwise = case A.index ary i of
                (L kx v)
                    | k == kx   -> Just v
                    | otherwise -> go k ary (i+1) n
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE lookupInArray #-}
    #endif

    -- | /O(n)/ Lookup the value associated with the given key in this
    -- array.  Returns 'Nothing' if the key wasn't found.
    indexOf :: Eq k => k -> A.Array (Leaf k v) -> Maybe Int
    indexOf k0 ary0 = go k0 ary0 0 (A.length ary0)
      where
        go !k !ary !i !n
            | i >= n    = Nothing
            | otherwise = case A.index ary i of
                (L kx _)
                    | k == kx   -> Just i
                    | otherwise -> go k ary (i+1) n
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE indexOf #-}
    #endif

    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   -> A.update ary i (L k (f y))
                         | 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   -> A.update ary i (L k (f v y))
                         | otherwise -> go k v ary (i+1) n
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE updateOrSnocWith #-}
    #endif

    updateOrConcatWith :: Eq k => (v -> v -> v) -> A.Array (Leaf k v) -> A.Array (Leaf k v) -> A.Array (Leaf k v)
    updateOrConcatWith f ary1 ary2 = A.run $ do
        -- first: look up the position of each element of ary2 in ary1
        let indices = A.map (\(L k _) -> indexOf k ary1) ary2
        -- that tells us how large the overlap is:
        -- count number of Nothing constructors
        let nOnly2 = A.foldl' (\n -> maybe (n+1) (const n)) 0 indices
        let n1 = A.length ary1
        let n2 = A.length ary2
        -- copy over all elements from ary1
        mary <- A.new_ (n1 + nOnly2)
        A.copy ary1 0 mary 0 n1
        -- append or update all elements from ary2
        let go !iEnd !i2
              | i2 >= n2 = return ()
              | otherwise = case A.index indices i2 of
                   Just i1 -> do -- key occurs in both arrays, store combination in position i1
                                 L k v1 <- A.indexM ary1 i1
                                 L _ v2 <- A.indexM ary2 i2
                                 A.write mary i1 (L k (f v1 v2))
                                 go iEnd (i2+1)
                   Nothing -> do -- key is only in ary2, append to end
                                 A.write mary iEnd =<< A.indexM ary2 i2
                                 go (iEnd+1) (i2+1)
        go n1 0
        return mary
    #if __GLASGOW_HASKELL__ >= 700
    {-# INLINABLE updateOrConcatWith #-}
    #endif

    ------------------------------------------------------------------------
    -- Manually unrolled loops

    -- | /O(n)/ Update the element at the given position in this array.
    update16 :: A.Array e -> Int -> e -> A.Array e
    update16 ary idx b = runST (update16M ary idx b)
    {-# INLINE update16 #-}

    -- | /O(n)/ Update the element at the given position in this array.
    update16M :: A.Array e -> Int -> e -> ST s (A.Array e)
    update16M ary idx b = do
        mary <- clone16 ary
        A.write mary idx b
        A.unsafeFreeze mary
    {-# INLINE update16M #-}

    -- | /O(n)/ Update the element at the given position in this array, by applying a function to it.
    update16With' :: A.Array e -> Int -> (e -> e) -> A.Array e
    update16With' ary idx f = update16 ary idx $! f (A.index ary idx)
    {-# INLINE update16With' #-}

    -- | Unsafely clone an array of 16 elements.  The length of the input
    -- array is not checked.
    clone16 :: A.Array e -> ST s (A.MArray s e)
    clone16 ary =
    #if __GLASGOW_HASKELL__ >= 702
        A.thaw ary 0 16
    #else
        do mary <- A.new_ 16
           A.indexM ary 0 >>= A.write mary 0
           A.indexM ary 1 >>= A.write mary 1
           A.indexM ary 2 >>= A.write mary 2
           A.indexM ary 3 >>= A.write mary 3
           A.indexM ary 4 >>= A.write mary 4
           A.indexM ary 5 >>= A.write mary 5
           A.indexM ary 6 >>= A.write mary 6
           A.indexM ary 7 >>= A.write mary 7
           A.indexM ary 8 >>= A.write mary 8
           A.indexM ary 9 >>= A.write mary 9
           A.indexM ary 10 >>= A.write mary 10
           A.indexM ary 11 >>= A.write mary 11
           A.indexM ary 12 >>= A.write mary 12
           A.indexM ary 13 >>= A.write mary 13
           A.indexM ary 14 >>= A.write mary 14
           A.indexM ary 15 >>= A.write mary 15
           return mary
    #endif

    ------------------------------------------------------------------------
    -- Bit twiddling

    bitsPerSubkey :: Int
    bitsPerSubkey = 4

    maxChildren :: Int
    maxChildren = fromIntegral $ 1 `unsafeShiftL` bitsPerSubkey

    subkeyMask :: Bitmap
    subkeyMask = 1 `unsafeShiftL` bitsPerSubkey - 1

    sparseIndex :: Bitmap -> Bitmap -> Int
    sparseIndex b m = popCount (b .&. (m - 1))

    mask :: Word -> Shift -> Bitmap
    mask w s = 1 `unsafeShiftL` index w s
    {-# INLINE mask #-}

    -- | Mask out the 'bitsPerSubkey' bits used for indexing at this level
    -- of the tree.
    index :: Hash -> Shift -> Int
    index w s = fromIntegral $ (unsafeShiftR w s) .&. subkeyMask
    {-# INLINE index #-}

    -- | A bitmask with the 'bitsPerSubkey' least significant bits set.
    fullNodeMask :: Bitmap
    fullNodeMask = complement (complement 0 `unsafeShiftL` maxChildren)
    {-# INLINE fullNodeMask #-}

    -- | Check if two the two arguments are the same value.  N.B. This
    -- function might give false negatives (due to GC moving objects.)
    ptrEq :: a -> a -> Bool
    #if defined(__GLASGOW_HASKELL__)
    ptrEq x y = reallyUnsafePtrEquality# x y ==# 1#
    #else
    ptrEq _ _ = False
    #endif
    {-# INLINE ptrEq #-}