{-# LANGUAGE UnboxedTuples #-}

module Data.TrieMap (
	-- * Map type
	TKey,
	TMap,
	-- * Location type
	TLocation,
	-- ** Components
	key,
	before,
	after,
	-- ** Locations in maps
	search,
	index,
	minLocation,
	maxLocation,
	-- ** Building maps
	assign,
	clear,
	-- * Operators
	(!),
	(\\),
	-- * Query
	null,
	size,
	member,
	notMember,
	lookup,
	findWithDefault,
	-- * Construction
	empty,
	singleton,
	-- ** Insertion
	insert,
	insertWith,
	insertWithKey,
	insertLookupWithKey,
	-- ** Delete/Update
	delete,
	adjust,
	adjustWithKey,
	update,
	updateWithKey,
	alter,
	-- * Combine
	-- ** Union
	union,
	unionWith,
	unionWithKey,
	unionMaybeWith,
	unionMaybeWithKey,
	symmetricDifference,
	-- ** Difference
	difference,
	differenceWith,
	differenceWithKey,
	-- ** Intersection
	intersection,
	intersectionWith,
	intersectionWithKey,
	intersectionMaybeWith,
	intersectionMaybeWithKey,
	-- * Traversal
	-- ** Map
	map,
	mapWithKey,
	mapKeys,
	mapKeysWith,
	mapKeysMonotonic,
	-- ** Traverse
	traverseWithKey,
	-- ** Fold
-- 	fold,
	foldrWithKey,
	foldlWithKey,
	-- * Conversion
	elems,
	keys,
	keysSet,
	assocs,
	-- ** Lists
	fromList,
	fromListWith,
	fromListWithKey,
	-- ** Ordered lists
	fromAscList,
	fromAscListWith,
	fromAscListWithKey,
	fromDistinctAscList,
	-- * Filter
	filter,
	filterWithKey,
	partition,
	partitionWithKey,
	mapMaybe,
	mapMaybeWithKey,
	mapEither,
	mapEitherWithKey,
	split,
	splitLookup,
	-- * Submap
	isSubmapOf,
	isSubmapOfBy,
	-- * Indexed
	lookupIndex,
	findIndex,
	elemAt,
	updateAt,
	deleteAt,
	-- * Min/Max
	findMin,
	findMax,
	deleteMin,
	deleteMax,
	deleteFindMin,
	deleteFindMax,
	updateMin,
	updateMax,
	updateMinWithKey,
	updateMaxWithKey,
	minView,
	maxView,
	minViewWithKey,
	maxViewWithKey
	) where

import Control.Monad.Ends

import Data.TrieMap.Class
import Data.TrieMap.Class.Instances()
import Data.TrieMap.TrieKey
import Data.TrieMap.Representation
import Data.TrieMap.Representation.Instances ()
import Data.TrieMap.Sized
import Data.TrieMap.Utils

import Control.Applicative hiding (empty)
import Control.Monad
import qualified Data.Foldable as F
import Data.Maybe hiding (mapMaybe)
import Data.Monoid(Monoid(..))

import GHC.Exts (build)

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

instance (Show k, Show a, TKey k) => Show (TMap k a) where
	show m = "fromList " ++ show (assocs m)

instance (Eq k, TKey k, Eq a) => Eq (TMap k a) where
	m1 == m2 = assocs m1 == assocs m2

instance (Ord k, TKey k, Ord a) => Ord (TMap k a) where
	m1 `compare` m2 = assocs m1 `compare` assocs m2

instance TKey k => Monoid (TMap k a) where
	mempty = empty
	mappend = union

-- | A 'TLocation' represents a 'TMap' with a \"hole\" at a particular key position.
-- 
-- 'TLocation's are used for element-wise operations on maps (insertion, deletion and update) in a two-stage process:
-- 
-- 1. A 'TLocation' (and the value at that position, if any) is obtained from a 'TMap' by searching or indexing.
-- 2. A new 'TMap' is made from a 'TLocation' by either filling the hole with a value ('assign') or erasing it ('clear').
data TLocation k a = TLoc k (Hole (Rep k) (Assoc k a))

{-# INLINE empty #-}
-- | /O(1)/. The empty map.
empty :: TKey k => TMap k a
empty = TMap emptyM

-- | /O(1)/. A map with a single element.
{-# INLINE singleton #-}
singleton :: TKey k => k -> a -> TMap k a
singleton k a = TMap (singletonM (toRep k) (Assoc k a))

-- | /O(1)/. Is the map empty?
{-# INLINE null #-}
null :: TKey k => TMap k a -> Bool
null (TMap m) = nullM m

-- | Lookup the value at a key in the map.
-- 
-- The function will return the corresponding value as @('Just' value)@, or 'Nothing' if the key isn't in the map.
{-# INLINE lookup #-}
lookup :: TKey k => k -> TMap k a -> Maybe a
lookup k (TMap m) = option (lookupM (toRep k) m) Nothing (Just . getValue)

-- | The expression @('findWithDefault' def k map)@ returns the value at key @k@ or returns default value @def@
-- when the key is not in the map.
{-# INLINE findWithDefault #-}
findWithDefault :: TKey k => a -> k -> TMap k a -> a
findWithDefault a = fromMaybe a .: lookup

-- | Find the value at a key. Calls 'error' when the element can not be found.
{-# INLINE (!) #-}
(!) :: TKey k => TMap k a -> k -> a
m ! k = fromMaybe (error "Element not found") (lookup k m)

-- | 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 a 'TMap'. In short:
-- @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
{-# INLINE alter #-}
alter :: TKey k => (Maybe a -> Maybe a) -> k -> TMap k a -> TMap k a
alter f k (TMap m) = TMap $ searchMC (toRep k) m nomatch match where
  nomatch hole = case f Nothing of
      Nothing	-> m
      Just a'	-> assignM (Assoc k a') hole
  match (Assoc _ a) hole = fillHoleM (Assoc k <$> f (Just a)) hole

-- | Insert a new key and value in the map.
-- If the key is already present in the map, the associated value is
-- replaced with the supplied value. 'insert' is equivalent to
-- @'insertWith' 'const'@.
--
-- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')]
-- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')]
-- > insert 5 'x' empty                         == singleton 5 'x'
{-# INLINE insert #-}
insert :: TKey k => k -> a -> TMap k a -> TMap k a
insert = insertWith const

-- | Insert with a function, combining new value and old value.
-- @'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 the pair @(key, f new_value old_value)@.
--
-- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")]
-- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
-- > insertWith (++) 5 "xxx" empty                         == singleton 5 "xxx"
{-# INLINE insertWith #-}
insertWith :: TKey k => (a -> a -> a) -> k -> a -> TMap k a -> TMap k a
insertWith = insertWithKey . const

-- | Insert with a function, combining key, new value and old value.
-- @'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 the pair @(key,f key new_value old_value)@.
-- Note that the key passed to f is the same key passed to 'insertWithKey'.
--
-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
-- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")]
-- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")]
-- > insertWithKey f 5 "xxx" empty                         == singleton 5 "xxx"
{-# INLINE insertWithKey #-}
insertWithKey :: TKey k => (k -> a -> a -> a) -> k -> a -> TMap k a -> TMap k a
insertWithKey f k a (TMap m) =
  TMap (insertWithM (\ (Assoc _ a0) -> Assoc k (f k a a0)) (toRep k) (Assoc k a) m)

-- | Combines insert operation with old value retrieval.
-- 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@).
--
-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value
-- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")])
-- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing,  fromList [(3, "b"), (5, "a"), (7, "xxx")])
-- > insertLookupWithKey f 5 "xxx" empty                         == (Nothing,  singleton 5 "xxx")
{-# INLINE insertLookupWithKey #-}
insertLookupWithKey :: TKey k => (k -> a -> a -> a) -> k -> a -> TMap k a -> (Maybe a, TMap k a)
insertLookupWithKey f k a m = case search k m of
	(a', hole)	-> (a', assign (maybe a (f k a) a') hole)

-- | 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 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
-- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
-- > delete 5 empty                         == empty
{-# INLINE delete #-}
delete :: TKey k => k -> TMap k a -> TMap k a
delete k m = case search k m of
	(Nothing, _)	-> m
	(Just{}, hole)	-> clear hole

-- | Update a value at a specific key with the result of the provided function.
-- When the key is not a member of the map, the original map is returned.
--
-- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
-- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
-- > adjust ("new " ++) 7 empty                         == empty
{-# INLINE adjust #-}
adjust :: TKey k => (a -> a) -> k -> TMap k a -> TMap k a
adjust = adjustWithKey . const

-- | Adjust a value at a specific key. When the key is not
-- a member of the map, the original map is returned.
--
-- > let f key x = (show key) ++ ":new " ++ x
-- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
-- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
-- > adjustWithKey f 7 empty                         == empty
{-# INLINE adjustWithKey #-}
adjustWithKey :: TKey k => (k -> a -> a) -> k -> TMap k a -> TMap k a
adjustWithKey f k m = case search k m of
	(Nothing, _)	-> m
	(Just a, hole)	-> assign (f k a) hole

-- | 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@.
--
-- > let f x = if x == "a" then Just "new a" else Nothing
-- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")]
-- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
-- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
{-# INLINE update #-}
update :: TKey k => (a -> Maybe a) -> k -> TMap k a -> TMap k a
update f = updateWithKey (const f)

-- | The expression (@'updateWithKey' 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@.
--
-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing
-- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")]
-- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")]
-- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
{-# INLINE updateWithKey #-}
updateWithKey :: TKey k => (k -> a -> Maybe a) -> k -> TMap k a -> TMap k a
updateWithKey f k m = case search k m of
	(Nothing, _)	-> m
	(Just a, hole)	-> fillHole (f k a) hole

-- | Post-order fold.  The function will be applied from the lowest
-- value to the highest.
{-# INLINE foldrWithKey #-}
foldrWithKey :: TKey k => (k -> a -> b -> b) -> b -> TMap k a -> b
foldrWithKey f z (TMap m) = F.foldr (\ (Assoc k a) -> f k a) z m

-- | Pre-order fold.  The function will be applied from the highest
-- value to the lowest.
{-# INLINE foldlWithKey #-}
foldlWithKey :: TKey k => (b -> k -> a -> b) -> b -> TMap k a -> b
foldlWithKey f z (TMap m) = F.foldl (\ z (Assoc k a) -> f z k a) z m

-- | Map each key\/element pair to an action, evaluate these actions from left to right, and collect the results.
{-# INLINE traverseWithKey #-}
traverseWithKey :: (TKey k, Applicative f) => (k -> a -> f b) -> TMap k a -> f (TMap k b)
traverseWithKey f (TMap m) = TMap <$> traverseM (\ (Assoc k a) -> Assoc k <$> f k a) m

-- | Map a function over all values in the map.
--
-- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")]
{-# INLINE map #-}
map :: TKey k => (a -> b) -> TMap k a -> TMap k b
map = fmap

-- | Map a function over all values in the map.
--
-- > let f key x = (show key) ++ ":" ++ x
-- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")]
{-# INLINEABLE mapWithKey #-}
mapWithKey :: TKey k => (k -> a -> b) -> TMap k a -> TMap k b
mapWithKey f (TMap m) = TMap (fmapM (\ (Assoc k a) -> Assoc k (f k a)) m)

-- |
-- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@.
-- 
-- The size of the result may be smaller if @f@ maps two or more distinct
-- keys to the same new key.  In this case the value at the smallest of
-- these keys is retained.
--
-- > mapKeys (+ 1) (fromList [(5,"a"), (3,"b")])                        == fromList [(4, "b"), (6, "a")]
-- > mapKeys (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "c"
-- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c"
{-# INLINE mapKeys #-}
mapKeys :: (TKey k, TKey k') => (k -> k') -> TMap k a -> TMap k' a
mapKeys = mapKeysWith const

-- |
-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
-- 
-- The size of the result may be smaller if @f@ maps two or more distinct
-- keys to the same new key.  In this case the associated values will be
-- combined using @c@.
--
-- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab"
-- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab"
{-# INLINE mapKeysWith #-}
mapKeysWith :: (TKey k, TKey k') => (a -> a -> a) -> (k -> k') -> TMap k a -> TMap k' a
mapKeysWith g f m = fromListWith g [(f k, a) | (k, a) <- assocs m]

-- |
-- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@
-- is strictly monotonic.
-- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@.
-- /The precondition is not checked./
-- Semi-formally, we have:
-- 
-- > and [x < y ==> f x < f y | x <- ls, y <- ls] 
-- >                     ==> mapKeysMonotonic f s == mapKeys f s
-- >     where ls = keys s
--
-- This means that @f@ maps distinct original keys to distinct resulting keys.
-- This function has better performance than 'mapKeys'.
--
-- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")]
{-# INLINE mapKeysMonotonic #-}
mapKeysMonotonic :: (TKey k, TKey k') => (k -> k') -> TMap k a -> TMap k' a
mapKeysMonotonic f m = fromDistinctAscList [(f k, a) | (k, a) <- assocs m]

-- |
-- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@. 
-- It prefers @t1@ when duplicate keys are encountered,
-- i.e. (@'union' == 'unionWith' 'const'@).
-- The implementation uses the efficient /hedge-union/ algorithm.
-- Hedge-union is more efficient on (bigset \``union`\` smallset).
--
-- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")]
{-# INLINE union #-}
union :: TKey k => TMap k a -> TMap k a -> TMap k a
union = unionWith const

-- | /O(n+m)/. Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
--
-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")]
{-# INLINE unionWith #-}
unionWith :: TKey k => (a -> a -> a) -> TMap k a -> TMap k a -> TMap k a
unionWith = unionWithKey . const

-- Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
-- Hedge-union is more efficient on (bigset \``union`\` smallset).
--
-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value
-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
{-# INLINE unionWithKey #-}
unionWithKey :: TKey k => (k -> a -> a -> a) -> TMap k a -> TMap k a -> TMap k a
unionWithKey f = unionMaybeWithKey (\ k a b -> Just (f k a b))

-- | Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
{-# INLINE unionMaybeWith #-}
unionMaybeWith :: TKey k => (a -> a -> Maybe a) -> TMap k a -> TMap k a -> TMap k a
unionMaybeWith = unionMaybeWithKey . const

-- | Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
-- Hedge-union is more efficient on (bigset \``union`\` smallset).
--
-- > let f key left_value right_value = Just ((show key) ++ ":" ++ left_value ++ "|" ++ right_value)
-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
{-# INLINEABLE unionMaybeWithKey #-}
unionMaybeWithKey :: TKey k => (k -> a -> a -> Maybe a) -> TMap k a -> TMap k a -> TMap k a
unionMaybeWithKey f (TMap m1) (TMap m2) = TMap (unionM f' m1 m2) where
	f' (Assoc k a) (Assoc _ b) = Assoc k <$> f k a b

-- | 'symmetricDifference' is equivalent to @'unionMaybeWith' (\ _ _ -> Nothing)@.
{-# INLINE symmetricDifference #-}
symmetricDifference :: TKey k => TMap k a -> TMap k a -> TMap k a
symmetricDifference = unionMaybeWith (\ _ _ -> Nothing)

-- | Intersection of two maps.
-- Return data in the first map for the keys existing in both maps.
-- (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@).
--
-- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a"
{-# INLINE intersection #-}
intersection :: TKey k => TMap k a -> TMap k b -> TMap k a
intersection = intersectionWith const

-- | Intersection with a combining function.
--
-- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA"
{-# INLINE intersectionWith #-}
intersectionWith :: TKey k => (a -> b -> c) -> TMap k a -> TMap k b -> TMap k c
intersectionWith = intersectionWithKey . const

-- | Intersection with a combining function.
-- Intersection is more efficient on (bigset \``intersection`\` smallset).
--
-- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar
-- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A"
{-# INLINE intersectionWithKey #-}
intersectionWithKey :: TKey k => (k -> a -> b -> c) -> TMap k a -> TMap k b -> TMap k c
intersectionWithKey f = intersectionMaybeWithKey (\ k a b -> Just (f k a b))

-- | @'intersectionMaybeWith' f m1 m2@ is equivalent to
-- @'mapMaybe' 'id' ('intersectionWith' f m1 m2)@.
{-# INLINE intersectionMaybeWith #-}
intersectionMaybeWith :: TKey k => (a -> b -> Maybe c) -> TMap k a -> TMap k b -> TMap k c
intersectionMaybeWith = intersectionMaybeWithKey . const

-- | @'intersectionMaybeWithKey' f m1 m2@ is equivalent to
-- @'mapMaybe' 'id' ('intersectionWithKey' f m1 m2)@.
{-# INLINEABLE intersectionMaybeWithKey #-}
intersectionMaybeWithKey :: TKey k => (k -> a -> b -> Maybe c) -> TMap k a -> TMap k b -> TMap k c
intersectionMaybeWithKey f (TMap m1) (TMap m2) = TMap (isectM f' m1 m2) where
	f' (Assoc k a) (Assoc _ b) = Assoc k <$> f k a b

-- | Difference of two maps. 
-- Return elements of the first map not existing in the second map.
-- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
--
-- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b"
{-# INLINE difference #-}
difference :: TKey k => TMap k a -> TMap k b -> TMap k a
difference = differenceWith (\ _ _ -> Nothing)

-- | Same as 'difference'.
(\\) :: TKey k => TMap k a -> TMap k b -> TMap k a
(\\) = difference

-- | Difference with a combining function. 
-- When two equal keys are
-- encountered, the combining function is applied to the values of these keys.
-- 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@. 
-- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
--
-- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing
-- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")])
-- >     == singleton 3 "b:B"
{-# INLINE differenceWith #-}
differenceWith :: TKey k => (a -> b -> Maybe a) -> TMap k a -> TMap k b -> TMap k a
differenceWith = differenceWithKey . const

-- | 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@. 
-- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
--
-- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing
-- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")])
-- >     == singleton 3 "3:b|B"
{-# INLINEABLE differenceWithKey #-}
differenceWithKey :: TKey k => (k -> a -> b -> Maybe a) -> TMap k a -> TMap k b -> TMap k a
differenceWithKey f (TMap m1) (TMap m2) = TMap (diffM f' m1 m2) where
	f' (Assoc k a) (Assoc _ b) = Assoc k <$> f k a b

-- | Retrieves the value associated with minimal key of the
-- map, and the map stripped of that element, or 'Nothing' if passed an
-- empty map.
--
-- > minView (fromList [(5,"a"), (3,"b")]) == Just ("b", singleton 5 "a")
-- > minView empty == Nothing
{-# INLINE minView #-}
minView :: TKey k => TMap k a -> Maybe (a, TMap k a)
minView = fmap (fmap after) . minLocation

-- | Retrieves the value associated with maximal key of the
-- map, and the map stripped of that element, or 'Nothing' if passed an
--
-- > maxView (fromList [(5,"a"), (3,"b")]) == Just ("a", singleton 3 "b")
-- > maxView empty == Nothing
{-# INLINE maxView #-}
maxView :: TKey k => TMap k a -> Maybe (a, TMap k a)
maxView = fmap (fmap before) . maxLocation

-- | The minimal key of the map. Calls 'error' if the map is empty.
--
-- > findMin (fromList [(5,"a"), (3,"b")]) == (3,"b")
-- > findMin empty                            Error: empty map has no minimal element
{-# INLINE findMin #-}
findMin :: TKey k => TMap k a -> (k, a)
findMin = maybe (error "empty map has no minimal element") fst . minViewWithKey

-- | The maximal key of the map. Calls 'error' if the map is empty.
--
-- > findMax (fromList [(5,"a"), (3,"b")]) == (5,"a")
-- > findMax empty                            Error: empty map has no maximal element
{-# INLINE findMax #-}
findMax :: TKey k => TMap k a -> (k, a)
findMax = maybe (error "empty map has no maximal element") fst . maxViewWithKey

-- | Delete the minimal key. Returns an empty map if the map is empty.
--
-- > deleteMin (fromList [(5,"a"), (3,"b"), (7,"c")]) == fromList [(5,"a"), (7,"c")]
-- > deleteMin empty == empty
{-# INLINE deleteMin #-}
deleteMin :: TKey k => TMap k a -> TMap k a
deleteMin m = maybe m snd (minViewWithKey m)

-- | Delete the maximal key. Returns an empty map if the map is empty.
--
-- > deleteMax (fromList [(5,"a"), (3,"b"), (7,"c")]) == fromList [(3,"b"), (5,"a")]
-- > deleteMax empty == empty
{-# INLINE deleteMax #-}
deleteMax :: TKey k => TMap k a -> TMap k a
deleteMax m = maybe m snd (maxViewWithKey m)

-- | Update the value at the minimal key.
--
-- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")]
-- > updateMin (\ _ -> Nothing)         (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
{-# INLINE updateMin #-}
updateMin :: TKey k => (a -> Maybe a) -> TMap k a -> TMap k a
updateMin = updateMinWithKey . const

-- | Update the value at the maximal key.
--
-- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")]
-- > updateMax (\ _ -> Nothing)         (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
{-# INLINE updateMax #-}
updateMax :: TKey k => (a -> Maybe a) -> TMap k a -> TMap k a
updateMax = updateMaxWithKey . const

{-# INLINE updateHelper #-}
updateHelper :: (TKey k, Functor m, MonadPlus m) =>
  (k -> a -> Maybe a) -> TMap k a -> m (Maybe (Assoc k a), Hole (Rep k) (Assoc k a))
updateHelper f (TMap m) = fmap (\ (Assoc k a, loc) -> (Assoc k <$> f k a, loc)) (extractHoleM m)

-- | Update the value at the minimal key.
--
-- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")]
-- > updateMinWithKey (\ _ _ -> Nothing)                     (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
{-# INLINEABLE updateMinWithKey #-}
updateMinWithKey :: TKey k => (k -> a -> Maybe a) -> TMap k a -> TMap k a
updateMinWithKey f m = fromMaybe m $ do
	(a, loc) <- getFirst $ updateHelper f m
	return (TMap (afterMM a loc))

-- | Update the value at the maximal key.
--
-- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")]
-- > updateMaxWithKey (\ _ _ -> Nothing)                     (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
{-# INLINEABLE updateMaxWithKey #-}
updateMaxWithKey :: TKey k => (k -> a -> Maybe a) -> TMap k a -> TMap k a
updateMaxWithKey f m = fromMaybe m $ do
	(a, loc) <- getLast $ updateHelper f m
	return (TMap (beforeMM a loc))

-- | Delete and find the minimal element.
--
-- > deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((3,"b"), fromList[(5,"a"), (10,"c")]) 
-- > deleteFindMin                                            Error: can not return the minimal element of an empty map
{-# INLINEABLE deleteFindMin #-}
deleteFindMin :: TKey k => TMap k a -> ((k, a), TMap k a)
deleteFindMin m = fromMaybe (error "Cannot return the minimal element of an empty map") (minViewWithKey m)

-- | Delete and find the minimal element.
--
-- > deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((3,"b"), fromList[(5,"a"), (10,"c")]) 
-- > deleteFindMin                                            Error: can not return the minimal element of an empty map
{-# INLINEABLE deleteFindMax #-}
deleteFindMax :: TKey k => TMap k a -> ((k, a), TMap k a)
deleteFindMax m = fromMaybe (error "Cannot return the maximal element of an empty map") (maxViewWithKey m)

-- | Retrieves the minimal (key,value) pair of the map, and
-- the map stripped of that element, or 'Nothing' if passed an empty map.
--
-- > minViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((3,"b"), singleton 5 "a")
-- > minViewWithKey empty == Nothing
minViewWithKey :: TKey k => TMap k a -> Maybe ((k, a), TMap k a)
{-# INLINE minViewWithKey #-}
minViewWithKey m = do
	(a, loc) <- minLocation m
	return ((key loc, a), after loc)

-- | /O(log n)/. Retrieves the maximal (key,value) pair of the map, and
-- the map stripped of that element, or 'Nothing' if passed an empty map.
--
-- > maxViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((5,"a"), singleton 3 "b")
-- > maxViewWithKey empty == Nothing
{-# INLINE maxViewWithKey #-}
maxViewWithKey :: TKey k => TMap k a -> Maybe ((k, a), TMap k a)
maxViewWithKey m = do
	(a, loc) <- maxLocation m
	return ((key loc, a), before loc)

-- |
-- Return all elements of the map in the ascending order of their keys.
--
-- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"]
-- > elems empty == []
{-# INLINE elems #-}
elems :: TKey k => TMap k a -> [a]
elems m = build (\ c n -> foldrWithKey (\ _ a -> c a) n m)

-- | Return all keys of the map in ascending order.
--
-- > keys (fromList [(5,"a"), (3,"b")]) == [3,5]
-- > keys empty == []
{-# INLINE keys #-}
keys :: TKey k => TMap k a -> [k]
keys m = build (\ c n -> foldrWithKey (\ k _ -> c k) n m)

-- | Return all key\/value pairs in the map in ascending key order.
--
-- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
-- > assocs empty == []
{-# INLINE assocs #-}
assocs :: TKey k => TMap k a -> [(k, a)]
assocs m = build (\ c n -> foldrWithKey (curry c) n m)

-- | Map values and separate the 'Left' and 'Right' results.
--
-- > let f a = if a < "c" then Left a else Right a
-- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
-- >     == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")])
-- >
-- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
-- >     == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
{-# INLINE mapEither #-}
mapEither :: TKey k => (a -> Either b c) -> TMap k a -> (TMap k b, TMap k c)
mapEither = mapEitherWithKey . const

-- | Map keys\/values and separate the 'Left' and 'Right' results.
--
-- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a)
-- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
-- >     == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")])
-- >
-- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")])
-- >     == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")])
{-# INLINEABLE mapEitherWithKey #-}
mapEitherWithKey :: TKey k => (k -> a -> Either b c) -> TMap k a -> (TMap k b, TMap k c)
mapEitherWithKey f (TMap m) = case mapEitherM f' m of
	(# mL, mR #) -> (TMap mL, TMap mR) 
	where	f' (Assoc k a) = case f k a of
			Left b	-> (# Just (Assoc k b), Nothing #)
			Right c	-> (# Nothing, Just (Assoc k c) #)

-- | /O(n)/. Map values and collect the 'Just' results.
--
-- > let f x = if x == "a" then Just "new a" else Nothing
-- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a"
{-# INLINE mapMaybe #-}
mapMaybe :: TKey k => (a -> Maybe b) -> TMap k a -> TMap k b
mapMaybe = mapMaybeWithKey . const

-- | Map keys\/values and collect the 'Just' results.
--
-- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing
-- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3"
{-# INLINEABLE mapMaybeWithKey #-}
mapMaybeWithKey :: TKey k => (k -> a -> Maybe b) -> TMap k a -> TMap k b
mapMaybeWithKey f (TMap m) = TMap (mapMaybeM (\ (Assoc k a) -> Assoc k <$> f k a) m)

-- | Partition the map according to a predicate. The first
-- map contains all elements that satisfy the predicate, the second all
-- elements that fail the predicate. See also 'split'.
--
-- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
-- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
-- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
{-# INLINE partition #-}
partition :: TKey k => (a -> Bool) -> TMap k a -> (TMap k a, TMap k a)
partition = partitionWithKey . const

-- | Partition the map according to a predicate. The first
-- map contains all elements that satisfy the predicate, the second all
-- elements that fail the predicate. See also 'split'.
--
-- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
-- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty)
-- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")])
{-# INLINE partitionWithKey #-}
partitionWithKey :: TKey k => (k -> a -> Bool) -> TMap k a -> (TMap k a, TMap k a)
partitionWithKey p = mapEitherWithKey (\ k a -> (if p k a then Left else Right) a)

-- | Filter all values that satisfy the predicate.
--
-- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
-- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty
-- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty
{-# INLINE filter #-}
filter :: TKey k => (a -> Bool) -> TMap k a -> TMap k a
filter = filterWithKey . const

-- | Filter all keys\/values that satisfy the predicate.
--
-- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
{-# INLINE filterWithKey #-}
filterWithKey :: TKey k => (k -> a -> Bool) -> TMap k a -> TMap k a
filterWithKey p = mapMaybeWithKey (\ k a -> if p k a then Just a else Nothing)

-- | The expression (@'split' k map@) is a pair @(map1,map2)@ where
-- the keys in @map1@ are smaller than @k@ and the keys in @map2@ larger than @k@.
-- Any key equal to @k@ is found in neither @map1@ nor @map2@.
--
-- > split 2 (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3,"b"), (5,"a")])
-- > split 3 (fromList [(5,"a"), (3,"b")]) == (empty, singleton 5 "a")
-- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a")
-- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty)
-- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty)
{-# INLINE split #-}
split :: TKey k => k -> TMap k a -> (TMap k a, TMap k a)
split k m = case splitLookup k m of
	(mL, _, mR) -> (mL, mR)

-- | The expression (@'splitLookup' k map@) splits a map just
-- like 'split' but also returns @'lookup' k map@.
--
-- > splitLookup 2 (fromList [(5,"a"), (3,"b")]) == (empty, Nothing, fromList [(3,"b"), (5,"a")])
-- > splitLookup 3 (fromList [(5,"a"), (3,"b")]) == (empty, Just "b", singleton 5 "a")
-- > splitLookup 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Nothing, singleton 5 "a")
-- > splitLookup 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Just "a", empty)
-- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty)
{-# INLINE splitLookup #-}
splitLookup :: TKey k => k -> TMap k a -> (TMap k a, Maybe a, TMap k a)
splitLookup k m = case search k m of
	(x, hole) -> (before hole, x, after hole)

-- | 
-- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
{-# INLINE isSubmapOf #-}
isSubmapOf :: (TKey k, Eq a) => TMap k a -> TMap k a -> Bool
isSubmapOf = isSubmapOfBy (==)

{- |
 The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if
 all keys in @t1@ are in tree @t2@, and when @f@ returns 'True' when
 applied to their respective values. For example, the following 
 expressions are all 'True':
 
 > isSubmapOfBy (==) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
 > isSubmapOfBy (<=) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
 > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1),('b',2)])

 But the following are all 'False':
 
 > isSubmapOfBy (==) (fromList [('a',2)]) (fromList [('a',1),('b',2)])
 > isSubmapOfBy (<)  (fromList [('a',1)]) (fromList [('a',1),('b',2)])
 > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1)])
 
-}
{-# INLINEABLE isSubmapOfBy #-}
isSubmapOfBy :: TKey k => (a -> b -> Bool) -> TMap k a -> TMap k b -> Bool
isSubmapOfBy (<=) (TMap m1) (TMap m2) = isSubmapM (<<=) m1 m2 where
	Assoc _ a <<= Assoc _ b = a <= b

-- | Build a map from a list of key\/value pairs. See also 'fromAscList'.
-- If the list contains more than one value for the same key, the last value
-- for the key is retained.
--
-- > fromList [] == empty
-- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")]
-- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")]
{-# INLINE fromList #-}
fromList :: TKey k => [(k, a)] -> TMap k a
fromList = fromListWith const

-- | Build a map from an ascending list in linear time.
-- /The precondition (input list is ascending) is not checked./
--
-- > fromAscList [(3,"b"), (5,"a")]          == fromList [(3, "b"), (5, "a")]
-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
{-# INLINE fromAscList #-}
fromAscList :: TKey k => [(k, a)] -> TMap k a
fromAscList = fromAscListWith const

-- | Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
--
-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")]
-- > fromListWith (++) [] == empty
{-# INLINE fromListWith #-}
fromListWith :: TKey k => (a -> a -> a) -> [(k, a)] -> TMap k a
fromListWith = fromListWithKey . const

-- | Build a map from an ascending list in linear time with a combining function for equal keys.
-- /The precondition (input list is ascending) is not checked./
--
-- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
{-# INLINE fromAscListWith #-}
fromAscListWith :: TKey k => (a -> a -> a) -> [(k, a)] -> TMap k a
fromAscListWith = fromAscListWithKey . const

-- | Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
--
-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")]
-- > fromListWith (++) [] == empty
{-# INLINEABLE fromListWithKey #-}
fromListWithKey :: TKey k => (k -> a -> a -> a) -> [(k, a)] -> TMap k a
fromListWithKey f xs = TMap (fromListM f' [(toRep k, Assoc k a) | (k, a) <- xs])
	where f' (Assoc k a) (Assoc _ b) = Assoc k (f k a b)

-- | Build a map from an ascending list in linear time.
-- /The precondition (input list is ascending) is not checked./
--
-- > fromAscList [(3,"b"), (5,"a")]          == fromList [(3, "b"), (5, "a")]
-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
{-# INLINEABLE fromAscListWithKey #-}
fromAscListWithKey :: TKey k => (k -> a -> a -> a) -> [(k, a)] -> TMap k a
fromAscListWithKey f xs = TMap (fromAscListM f' [(toRep k, Assoc k a) | (k, a) <- xs])
	where f' (Assoc k a) (Assoc _ b) = Assoc k (f k a b)

-- | Build a map from an ascending list of distinct elements in linear time.
-- /The precondition is not checked./
--
-- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
{-# INLINEABLE fromDistinctAscList #-}
fromDistinctAscList :: TKey k => [(k, a)] -> TMap k a
fromDistinctAscList xs = TMap (fromDistAscListM [(toRep k, Assoc k a) | (k, a) <- xs])

-- | /O(1)/. The number of elements in the map.
--
-- > size empty                                   == 0
-- > size (singleton 1 'a')                       == 1
-- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3
{-# INLINE size #-}
size :: TKey k => TMap k a -> Int
size (TMap m) = getSize m

-- | Is the key a member of the map? See also 'notMember'.
--
-- > member 5 (fromList [(5,'a'), (3,'b')]) == True
-- > member 1 (fromList [(5,'a'), (3,'b')]) == False
{-# INLINE member #-}
member :: TKey k => k -> TMap k a -> Bool
member = isJust .: lookup

-- | Is the key not a member of the map? See also 'member'.
--
-- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False
-- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True
{-# INLINE notMember #-}
notMember :: TKey k => k -> TMap k a -> Bool
notMember = not .: member

-- | The set of all keys of the map.
--
-- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.TrieSet.fromList [3,5]
-- > keysSet empty == Data.TrieSet.empty
{-# INLINE keysSet #-}
keysSet :: TKey k => TMap k a -> TSet k
keysSet (TMap m) = TSet (fmapM (\ (Assoc k _) -> Elem k) m)

-- | /O(1)/.  The key marking the position of the \"hole\" in the map.
{-# INLINE key #-}
key :: TKey k => TLocation k a -> k
key (TLoc k _) = k

-- | @'before' loc@ is the submap with keys less than @'key' loc@.
{-# INLINE before #-}
before :: TKey k => TLocation k a -> TMap k a
before (TLoc _ hole) = TMap (beforeM hole)

-- | @'after' loc@ is the submap with keys greater than @'key' loc@.
{-# INLINE after #-}
after :: TKey k => TLocation k a -> TMap k a
after (TLoc _ hole) = TMap (afterM hole)

-- | Search the map for the given key, returning the
-- corresponding value (if any) and an updatable location for that key.
--
-- Properties:
--
-- @
-- case 'search' k m of
--     (Nothing, loc) -> 'key' loc == k && 'clear' loc == m
--     (Just v,  loc) -> 'key' loc == k && 'assign' v loc == m
-- @
--
-- @'lookup' k m == 'fst' ('search' k m)@
{-# INLINE search #-}
search :: TKey k => k -> TMap k a -> (Maybe a, TLocation k a)
search k (TMap m) = searchMC (toRep k) m nomatch match where
  nomatch hole = (Nothing, TLoc k hole)
  match (Assoc k a) hole = (Just a, TLoc k hole)

-- | Return the value and an updatable location for the
-- /i/th key in the map.  Calls 'error' if /i/ is out of range.
--
-- Properties:
--
-- @
-- 0 \<= i && i \< 'size' m ==>
--     let (v, loc) = 'index' i m in
--         'size' ('before' loc) == i && 'assign' v loc == m
-- @
--
-- @'elemAt' i m == let (v, loc) = 'index' i m in ('key' loc, v)@
{-# INLINEABLE index #-}
index :: TKey k => Int -> TMap k a -> (a, TLocation k a)
index i m
	| i < 0 || i >= size m
		= error "TrieMap.index: index out of range"
index i (TMap m) = case indexM i m of
	(# _, Assoc k a, hole #) -> (a, TLoc k hole)

{-# INLINE extract #-}
extract :: (TKey k, Functor m, MonadPlus m) => TMap k a -> m (a, TLocation k a)
extract (TMap m) = fmap (\ (Assoc k a, hole) -> (a, TLoc k hole)) (extractHoleM m)

-- | /O(log n)/. Return the value and an updatable location for the
-- least key in the map, or 'Nothing' if the map is empty.
--
-- Properties:
--
-- @
-- 'size' m > 0 ==>
--     let Just (v, loc) = 'minLocation' i m in
--         'size' (`before` loc) == 0 && 'assign' v loc == m
-- @
--
-- @'findMin' m == let Just (v, loc) = 'minLocation' i m in ('key' loc, v)@
{-# INLINEABLE minLocation #-}
minLocation :: TKey k => TMap k a -> Maybe (a, TLocation k a)
minLocation = getFirst . extract

-- | Return the value and an updatable location for the
-- greatest key in the map, or 'Nothing' if the map is empty.
--
-- Properties:
--
-- @
-- 'size' m > 0 ==>
--     let Just (v, loc) = 'maxLocation' i m in
--         'size' (`after` loc) == 0 && 'assign' v loc == m
-- @
--
-- @'findMax' m == let Just (v, loc) = 'maxLocation' i m in ('key' loc, v)@
{-# INLINEABLE maxLocation #-}
maxLocation :: TKey k => TMap k a -> Maybe (a, TLocation k a)
maxLocation = getLast . extract

-- | Return a map obtained by placing the given value
-- at the location (replacing an existing value, if any).
--
-- @'assign' v loc == 'before' loc `union` 'singleton' ('key' loc) v `union` 'after' loc@
{-# INLINE assign #-}
assign :: TKey k => a -> TLocation k a -> TMap k a
assign a (TLoc k hole) = TMap (assignM (Assoc k a) hole)

-- | Return a map obtained by erasing the location.
--
-- @'clear' loc == 'before' loc `union` 'after' loc@
{-# INLINE clear #-}
clear :: TKey k => TLocation k a -> TMap k a
clear (TLoc _ hole) = TMap (clearM hole)

{-# INLINE fillHole #-}
fillHole :: TKey k => Maybe a -> TLocation k a -> TMap k a
fillHole = maybe clear assign

-- | Return the /index/ of a key. The index is a number from
-- /0/ up to, but not including, the 'size' of the map. Calls 'error' when
-- the key is not a 'member' of the map.
--
-- > findIndex 2 (fromList [(5,"a"), (3,"b")])    Error: element is not in the map
-- > findIndex 3 (fromList [(5,"a"), (3,"b")]) == 0
-- > findIndex 5 (fromList [(5,"a"), (3,"b")]) == 1
-- > findIndex 6 (fromList [(5,"a"), (3,"b")])    Error: element is not in the map
{-# INLINEABLE findIndex #-}
findIndex :: TKey k => k -> TMap k a -> Int
findIndex k m = fromMaybe (error "TrieMap.findIndex: key is not in the map") (lookupIndex k m)

-- | Lookup the /index/ of a key. The index is a number from
-- /0/ up to, but not including, the 'size' of the map.
--
-- > lookupIndex 2 (fromList [(5,"a"), (3,"b")]) == Nothing
-- > lookupIndex 3 (fromList [(5,"a"), (3,"b")]) == Just 0
-- > lookupIndex 5 (fromList [(5,"a"), (3,"b")]) == Just 1
-- > lookupIndex 6 (fromList [(5,"a"), (3,"b")]) == Nothing
{-# INLINEABLE lookupIndex #-}
lookupIndex :: TKey k => k -> TMap k a -> Maybe Int
lookupIndex k m = case search k m of
	(Nothing, _)	-> Nothing
	(_, hole)	-> Just $ size (before hole)

-- | Retrieve an element by /index/. Calls 'error' when an
-- invalid index is used.
--
-- > elemAt 0 (fromList [(5,"a"), (3,"b")]) == (3,"b")
-- > elemAt 1 (fromList [(5,"a"), (3,"b")]) == (5, "a")
-- > elemAt 2 (fromList [(5,"a"), (3,"b")])    Error: index out of range
{-# INLINEABLE elemAt #-}
elemAt :: TKey k => Int -> TMap k a -> (k, a)
elemAt i m = case index i m of
	(a, hole) -> (key hole, a)

-- | Update the element at /index/. Calls 'error' when an
-- invalid index is used.
--
-- > updateAt (\ _ _ -> Just "x") 0    (fromList [(5,"a"), (3,"b")]) == fromList [(3, "x"), (5, "a")]
-- > updateAt (\ _ _ -> Just "x") 1    (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "x")]
-- > updateAt (\ _ _ -> Just "x") 2    (fromList [(5,"a"), (3,"b")])    Error: index out of range
-- > updateAt (\ _ _ -> Just "x") (-1) (fromList [(5,"a"), (3,"b")])    Error: index out of range
-- > updateAt (\_ _  -> Nothing)  0    (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
-- > updateAt (\_ _  -> Nothing)  1    (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
-- > updateAt (\_ _  -> Nothing)  2    (fromList [(5,"a"), (3,"b")])    Error: index out of range
-- > updateAt (\_ _  -> Nothing)  (-1) (fromList [(5,"a"), (3,"b")])    Error: index out of range
{-# INLINEABLE updateAt #-}
updateAt :: TKey k => (k -> a -> Maybe a) -> Int -> TMap k a -> TMap k a
updateAt f i m = case index i m of
	(a, hole) -> fillHole (f (key hole) a) hole

-- | Delete the element at /index/.
-- Defined as (@'deleteAt' i map = 'updateAt' (\k x -> 'Nothing') i map@).
--
-- > deleteAt 0  (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
-- > deleteAt 1  (fromList [(5,"a"), (3,"b")]) == singleton 3 "b"
-- > deleteAt 2 (fromList [(5,"a"), (3,"b")])     Error: index out of range
-- > deleteAt (-1) (fromList [(5,"a"), (3,"b")])  Error: index out of range
{-# INLINEABLE deleteAt #-}
deleteAt :: TKey k => Int -> TMap k a -> TMap k a
deleteAt i m = clear (snd (index i m))