Skip to content
Browse files

Migration to git

  • Loading branch information...
0 parents commit c3e85cd99a7d12b5ca404185d59dfebbc90dacdd @lowasser committed Jan 26, 2011
1,092 Data/TrieMap.hs
@@ -0,0 +1,1092 @@
+{-# LANGUAGE TypeFamilies, FlexibleContexts, UnboxedTuples, RecordWildCards #-}
+
+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.Representation
+import Data.TrieMap.Representation.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').
+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) = getValue <$> 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.
+{-# 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 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"
+{-# 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) = foldrM (\ (Assoc k a) -> f k a) m z
+
+-- | 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) = foldlM (\ z (Assoc k a) -> f z k a) m z
+
+-- | 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 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 (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, MonadPlus m) => (k -> a -> Maybe a) -> TMap k a -> m (Maybe (Assoc k a), Hole (Rep k) (Assoc k a))
+updateHelper f (TMap m) = do
+ (Assoc k a, loc) <- extractHoleM m
+ return (Assoc k <$> f k 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)
+
+-- | 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 m = TSet (() <$ 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 Nothing 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 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)@
+{-# INLINE search #-}
+search :: TKey k => k -> TMap k a -> (Maybe a, TLocation k a)
+search k (TMap m) = case searchM (toRep k) m of
+ (# Just (Assoc k a), hole #) -> (Just a, TLoc k hole)
+ (# _, hole #) -> (Nothing, 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 (unbox i) m of
+ (# _, Assoc k a, hole #) -> (a, TLoc k hole)
+
+{-# INLINE extract #-}
+extract :: (TKey k, MonadPlus m) => TMap k a -> m (a, TLocation k a)
+extract (TMap m) = do
+ (Assoc k a, hole) <- extractHoleM m
+ return (a, TLoc k 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@
+{-# INLINE assign #-}
+assign :: TKey k => a -> TLocation k a -> TMap k a
+assign a (TLoc k hole) = TMap (assignM (Just $ 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 (assignM Nothing 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))
68 Data/TrieMap/Applicative.hs
@@ -0,0 +1,68 @@
+{-# LANGUAGE StandaloneDeriving, GeneralizedNewtypeDeriving #-}
+
+module Data.TrieMap.Applicative where
+
+import Control.Applicative
+import Control.Monad
+
+import Data.Monoid hiding (Dual)
+
+instance Functor First where
+ fmap f (First m) = First (fmap f m)
+
+instance Functor Last where
+ fmap f (Last m) = Last (fmap f m)
+
+instance Monad First where
+ return = First . return
+ First m >>= k = First (m >>= getFirst . k)
+
+instance Monad Last where
+ return = Last . return
+ Last m >>= k = Last (m >>= getLast . k)
+
+instance MonadPlus First where
+ mzero = mempty
+ mplus = mappend
+
+instance MonadPlus Last where
+ mzero = mempty
+ mplus = mappend
+
+(.:) :: (c -> d) -> (a -> b -> c) -> a -> b -> d
+(f .: g) x y = f (g x y)
+
+(<.>) :: Functor f => (b -> c) -> (a -> f b) -> a -> f c
+f <.> g = fmap f . g
+
+(<.:>) :: Functor f => (c -> d) -> (a -> b -> f c) -> a -> b -> f d
+(f <.:> g) x y = f <$> g x y
+
+instance Applicative First where
+ pure = return
+ (<*>) = ap
+
+instance Alternative First where
+ empty = mempty
+ (<|>) = mplus
+
+instance Applicative Last where
+ pure = return
+ (<*>) = ap
+
+instance Alternative Last where
+ empty = mempty
+ (<|>) = mplus
+
+newtype DualPlus f a = DualPlus {runDualPlus :: f a} deriving (Functor, Applicative, Monad)
+newtype Dual f a = Dual {runDual :: f a} deriving (Functor)
+
+instance Applicative f => Applicative (Dual f) where
+ pure = Dual . pure
+ Dual f <*> Dual a = Dual (a <**> f)
+ Dual f *> Dual g = Dual (g <* f)
+ Dual f <* Dual g = Dual (g *> f)
+
+instance MonadPlus m => MonadPlus (DualPlus m) where
+ mzero = DualPlus mzero
+ DualPlus m `mplus` DualPlus k = DualPlus (k `mplus` m)
33 Data/TrieMap/Class.hs
@@ -0,0 +1,33 @@
+{-# LANGUAGE TypeFamilies, FlexibleContexts, FlexibleInstances, UndecidableInstances #-}
+
+module Data.TrieMap.Class (TMap(..), TSet(..), TKey, Rep, TrieMap, TrieKey) where
+
+import Data.TrieMap.TrieKey
+import Data.TrieMap.Representation.Class
+import Data.TrieMap.Sized
+
+import Control.Applicative
+import Data.Foldable hiding (foldrM, foldlM)
+import Data.Traversable
+
+import Prelude hiding (foldr)
+
+newtype TMap k a = TMap {getTMap :: TrieMap (Rep k) (Assoc k a)}
+
+newtype TSet a = TSet (TMap a ())
+
+-- | @'TKey' k@ is a handy alias for @('Repr' k, 'TrieKey' ('Rep' k))@. To make a type an instance of 'TKey',
+-- use the methods available in "Data.TrieMap.Representation.TH" to generate a 'Repr' instance that will
+-- satisfy @'TrieKey' ('Rep' k)@.
+class (Repr k, TrieKey (Rep k)) => TKey k
+
+instance (Repr k, TrieKey (Rep k)) => TKey k
+
+instance TKey k => Functor (TMap k) where
+ fmap = fmapDefault
+
+instance TKey k => Foldable (TMap k) where
+ foldr f z (TMap m) = foldrM (\ (Assoc _ a) -> f a) m z
+
+instance TKey k => Traversable (TMap k) where
+ traverse f (TMap m) = TMap <$> traverseM (\ (Assoc k a) -> Assoc k <$> f a) m
14 Data/TrieMap/Class/Instances.hs
@@ -0,0 +1,14 @@
+module Data.TrieMap.Class.Instances where
+
+import Data.TrieMap.Class ()
+import Data.TrieMap.TrieKey ()
+import Data.TrieMap.Representation.Instances ()
+import Data.TrieMap.Sized ()
+import Data.TrieMap.ReverseMap ()
+import Data.TrieMap.RadixTrie ()
+import Data.TrieMap.IntMap ()
+import Data.TrieMap.OrdMap ()
+import Data.TrieMap.ProdMap ()
+import Data.TrieMap.UnionMap ()
+import Data.TrieMap.UnitMap()
+import Data.TrieMap.Key ()
317 Data/TrieMap/IntMap.hs
@@ -0,0 +1,317 @@
+{-# LANGUAGE UnboxedTuples, BangPatterns, TypeFamilies, PatternGuards, MagicHash, CPP #-}
+{-# OPTIONS -funbox-strict-fields #-}
+module Data.TrieMap.IntMap () where
+
+import Data.TrieMap.TrieKey
+import Data.TrieMap.Sized
+
+import Control.Applicative
+import Control.Monad hiding (join)
+
+import Data.Bits
+import Data.Maybe hiding (mapMaybe)
+import Data.Word
+
+import GHC.Exts
+
+import Prelude hiding (lookup, null, foldl, foldr)
+
+#include "MachDeps.h"
+type Nat = Word
+
+type Prefix = Word
+type Mask = Word
+type Key = Word
+type Size = Int#
+
+data Path a = Root
+ | LeftBin !Prefix !Mask !(Path a) !(TrieMap Word a)
+ | RightBin !Prefix !Mask !(TrieMap Word a) !(Path a)
+
+instance TrieKey Word where
+ (=?) = (==)
+ cmp = compare
+
+ data TrieMap Word a = Nil
+ | Tip !Size !Key a
+ | Bin !Size !Prefix !Mask !(TrieMap Word a) !(TrieMap Word a)
+ data Hole Word a = Hole !Key !(Path a)
+ emptyM = Nil
+ singletonM = singleton
+ getSimpleM Nil = Null
+ getSimpleM (Tip _ _ a) = Singleton a
+ getSimpleM _ = NonSimple
+ sizeM = size
+ lookupM = lookup
+ traverseM = traverse
+ foldrM = foldr
+ foldlM = foldl
+ fmapM = mapWithKey
+ mapMaybeM = mapMaybe
+ mapEitherM = mapEither
+ unionM = unionWith
+ isectM = intersectionWith
+ diffM = differenceWith
+ isSubmapM = isSubmapOfBy
+
+ singleHoleM k = Hole k Root
+ beforeM a (Hole k path) = before (singletonMaybe k a) path where
+ before t Root = t
+ before t (LeftBin _ _ path _) = before t path
+ before t (RightBin p m l path) = before (bin p m l t) path
+ afterM a (Hole k path) = after (singletonMaybe k a) path where
+ after t Root = t
+ after t (RightBin _ _ _ path) = after t path
+ after t (LeftBin p m path r) = after (bin p m t r) path
+ searchM !k = onSnd (Hole k) (search Root) where
+ search path t@(Bin _ p m l r)
+ | nomatch k p m = (# Nothing, branchHole k p path t #)
+ | zero k m
+ = search (LeftBin p m path r) l
+ | otherwise
+ = search (RightBin p m l path) r
+ search path t@(Tip _ ky y)
+ | k == ky = (# Just y, path #)
+ | otherwise = (# Nothing, branchHole k ky path t #)
+ search path _ = (# Nothing, path #)
+ indexM i# t = indexT i# t Root where
+ indexT _ Nil _ = indexFail ()
+ indexT i# (Tip _ kx x) path = (# i#, x, Hole kx path #)
+ indexT i# (Bin _ p m l r) path
+ | i# <# sl# = indexT i# l (LeftBin p m path r)
+ | otherwise = indexT (i# -# sl#) r (RightBin p m l path)
+ where !sl# = size l
+ extractHoleM = extractHole Root where
+ extractHole _ Nil = mzero
+ extractHole path (Tip _ kx x) = return (x, Hole kx path)
+ extractHole path (Bin _ p m l r) =
+ extractHole (LeftBin p m path r) l `mplus`
+ extractHole (RightBin p m l path) r
+ assignM v (Hole kx path) = assign (singletonM' kx v) path where
+ assign t Root = t
+ assign t (LeftBin p m path r) = assign (bin p m t r) path
+ assign t (RightBin p m l path) = assign (bin p m l t) path
+
+ {-# INLINE unifyM #-}
+ unifyM = unify
+
+branchHole :: Key -> Prefix -> Path a -> TrieMap Word a -> Path a
+branchHole !k !p path t
+ | zero k m = LeftBin p' m path t
+ | otherwise = RightBin p' m t path
+ where m = branchMask k p
+ p' = mask k m
+
+natFromInt :: Word -> Nat
+natFromInt = id
+
+intFromNat :: Nat -> Word
+intFromNat = id
+
+shiftRL :: Nat -> Key -> Nat
+-- #if __GLASGOW_HASKELL__
+{--------------------------------------------------------------------
+ GHC: use unboxing to get @shiftRL@ inlined.
+--------------------------------------------------------------------}
+-- shiftRL (W# x) (I# i)
+-- = W# (shiftRL# x i)
+-- #else
+shiftRL x i = shiftR x (fromIntegral i)
+-- #endif
+
+size :: TrieMap Word a -> Int#
+size Nil = 0#
+size (Tip sz _ _) = sz
+size (Bin sz _ _ _ _) = sz
+
+lookup :: Nat -> TrieMap Word a -> Maybe a
+lookup !k (Bin _ _ m l r) = lookup k (if zeroN k m then l else r)
+lookup k (Tip _ kx x)
+ | k == kx = Just x
+lookup _ _ = Nothing
+
+singleton :: Sized a => Key -> a -> TrieMap Word a
+singleton k a = Tip (getSize# a) k a
+
+singletonMaybe :: Sized a => Key -> Maybe a -> TrieMap Word a
+singletonMaybe k = maybe Nil (singleton k)
+
+traverse :: (Applicative f, Sized b) => (a -> f b) -> TrieMap Word a -> f (TrieMap Word b)
+traverse f t = case t of
+ Nil -> pure Nil
+ Tip _ kx x -> singleton kx <$> f x
+ Bin _ p m l r -> bin p m <$> traverse f l <*> traverse f r
+
+foldr :: (a -> b -> b) -> TrieMap Word a -> b -> b
+foldr f t
+ = case t of
+ Bin _ _ _ l r -> foldr f l . foldr f r
+ Tip _ _ x -> f x
+ Nil -> id
+
+foldl :: (b -> a -> b) -> TrieMap Word a -> b -> b
+foldl f t
+ = case t of
+ Bin _ _ _ l r -> foldl f r . foldl f l
+ Tip _ _ x -> flip f x
+ Nil -> id
+
+mapWithKey :: Sized b => (a -> b) -> TrieMap Word a -> TrieMap Word b
+mapWithKey f (Bin _ p m l r) = bin p m (mapWithKey f l) (mapWithKey f r)
+mapWithKey f (Tip _ kx x) = singleton kx (f x)
+mapWithKey _ _ = Nil
+
+mapMaybe :: Sized b => (a -> Maybe b) -> TrieMap Word a -> TrieMap Word b
+mapMaybe f (Bin _ p m l r) = bin p m (mapMaybe f l) (mapMaybe f r)
+mapMaybe f (Tip _ kx x) = singletonMaybe kx (f x)
+mapMaybe _ _ = Nil
+
+mapEither :: (Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) ->
+ TrieMap Word a -> (# TrieMap Word b, TrieMap Word c #)
+mapEither f (Bin _ p m l r) = both (bin p m lL) (bin p m lR) (mapEither f) r
+ where !(# lL, lR #) = mapEither f l
+mapEither f (Tip _ kx x) = both (singletonMaybe kx) (singletonMaybe kx) f x
+mapEither _ _ = (# Nil, Nil #)
+
+unionWith :: Sized a => (a -> a -> Maybe a) -> TrieMap Word a -> TrieMap Word a -> TrieMap Word a
+unionWith _ Nil t = t
+unionWith _ t Nil = t
+unionWith f (Tip _ k x) t = alterM (maybe (Just x) (f x)) k t
+unionWith f t (Tip _ k x) = alterM (maybe (Just x) (flip f x)) k t
+unionWith f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2)
+ | shorter m1 m2 = union1
+ | shorter m2 m1 = union2
+ | p1 == p2 = bin p1 m1 (unionWith f l1 l2) (unionWith f r1 r2)
+ | otherwise = join p1 t1 p2 t2
+ where
+ union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2
+ | zero p2 m1 = bin p1 m1 (unionWith f l1 t2) r1
+ | otherwise = bin p1 m1 l1 (unionWith f r1 t2)
+
+ union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2
+ | zero p1 m2 = bin p2 m2 (unionWith f t1 l2) r2
+ | otherwise = bin p2 m2 l2 (unionWith f t1 r2)
+
+intersectionWith :: Sized c => (a -> b -> Maybe c) -> TrieMap Word a -> TrieMap Word b -> TrieMap Word c
+intersectionWith _ Nil _ = Nil
+intersectionWith _ _ Nil = Nil
+intersectionWith f (Tip _ k x) t2
+ = singletonMaybe k (lookup (natFromInt k) t2 >>= f x)
+intersectionWith f t1 (Tip _ k y)
+ = singletonMaybe k (lookup (natFromInt k) t1 >>= flip f y)
+intersectionWith f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2)
+ | shorter m1 m2 = intersection1
+ | shorter m2 m1 = intersection2
+ | p1 == p2 = bin p1 m1 (intersectionWith f l1 l2) (intersectionWith f r1 r2)
+ | otherwise = Nil
+ where
+ intersection1 | nomatch p2 p1 m1 = Nil
+ | zero p2 m1 = intersectionWith f l1 t2
+ | otherwise = intersectionWith f r1 t2
+
+ intersection2 | nomatch p1 p2 m2 = Nil
+ | zero p1 m2 = intersectionWith f t1 l2
+ | otherwise = intersectionWith f t1 r2
+
+differenceWith :: Sized a => (a -> b -> Maybe a) -> TrieMap Word a -> TrieMap Word b -> TrieMap Word a
+differenceWith _ Nil _ = Nil
+differenceWith _ t Nil = t
+differenceWith f t1@(Tip _ k x) t2
+ = maybe t1 (singletonMaybe k . f x) (lookup (natFromInt k) t2)
+differenceWith f t (Tip _ k y) = alterM (>>= flip f y) k t
+differenceWith f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2)
+ | shorter m1 m2 = difference1
+ | shorter m2 m1 = difference2
+ | p1 == p2 = bin p1 m1 (differenceWith f l1 l2) (differenceWith f r1 r2)
+ | otherwise = t1
+ where
+ difference1 | nomatch p2 p1 m1 = t1
+ | zero p2 m1 = bin p1 m1 (differenceWith f l1 t2) r1
+ | otherwise = bin p1 m1 l1 (differenceWith f r1 t2)
+
+ difference2 | nomatch p1 p2 m2 = t1
+ | zero p1 m2 = differenceWith f t1 l2
+ | otherwise = differenceWith f t1 r2
+
+isSubmapOfBy :: LEq a b -> LEq (TrieMap Word a) (TrieMap Word b)
+isSubmapOfBy (<=) t1@(Bin _ p1 m1 l1 r1) (Bin _ p2 m2 l2 r2)
+ | shorter m1 m2 = False
+ | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy (<=) t1 l2
+ else isSubmapOfBy (<=) t1 r2)
+ | otherwise = (p1==p2) && isSubmapOfBy (<=) l1 l2 && isSubmapOfBy (<=) r1 r2
+isSubmapOfBy _ (Bin _ _ _ _ _) _
+ = False
+isSubmapOfBy (<=) (Tip _ k x) t
+ = maybe False (x <=) (lookup (natFromInt k) t)
+isSubmapOfBy _ Nil _
+ = True
+
+mask :: Key -> Mask -> Prefix
+mask i m
+ = maskW (natFromInt i) (natFromInt m)
+
+zero :: Key -> Mask -> Bool
+zero i m
+ = (natFromInt i) .&. (natFromInt m) == 0
+
+nomatch,match :: Key -> Prefix -> Mask -> Bool
+nomatch i p m
+ = (mask i m) /= p
+
+match i p m
+ = (mask i m) == p
+
+zeroN :: Nat -> Nat -> Bool
+zeroN i m = (i .&. m) == 0
+
+maskW :: Nat -> Nat -> Prefix
+maskW i m
+ = intFromNat (i .&. (complement (m-1) `xor` m))
+
+shorter :: Mask -> Mask -> Bool
+shorter m1 m2
+ = (natFromInt m1) > (natFromInt m2)
+
+branchMask :: Prefix -> Prefix -> Mask
+branchMask p1 p2
+ = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
+
+highestBitMask :: Nat -> Nat
+highestBitMask x0
+ = case (x0 .|. shiftRL x0 1) of
+ x1 -> case (x1 .|. shiftRL x1 2) of
+ x2 -> case (x2 .|. shiftRL x2 4) of
+ x3 -> case (x3 .|. shiftRL x3 8) of
+ x4 -> case (x4 .|. shiftRL x4 16) of
+#if WORD_SIZE_IN_BITS > 32
+ x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms
+ x6 -> (x6 `xor` (shiftRL x6 1))
+#else
+ x5 -> x5 `xor` shiftRL x5 1
+#endif
+
+{-# INLINE join #-}
+join :: Prefix -> TrieMap Word a -> Prefix -> TrieMap Word a -> TrieMap Word a
+join p1 t1 p2 t2
+ | zero p1 m = bin p m t1 t2
+ | otherwise = bin p m t2 t1
+ where
+ m = branchMask p1 p2
+ p = mask p1 m
+
+bin :: Prefix -> Mask -> TrieMap Word a -> TrieMap Word a -> TrieMap Word a
+bin _ _ l Nil = l
+bin _ _ Nil r = r
+bin p m l r = Bin (size l +# size r) p m l r
+
+{-# INLINE unify #-}
+unify :: Sized a => Key -> a -> Key -> a -> Unified Word a
+unify k1 _ k2 _
+ | k1 == k2 = Left (Hole k1 Root)
+unify k1 a1 k2 a2 = Right (if zero k1 m then outBin t1 t2 else outBin t2 t1)
+ where !s1# = getSize# a1
+ !s2# = getSize# a2
+ t1 = Tip s1# k1 a1
+ t2 = Tip s2# k2 a2
+ m = branchMask k1 k2
+ outBin = Bin (s1# +# s2#) (mask k1 m) m
52 Data/TrieMap/Key.hs
@@ -0,0 +1,52 @@
+{-# LANGUAGE TypeFamilies, UnboxedTuples #-}
+
+module Data.TrieMap.Key () where
+
+import Control.Applicative
+
+import Data.TrieMap.Class
+import Data.TrieMap.TrieKey
+import Data.TrieMap.Representation.Class
+import Data.TrieMap.Modifiers
+
+import Data.TrieMap.ProdMap()
+import Data.TrieMap.UnionMap()
+import Data.TrieMap.IntMap()
+import Data.TrieMap.OrdMap()
+import Data.TrieMap.RadixTrie()
+
+instance TKey k => TrieKey (Key k) where
+ Key k1 =? Key k2 = toRep k1 =? toRep k2
+ Key k1 `cmp` Key k2 = toRep k1 `cmp` toRep k2
+
+ newtype TrieMap (Key k) a = KeyMap (TrieMap (Rep k) a)
+ newtype Hole (Key k) a = KeyHole (Hole (Rep k) a)
+
+ emptyM = KeyMap emptyM
+ singletonM (Key k) a = KeyMap (singletonM (toRep k) a)
+ getSimpleM (KeyMap m) = getSimpleM m
+ sizeM (KeyMap m) = sizeM m
+ lookupM (Key k) (KeyMap m) = lookupM (toRep k) m
+ traverseM f (KeyMap m) = KeyMap <$> traverseM f m
+ foldrM f (KeyMap m) = foldrM f m
+ foldlM f (KeyMap m) = foldlM f m
+ fmapM f (KeyMap m) = KeyMap (fmapM f m)
+ mapMaybeM f (KeyMap m) = KeyMap (mapMaybeM f m)
+ mapEitherM f (KeyMap m) = both KeyMap KeyMap (mapEitherM f) m
+ unionM f (KeyMap m1) (KeyMap m2) = KeyMap (unionM f m1 m2)
+ isectM f (KeyMap m1) (KeyMap m2) = KeyMap (isectM f m1 m2)
+ diffM f (KeyMap m1) (KeyMap m2) = KeyMap (diffM f m1 m2)
+ isSubmapM (<=) (KeyMap m1) (KeyMap m2) = isSubmapM (<=) m1 m2
+
+ singleHoleM (Key k) = KeyHole (singleHoleM (toRep k))
+ beforeM a (KeyHole hole) = KeyMap (beforeM a hole)
+ afterM a (KeyHole hole) = KeyMap (afterM a hole)
+ searchM (Key k) (KeyMap m) = onSnd KeyHole (searchM (toRep k)) m
+ indexM i (KeyMap m) = case indexM i m of
+ (# i', v, hole #) -> (# i', v, KeyHole hole #)
+ extractHoleM (KeyMap m) = do
+ (v, hole) <- extractHoleM m
+ return (v, KeyHole hole)
+ assignM v (KeyHole hole) = KeyMap (assignM v hole)
+
+ unifyM (Key k1) a1 (Key k2) a2 = either (Left . KeyHole) (Right . KeyMap) (unifyM (toRep k1) a1 (toRep k2) a2)
21 Data/TrieMap/Modifiers.hs
@@ -0,0 +1,21 @@
+{-# LANGUAGE FlexibleContexts, UndecidableInstances, TypeFamilies #-}
+module Data.TrieMap.Modifiers where
+
+import Data.TrieMap.Representation.Class
+
+newtype Ordered a = Ord {unOrd :: a} deriving (Eq, Ord)
+newtype Rev k = Rev {getRev :: k} deriving (Eq)
+instance Ord k => Ord (Rev k) where
+ compare (Rev a) (Rev b) = compare b a
+
+instance Functor Ordered where
+ fmap f (Ord a) = Ord (f a)
+
+instance Functor Rev where
+ fmap f (Rev a) = Rev (f a)
+
+newtype Key k = Key {getKey :: k}
+
+instance Repr k => Repr (Key k) where
+ type Rep (Key k) = Rep k
+ toRep (Key k) = toRep k
391 Data/TrieMap/OrdMap.hs
@@ -0,0 +1,391 @@
+{-# LANGUAGE BangPatterns, UnboxedTuples, TypeFamilies, PatternGuards, MagicHash, CPP, TupleSections #-}
+
+module Data.TrieMap.OrdMap () where
+
+import Data.TrieMap.TrieKey
+import Data.TrieMap.Sized
+import Data.TrieMap.Modifiers
+
+import Control.Applicative
+import Control.Monad hiding (join, fmap)
+
+import Prelude hiding (lookup, foldr, foldl, fmap)
+
+import GHC.Exts
+
+#define DELTA 5#
+#define RATIO 2#
+
+type OrdMap k = TrieMap (Ordered k)
+
+data Path k a =
+ Root
+ | LeftBin k a !(Path k a) !(OrdMap k a)
+ | RightBin k a !(OrdMap k a) !(Path k a)
+
+singletonMaybe :: Sized a => k -> Maybe a -> OrdMap k a
+singletonMaybe k = maybe Tip (singleton k)
+
+instance Ord k => TrieKey (Ordered k) where
+ Ord k1 =? Ord k2 = k1 == k2
+ Ord k1 `cmp` Ord k2 = k1 `compare` k2
+
+ data TrieMap (Ordered k) a = Tip
+ | Bin Int# k a !(OrdMap k a) !(OrdMap k a)
+ data Hole (Ordered k) a =
+ Empty k !(Path k a)
+ | Full k !(Path k a) !(OrdMap k a) !(OrdMap k a)
+ emptyM = Tip
+ singletonM (Ord k) = singleton k
+ lookupM (Ord k) = lookup k
+ getSimpleM Tip = Null
+ getSimpleM (Bin _ _ a Tip Tip) = Singleton a
+ getSimpleM _ = NonSimple
+ sizeM = size#
+ traverseM = traverse
+ foldrM = foldr
+ foldlM = foldl
+ fmapM = fmap
+ mapMaybeM = mapMaybe
+ mapEitherM = mapEither
+ isSubmapM = isSubmap
+ fromAscListM f xs = fromAscList f [(k, a) | (Ord k, a) <- xs]
+ fromDistAscListM xs = fromDistinctAscList [(k, a) | (Ord k, a) <- xs]
+ unionM _ Tip m2 = m2
+ unionM _ m1 Tip = m1
+ unionM f m1 m2 = hedgeUnion f (const LT) (const GT) m1 m2
+ isectM = isect
+ diffM _ Tip _ = Tip
+ diffM _ m1 Tip = m1
+ diffM f m1 m2 = hedgeDiff f (const LT) (const GT) m1 m2
+
+ singleHoleM (Ord k) = Empty k Root
+ beforeM a (Empty k path) = before (singletonMaybe k a) path
+ beforeM a (Full k path l _) = before t path
+ where t = case a of
+ Nothing -> l
+ Just a -> insertMax k a l
+ afterM a (Empty k path) = after (singletonMaybe k a) path
+ afterM a (Full k path _ r) = after t path
+ where t = case a of
+ Nothing -> r
+ Just a -> insertMin k a r
+ searchM (Ord k) = search k Root
+ indexM i# = indexT Root i# where
+ indexT path i# (Bin _ kx x l r)
+ | i# <# sl# = indexT (LeftBin kx x path r) i# l
+ | i# <# sx# = (# i# -# sl#, x, Full kx path l r #)
+ | otherwise = indexT (RightBin kx x l path) (i# -# sx#) r
+ where !sl# = size# l
+ !sx# = getSize# x +# sl#
+ indexT _ _ _ = indexFail ()
+ extractHoleM = extractHole Root where
+ extractHole path (Bin _ kx x l r) =
+ extractHole (LeftBin kx x path r) l `mplus`
+ return (x, Full kx path l r) `mplus`
+ extractHole (RightBin kx x l path) r
+ extractHole _ _ = mzero
+ assignM x (Empty k path) = rebuild (maybe Tip (singleton k) x) path
+ assignM x (Full k path l r) = rebuild (joinMaybe k x l r) path
+
+ unifyM (Ord k1) a1 (Ord k2) a2 = case compare k1 k2 of
+ EQ -> Left $ Empty k1 Root
+ LT -> Right $ bin k1 a1 Tip (singleton k2 a2)
+ GT -> Right $ bin k1 a1 (singleton k2 a2) Tip
+
+rebuild :: Sized a => OrdMap k a -> Path k a -> OrdMap k a
+rebuild t Root = t
+rebuild t (LeftBin kx x path r) = rebuild (balance kx x t r) path
+rebuild t (RightBin kx x l path) = rebuild (balance kx x l t) path
+
+lookup :: Ord k => k -> OrdMap k a -> Maybe a
+lookup k (Bin _ k' v l r) = case compare k k' of
+ LT -> lookup k l
+ EQ -> Just v
+ GT -> lookup k r
+lookup _ _ = Nothing
+
+singleton :: Sized a => k -> a -> OrdMap k a
+singleton k a = Bin (getSize# a) k a Tip Tip
+
+traverse :: (Applicative f, Sized b) => (a -> f b) -> OrdMap k a -> f (OrdMap k b)
+traverse _ Tip = pure Tip
+traverse f (Bin _ k a l r) = balance k <$> f a <*> traverse f l <*> traverse f r
+
+foldr :: (a -> b -> b) -> OrdMap k a -> b -> b
+foldr _ Tip = id
+foldr f (Bin _ _ a l r) = foldr f l . f a . foldr f r
+
+foldl :: (b -> a -> b) -> OrdMap k a -> b -> b
+foldl _ Tip = id
+foldl f (Bin _ _ a l r) = foldl f r . flip f a . foldl f l
+
+fmap :: (Ord k, Sized b) => (a -> b) -> OrdMap k a -> OrdMap k b
+fmap f (Bin _ k a l r) = join k (f a) (fmap f l) (fmap f r)
+fmap _ _ = Tip
+
+mapMaybe :: (Ord k, Sized b) => (a -> Maybe b) -> OrdMap k a -> OrdMap k b
+mapMaybe f (Bin _ k a l r) = joinMaybe k (f a) (mapMaybe f l) (mapMaybe f r)
+mapMaybe _ _ = Tip
+
+mapEither :: (Ord k, Sized b, Sized c) => (a -> (# Maybe b, Maybe c #)) ->
+ OrdMap k a -> (# OrdMap k b, OrdMap k c #)
+mapEither f (Bin _ k a l r) = (# joinMaybe k aL lL rL, joinMaybe k aR lR rR #)
+ where !(# aL, aR #) = f a; !(# lL, lR #) = mapEither f l; !(# rL, rR #) = mapEither f r
+mapEither _ _ = (# Tip, Tip #)
+
+splitLookup :: (Ord k, Sized a) => SplitMap a x -> k -> OrdMap k a -> (# OrdMap k a, Maybe x, OrdMap k a #)
+splitLookup f k m = case m of
+ Tip -> (# Tip, Nothing, Tip #)
+ Bin _ kx x l r -> case compare k kx of
+ LT -> let !(# lL, ans, lR #) = splitLookup f k l in (# lL, ans, join kx x lR r #)
+ EQ -> let !(# xL, ans, xR #) = f x in
+ (# maybe l (\ xL -> insertMax kx xL l) xL, ans, maybe r (\ xR -> insertMin kx xR r) xR #)
+ GT -> let !(# rL, ans, rR #) = splitLookup f k r in (# join kx x l rL, ans, rR #)
+
+isSubmap :: (Ord k, Sized a, Sized b) => LEq a b -> LEq (OrdMap k a) (OrdMap k b)
+isSubmap _ Tip _ = True
+isSubmap _ _ Tip = False
+isSubmap (<=) (Bin _ kx x l r) t = case found of
+ Nothing -> False
+ Just (Elem y) -> x <= y && isSubmap (<=) l lt && isSubmap (<=) r gt
+ where !(# lt, found, gt #) = splitLookup (\ x -> (# Nothing, Just (Elem x), Nothing #)) kx t
+
+fromAscList :: (Eq k, Sized a) => (a -> a -> a) -> [(k, a)] -> OrdMap k a
+fromAscList f xs = fromDistinctAscList (combineEq xs) where
+ combineEq (x:xs) = combineEq' x xs
+ combineEq [] = []
+
+ combineEq' z [] = [z]
+ combineEq' (kz, zz) (x@(kx, xx):xs)
+ | kz == kx = combineEq' (kx, f xx zz) xs
+ | otherwise = (kz,zz):combineEq' x xs
+
+fromDistinctAscList :: Sized a => [(k, a)] -> OrdMap k a
+fromDistinctAscList xs = build const (length xs) xs
+ where
+ -- 1) use continutations so that we use heap space instead of stack space.
+ -- 2) special case for n==5 to build bushier trees.
+ build c 0 xs' = c Tip xs'
+ build c 5 xs' = case xs' of
+ ((k1,x1):(k2,x2):(k3,x3):(k4,x4):(k5,x5):xx)
+ -> c (bin k4 x4 (bin k2 x2 (singleton k1 x1) (singleton k3 x3)) (singleton k5 x5)) xx
+ _ -> error "fromDistinctAscList build"
+ build c n xs' = seq nr $ build (buildR nr c) nl xs'
+ where
+ nl = n `div` 2
+ nr = n - nl - 1
+
+ buildR n c l ((k,x):ys) = build (buildB l k x c) n ys
+ buildR _ _ _ [] = error "fromDistinctAscList buildR []"
+ buildB l k x c r zs = c (bin k x l r) zs
+
+hedgeUnion :: (Ord k, Sized a)
+ => (a -> a -> Maybe a)
+ -> (k -> Ordering) -> (k -> Ordering)
+ -> OrdMap k a -> OrdMap k a -> OrdMap k a
+hedgeUnion _ _ _ t1 Tip
+ = t1
+hedgeUnion _ cmplo cmphi Tip (Bin _ kx x l r)
+ = join kx x (filterGt cmplo l) (filterLt cmphi r)
+hedgeUnion f cmplo cmphi (Bin _ kx x l r) t2
+ = joinMaybe kx newx (hedgeUnion f cmplo cmpkx l lt)
+ (hedgeUnion f cmpkx cmphi r gt)
+ where
+ cmpkx k = compare kx k
+ lt = trim cmplo cmpkx t2
+ (found,gt) = trimLookupLo kx cmphi t2
+ newx = case found of
+ Nothing -> Just x
+ Just (_,y) -> f x y
+
+filterGt :: (Ord k, Sized a) => (k -> Ordering) -> OrdMap k a -> OrdMap k a
+filterGt _ Tip = Tip
+filterGt cmp (Bin _ kx x l r)
+ = case cmp kx of
+ LT -> join kx x (filterGt cmp l) r
+ GT -> filterGt cmp r
+ EQ -> r
+
+filterLt :: (Ord k, Sized a) => (k -> Ordering) -> OrdMap k a -> OrdMap k a
+filterLt _ Tip = Tip
+filterLt cmp (Bin _ kx x l r)
+ = case cmp kx of
+ LT -> filterLt cmp l
+ GT -> join kx x l (filterLt cmp r)
+ EQ -> l
+
+trim :: (k -> Ordering) -> (k -> Ordering) -> OrdMap k a -> OrdMap k a
+trim _ _ Tip = Tip
+trim cmplo cmphi t@(Bin _ kx _ l r)
+ = case cmplo kx of
+ LT -> case cmphi kx of
+ GT -> t
+ _ -> trim cmplo cmphi l
+ _ -> trim cmplo cmphi r
+
+trimLookupLo :: Ord k => k -> (k -> Ordering) -> OrdMap k a -> (Maybe (k,a), OrdMap k a)
+trimLookupLo _ _ Tip = (Nothing,Tip)
+trimLookupLo lo cmphi t@(Bin _ kx x l r)
+ = case compare lo kx of
+ LT -> case cmphi kx of
+ GT -> ((lo,) <$> lookup lo t, t)
+ _ -> trimLookupLo lo cmphi l
+ GT -> trimLookupLo lo cmphi r
+ EQ -> (Just (kx,x),trim (compare lo) cmphi r)
+
+isect :: (Ord k, Sized a, Sized b, Sized c) => (a -> b -> Maybe c) -> OrdMap k a -> OrdMap k b -> OrdMap k c
+isect f t1@Bin{} (Bin _ k2 x2 l2 r2)
+ = joinMaybe k2 (found >>= \ x1' -> f x1' x2) tl tr
+ where !(# found, hole #) = search k2 Root t1
+ tl = isect f (beforeM Nothing hole) l2
+ tr = isect f (afterM Nothing hole) r2
+isect _ _ _ = Tip
+
+hedgeDiff :: (Ord k, Sized a)
+ => (a -> b -> Maybe a)
+ -> (k -> Ordering) -> (k -> Ordering)
+ -> OrdMap k a -> OrdMap k b -> OrdMap k a
+hedgeDiff _ _ _ Tip _
+ = Tip
+hedgeDiff _ cmplo cmphi (Bin _ kx x l r) Tip
+ = join kx x (filterGt cmplo l) (filterLt cmphi r)
+hedgeDiff f cmplo cmphi t (Bin _ kx x l r)
+ = case found of
+ Nothing -> merge tl tr
+ Just (ky,y) ->
+ case f y x of
+ Nothing -> merge tl tr
+ Just z -> join ky z tl tr
+ where
+ cmpkx k = compare kx k
+ lt = trim cmplo cmpkx t
+ (found,gt) = trimLookupLo kx cmphi t
+ tl = hedgeDiff f cmplo cmpkx lt l
+ tr = hedgeDiff f cmpkx cmphi gt r
+
+joinMaybe :: (Ord k, Sized a) => k -> Maybe a -> OrdMap k a -> OrdMap k a -> OrdMap k a
+joinMaybe kx = maybe merge (join kx)
+
+join :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a
+join kx x Tip r = insertMin kx x r
+join kx x l Tip = insertMax kx x l
+join kx x l@(Bin sL# ky y ly ry) r@(Bin sR# kz z lz rz)
+ | DELTA *# sL# <=# sR# = balance kz z (join kx x l lz) rz
+ | DELTA *# sR# <=# sL# = balance ky y ly (join kx x ry r)
+ | otherwise = bin kx x l r
+
+-- insertMin and insertMax don't perform potentially expensive comparisons.
+insertMax,insertMin :: Sized a => k -> a -> OrdMap k a -> OrdMap k a
+insertMax kx x t
+ = case t of
+ Tip -> singleton kx x
+ Bin _ ky y l r
+ -> balance ky y l (insertMax kx x r)
+
+insertMin kx x t
+ = case t of
+ Tip -> singleton kx x
+ Bin _ ky y l r
+ -> balance ky y (insertMin kx x l) r
+
+{--------------------------------------------------------------------
+ [merge l r]: merges two trees.
+--------------------------------------------------------------------}
+merge :: Sized a => OrdMap k a -> OrdMap k a -> OrdMap k a
+merge Tip r = r
+merge l Tip = l
+merge l@(Bin sL# kx x lx rx) r@(Bin sR# ky y ly ry)
+ | DELTA *# sL# <=# sR# = balance ky y (merge l ly) ry
+ | DELTA *# sR# <=# sL# = balance kx x lx (merge rx r)
+ | otherwise = glue l r
+
+{--------------------------------------------------------------------
+ [glue l r]: glues two trees together.
+ Assumes that [l] and [r] are already balanced with respect to each other.
+--------------------------------------------------------------------}
+glue :: Sized a => OrdMap k a -> OrdMap k a -> OrdMap k a
+glue Tip r = r
+glue l Tip = l
+glue l r
+ | size# l ># size# r = let !(# f, l' #) = deleteFindMax (\ k a -> (# balance k a, Nothing #)) l in f l' r
+ | otherwise = let !(# f, r' #) = deleteFindMin (\ k a -> (# balance k a, Nothing #)) r in f l r'
+
+deleteFindMin :: Sized a => (k -> a -> (# x, Maybe a #)) -> OrdMap k a -> (# x, OrdMap k a #)
+deleteFindMin f t
+ = case t of
+ Bin _ k x Tip r -> onSnd (maybe r (\ y' -> bin k y' Tip r)) (f k) x
+ Bin _ k x l r -> onSnd (\ l' -> balance k x l' r) (deleteFindMin f) l
+ _ -> (# error "Map.deleteFindMin: can not return the minimal element of an empty fmap", Tip #)
+
+deleteFindMax :: Sized a => (k -> a -> (# x, Maybe a #)) -> OrdMap k a -> (# x, OrdMap k a #)
+deleteFindMax f t
+ = case t of
+ Bin _ k x l Tip -> onSnd (maybe l (\ y -> bin k y l Tip)) (f k) x
+ Bin _ k x l r -> onSnd (balance k x l) (deleteFindMax f) r
+ Tip -> (# error "Map.deleteFindMax: can not return the maximal element of an empty fmap", Tip #)
+
+size# :: OrdMap k a -> Int#
+size# Tip = 0#
+size# (Bin sz _ _ _ _) = sz
+
+balance :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a
+balance k x l r
+ | sR# >=# (DELTA *# sL#) = rotateL k x l r
+ | sL# >=# (DELTA *# sR#) = rotateR k x l r
+ | otherwise = Bin sX# k x l r
+ where
+ !sL# = size# l
+ !sR# = size# r
+ !sX# = sL# +# sR# +# getSize# x
+
+-- rotate
+rotateL :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a
+rotateL k x l r@(Bin _ _ _ ly ry)
+ | sL# <# (RATIO *# sR#) = singleL k x l r
+ | otherwise = doubleL k x l r
+ where !sL# = size# ly
+ !sR# = size# ry
+rotateL _ _ _ Tip = error "rotateL Tip"
+
+rotateR :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a
+rotateR k x l@(Bin _ _ _ ly ry) r
+ | sR# <# (RATIO *# sL#) = singleR k x l r
+ | otherwise = doubleR k x l r
+ where !sL# = size# ly
+ !sR# = size# ry
+rotateR _ _ _ _ = error "rotateR Tip"
+
+-- basic rotations
+singleL, singleR :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a
+singleL k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin k2 x2 (bin k1 x1 t1 t2) t3
+singleL k1 x1 t1 Tip = bin k1 x1 t1 Tip
+singleR k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3)
+singleR k1 x1 Tip t2 = bin k1 x1 Tip t2
+
+doubleL, doubleR :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a
+doubleL k1 x1 t1 (Bin _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4)
+doubleL k1 x1 t1 t2 = singleL k1 x1 t1 t2
+doubleR k1 x1 (Bin _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4)
+doubleR k1 x1 t1 t2 = singleR k1 x1 t1 t2
+
+bin :: Sized a => k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a
+bin k x l r
+ = Bin (size# l +# size# r +# getSize# x) k x l r
+
+before :: Sized a => OrdMap k a -> Path k a -> OrdMap k a
+before t (LeftBin _ _ path _) = before t path
+before t (RightBin k a l path) = before (join k a l t) path
+before t _ = t
+
+after :: Sized a => OrdMap k a -> Path k a -> OrdMap k a
+after t (LeftBin k a path r) = after (join k a t r) path
+after t (RightBin _ _ _ path) = after t path
+after t _ = t
+
+search :: Ord k => k -> Path k a -> OrdMap k a -> (# Maybe a, Hole (Ordered k) a #)
+search k path Tip = (# Nothing, Empty k path #)
+search k path (Bin _ kx x l r) = case compare k kx of
+ LT -> search k (LeftBin kx x path r) l
+ EQ -> (# Just x, Full k path l r #)
+ GT -> search k (RightBin kx x l path) r
68 Data/TrieMap/ProdMap.hs
@@ -0,0 +1,68 @@
+{-# LANGUAGE UnboxedTuples, TupleSections, PatternGuards, TypeFamilies #-}
+
+module Data.TrieMap.ProdMap () where
+
+import Data.TrieMap.Sized
+import Data.TrieMap.TrieKey
+
+import Control.Applicative
+
+import Data.Foldable hiding (foldlM, foldrM)
+import Data.Monoid
+
+import Data.Sequence ((|>))
+import qualified Data.Sequence as Seq
+
+instance (TrieKey k1, TrieKey k2) => TrieKey (k1, k2) where
+ (k11, k12) =? (k21, k22) = k11 =? k21 && k12 =? k22
+ (k11, k12) `cmp` (k21, k22) = (k11 `cmp` k21) `mappend` (k12 `cmp` k22)
+
+ newtype TrieMap (k1, k2) a = PMap (TrieMap k1 (TrieMap k2 a))
+ data Hole (k1, k2) a = PHole (Hole k1 (TrieMap k2 a)) (Hole k2 a)
+
+ emptyM = PMap emptyM
+ singletonM (k1, k2) = PMap . singletonM k1 . singletonM k2
+ getSimpleM (PMap m) = getSimpleM m >>= getSimpleM
+ sizeM (PMap m) = sizeM m
+ lookupM (k1, k2) (PMap m) = lookupM k1 m >>= lookupM k2
+ traverseM f (PMap m) = PMap <$> traverseM (traverseM f) m
+ foldrM f (PMap m) = foldrM (foldrM f) m
+ foldlM f (PMap m) = foldlM (flip $ foldlM f) m
+ fmapM f (PMap m) = PMap (fmapM (fmapM f) m)
+ mapMaybeM f (PMap m) = PMap (mapMaybeM (mapMaybeM' f) m)
+ mapEitherM f (PMap m) = both PMap PMap (mapEitherM (mapEitherM' f)) m
+ isSubmapM (<=) (PMap m1) (PMap m2) = isSubmapM (isSubmapM (<=)) m1 m2
+ unionM f (PMap m1) (PMap m2) = PMap (unionM (unionM' f) m1 m2)
+ isectM f (PMap m1) (PMap m2) = PMap (isectM (isectM' f) m1 m2)
+ diffM f (PMap m1) (PMap m2) = PMap (diffM (diffM' f) m1 m2)
+ fromAscListM f xs = PMap (fromDistAscListM
+ [(a, fromAscListM f ys) | (a, Elem ys) <- breakFst xs])
+ fromDistAscListM xs = PMap (fromDistAscListM
+ [(a, fromDistAscListM ys) | (a, Elem ys) <- breakFst xs])
+
+ singleHoleM (k1, k2) = PHole (singleHoleM k1) (singleHoleM k2)
+ assignM v (PHole hole1 hole2) = PMap (assignM (assignM' v hole2) hole1)
+ beforeM a (PHole hole1 hole2) = PMap (beforeM (beforeM' a hole2) hole1)
+ afterM a (PHole hole1 hole2) = PMap (afterM (afterM' a hole2) hole1)
+ searchM (k1, k2) (PMap m) = onSnd (PHole hole1) (searchM' k2) m'
+ where !(# m', hole1 #) = searchM k1 m
+ indexM i (PMap m) = onThird (PHole hole1) (indexM i') m'
+ where !(# i', m', hole1 #) = indexM i m
+ extractHoleM (PMap m) = do
+ (m', hole1) <- extractHoleM m
+ (v, hole2) <- extractHoleM m'
+ return (v, PHole hole1 hole2)
+
+ unifyM (k11, k12) a1 (k21, k22) a2 = case unifyM k11 (singletonM k12 a1) k21 (singletonM k22 a2) of
+ Left hole -> case unifyM k12 a1 k22 a2 of
+ Left hole' -> Left (PHole hole hole')
+ Right m' -> Right (PMap (assignM (Just m') hole))
+ Right m -> Right (PMap m)
+
+breakFst :: TrieKey k1 => [((k1, k2), a)] -> [(k1, Elem [(k2, a)])]
+breakFst [] = []
+breakFst (((a, b),v):xs) = breakFst' a (Seq.singleton (b, v)) xs where
+ breakFst' a vs (((a', b'), v'):xs)
+ | a =? a' = breakFst' a' (vs |> (b', v')) xs
+ | otherwise = (a, Elem $ toList vs):breakFst' a' (Seq.singleton (b', v')) xs
+ breakFst' a vs [] = [(a, Elem $ toList vs)]
128 Data/TrieMap/RadixTrie.hs
@@ -0,0 +1,128 @@
+{-# LANGUAGE BangPatterns, UnboxedTuples, TypeFamilies, MagicHash, FlexibleInstances #-}
+
+module Data.TrieMap.RadixTrie () where
+
+import Data.TrieMap.TrieKey
+import Data.TrieMap.Sized
+
+import Control.Applicative
+import Control.Monad
+
+import Foreign.Storable
+
+import Data.Maybe
+import Data.Monoid
+import Data.Ord
+import Data.Foldable (foldr, foldl)
+import Data.Vector.Generic hiding (Vector, cmp, foldl, foldr)
+import Data.Vector (Vector)
+import qualified Data.Vector as V
+import qualified Data.Vector.Storable as S
+import Data.Traversable
+import Data.Word
+
+import Data.TrieMap.RadixTrie.Slice
+import Data.TrieMap.RadixTrie.Edge
+
+import Prelude hiding (length, and, zip, zipWith, foldr, foldl)
+
+instance TrieKey k => TrieKey (Vector k) where
+ ks =? ls = length ks == length ls && and (zipWith (=?) ks ls)
+ ks `cmp` ls = V.foldr (\ (k, l) z -> (k `cmp` l) `mappend` z) (comparing length ks ls) (zip ks ls)
+
+ newtype TrieMap (Vector k) a = Radix (MEdge Vector k a)
+ newtype Hole (Vector k) a = Hole (EdgeLoc Vector k a)
+
+ emptyM = Radix Nothing
+ singletonM ks a = Radix (Just (singletonEdge (v2S ks) a))
+ getSimpleM (Radix Nothing) = Null
+ getSimpleM (Radix (Just e)) = getSimpleEdge e
+ sizeM (Radix m) = getSize# m
+ lookupM ks (Radix m) = m >>= lookupEdge ks
+
+ fmapM f (Radix m) = Radix (mapEdge f <$> m)
+ mapMaybeM f (Radix m) = Radix (m >>= mapMaybeEdge f)
+ mapEitherM f (Radix e) = both Radix Radix (mapEitherMaybe (mapEitherEdge f)) e
+ traverseM f (Radix m) = Radix <$> traverse (traverseEdge f) m
+
+ foldrM f (Radix m) z = foldr (foldrEdge f) z m
+ foldlM f (Radix m) z = foldl (foldlEdge f) z m
+
+ unionM f (Radix m1) (Radix m2) = Radix (unionMaybe (unionEdge f) m1 m2)
+ isectM f (Radix m1) (Radix m2) = Radix (isectMaybe (isectEdge f) m1 m2)
+ diffM f (Radix m1) (Radix m2) = Radix (diffMaybe (diffEdge f) m1 m2)
+
+ isSubmapM (<=) (Radix m1) (Radix m2) = subMaybe (isSubEdge (<=)) m1 m2
+
+ singleHoleM ks = Hole (singleLoc (v2S ks))
+ searchM ks (Radix (Just e)) = case searchEdge (v2S ks) e Root of
+ (a, loc) -> (# a, Hole loc #)
+ searchM ks _ = (# Nothing, singleHoleM ks #)
+ indexM i (Radix (Just e)) = case indexEdge i e Root of
+ (# i', a, loc #) -> (# i', a, Hole loc #)
+ indexM _ (Radix Nothing) = indexFail ()
+
+ assignM a (Hole loc) = Radix (fillHoleEdge a loc)
+
+ extractHoleM (Radix (Just e)) = do
+ (a, loc) <- extractEdgeLoc e Root
+ return (a, Hole loc)
+ extractHoleM _ = mzero
+
+ beforeM a (Hole loc) = Radix (beforeEdge a loc)
+ afterM a (Hole loc) = Radix (afterEdge a loc)
+
+ unifyM ks1 a1 ks2 a2 = either (Left . Hole) (Right . Radix . Just) (unifyEdge (v2S ks1) a1 (v2S ks2) a2)
+
+type WordVec = S.Vector Word
+
+vZipWith :: (Storable a, Storable b) => (a -> b -> c) -> S.Vector a -> S.Vector b -> Vector c
+vZipWith f xs ys = V.zipWith f (convert xs) (convert ys)
+
+instance TrieKey (S.Vector Word) where
+ ks =? ls = length ks == length ls && and (vZipWith (=?) ks ls)
+ ks `cmp` ls = V.foldr (\ (k, l) z -> (k `cmp` l) `mappend` z) (comparing length ks ls) (vZipWith (,) ks ls)
+
+ newtype TrieMap WordVec a = WRadix (MEdge S.Vector Word a)
+ newtype Hole WordVec a = WHole (EdgeLoc S.Vector Word a)
+
+ emptyM = WRadix Nothing
+ singletonM ks a = WRadix (Just (singletonEdge (v2S ks) a))
+ getSimpleM (WRadix Nothing) = Null
+ getSimpleM (WRadix (Just e)) = getSimpleEdge e
+ sizeM (WRadix m) = getSize# m
+ lookupM ks (WRadix m) = m >>= lookupEdge ks
+
+ fmapM f (WRadix m) = WRadix (mapEdge f <$> m)
+ mapMaybeM f (WRadix m) = WRadix (m >>= mapMaybeEdge f)
+ mapEitherM f (WRadix e) = both WRadix WRadix (mapEitherMaybe (mapEitherEdge f)) e
+ traverseM f (WRadix m) = WRadix <$> traverse (traverseEdge f) m
+
+ foldrM f (