{-# LANGUAGE TypeFamilies, FlexibleContexts, 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 Data.TrieMap.Class
import Data.TrieMap.Class.Instances()
import Data.TrieMap.TrieKey
import Data.TrieMap.Applicative
import Data.TrieMap.Rep
import Data.TrieMap.Rep.Instances ()
import Data.TrieMap.Sized

import Control.Applicative hiding (empty)
import Control.Monad
import Data.Maybe hiding (mapMaybe)
import Data.Monoid(Monoid(..), First(..), Last(..))

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').
newtype TLocation k a = TLoc (Hole (Rep k) (Elem a))

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

-- | /O(1)/. A map with a single element.
singleton :: TKey k => k -> a -> TMap k a
singleton k a = insert k a empty

-- | /O(1)/. Is the map empty?
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.
lookup :: TKey k => k -> TMap k a -> Maybe a
lookup k (TMap m) = getElem <$> lookupM (toRep k) m

-- | 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.
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.
(!) :: 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 m = case search k m of
	(Nothing, hole)	-> case f Nothing of
		Nothing	-> m
		Just a'	-> assign a' hole
	(a, hole)	-> fillHole (f 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 m = snd (insertLookupWithKey f 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"
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"
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.
foldrWithKey :: TKey k => (k -> a -> b -> b) -> b -> TMap k a -> b
foldrWithKey f z (TMap m) = foldrWithKeyM (\ k (Elem a) -> f (fromRep k) a) m z

-- | Pre-order fold.  The function will be applied from the highest
-- value to the lowest.
foldlWithKey :: TKey k => (b -> k -> a -> b) -> b -> TMap k a -> b
foldlWithKey f z (TMap m) = foldlWithKeyM (\ k z (Elem a) -> f z (fromRep k) a) m z

-- | Map each key\/element pair to an action, evaluate these actions from left to right, and collect the results.
traverseWithKey :: (TKey k, Applicative f) => (k -> a -> f b) -> TMap k a -> f (TMap k b)
traverseWithKey f (TMap m) = TMap <$> traverseWithKeyM (\ k (Elem a) -> Elem <$> f (fromRep 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 f = mapWithKey (const f)

-- | 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 (mapWithKeyM (\ k (Elem a) -> Elem (f (fromRep 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' k (Elem a) (Elem b) = Elem <$> f (fromRep k) a b

-- | 'symmetricDifference' is equivalent to @'unionMaybeWith' (\ _ _ -> Nothing)@.
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' k (Elem a) (Elem b) = Elem <$> f (fromRep 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"
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' k (Elem a) (Elem b) = Elem <$> f (fromRep 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, MonadPlus m) => (k -> a -> Maybe a) -> TMap k a -> m (Maybe (Elem a), Hole (Rep k) (Elem a))
updateHelper f (TMap m) = do
	(Elem a, loc) <- extractHoleM m
	return (Elem <$> f (fromRep (keyM loc)) a, loc)

-- | 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 (afterM 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 (afterM 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)

{-# INLINE minViewWithKey #-}
-- | 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)
minViewWithKey m = do
	(a, loc) <- minLocation m
	return ((key loc, a), after loc)

{-# INLINE maxViewWithKey #-}
-- | /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
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' k (Elem a) = case f (fromRep k) a of
			Left b	-> (# Just (Elem b), Nothing #)
			Right c	-> (# Nothing, Just (Elem 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 (\ k (Elem a) -> Elem <$> f (fromRep 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)
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
	Elem a <<= Elem 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, Elem a) | (k, a) <- xs])
	where f' k (Elem a) (Elem b) = Elem (f (fromRep 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, Elem a) | (k, a) <- xs])
	where f' k (Elem a) (Elem b) = Elem (f (fromRep 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, Elem 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
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 m = TSet (() <$ m)

-- | /O(1)/.  The key marking the position of the \"hole\" in the map.
key :: TKey k => TLocation k a -> k
key (TLoc hole) = fromRep (keyM hole)

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

-- | @'after' loc@ is the submap with keys greater than @'key' loc@.
after :: TKey k => TLocation k a -> TMap k a
after (TLoc hole) = TMap (afterM Nothing 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)@
search :: TKey k => k -> TMap k a -> (Maybe a, TLocation k a)
search k (TMap m) = case searchM (toRep k) m of
	(# a, hole #)	-> (getElem <$> a, TLoc 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 (unbox i) m of
	(# _, Elem a, hole #) -> (a, TLoc hole)

{-# INLINE extract #-}
extract :: (TKey k, MonadPlus m) => TMap k a -> m (a, TLocation k a)
extract (TMap m) = do
	(Elem a, hole) <- extractHoleM m
	return (a, TLoc hole)

-- | /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@
assign :: TKey k => a -> TLocation k a -> TMap k a
assign a (TLoc hole) = TMap (assignM (Elem a) hole)

-- | Return a map obtained by erasing the location.
--
-- @'clear' loc == 'before' loc `union` 'after' loc@
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))