diff --git a/Data/TrieMap.hs b/Data/TrieMap.hs --- a/Data/TrieMap.hs +++ b/Data/TrieMap.hs @@ -4,6 +4,20 @@ -- * Map type TKey, TMap, + -- * Location type + TLocation, + -- ** Components + key, + before, + after, + -- ** Locations in maps + search, + index, + minLocation, + maxLocation, + -- ** Building maps + assign, + clear, -- * Operators (!), (\\), @@ -16,12 +30,12 @@ findWithDefault, -- * Construction empty, --- showMap, singleton, -- ** Insertion insert, insertWith, insertWithKey, + insertLookupWithKey, -- ** Delete/Update delete, adjust, @@ -57,8 +71,7 @@ -- ** Traverse traverseWithKey, -- ** Fold - fold, - foldWithKey, +-- fold, foldrWithKey, foldlWithKey, -- * Conversion @@ -89,6 +102,12 @@ -- * Submap isSubmapOf, isSubmapOfBy, + -- * Indexed + lookupIndex, + findIndex, + elemAt, + updateAt, + deleteAt, -- * Min/Max findMin, findMax, @@ -115,7 +134,6 @@ import Data.TrieMap.Sized import Control.Applicative hiding (empty) -import Control.Arrow import Control.Monad import Data.Maybe hiding (mapMaybe) import Data.Monoid(Monoid(..), First(..), Last(..)) @@ -137,15 +155,23 @@ mempty = empty mappend = union --- | The empty map. +-- | 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 --- | A map with a single element. +-- | /O(1)/. A map with a single element. singleton :: TKey k => k -> a -> TMap k a singleton k a = insert k a empty --- | Is the map empty? +-- | /O(1)/. Is the map empty? null :: TKey k => TMap k a -> Bool null (TMap m) = nullM m @@ -167,237 +193,878 @@ -- | The expression @('alter' f k map)@ alters the value @x@ at @k@, or absence thereof. -- 'alter' can be used to insert, delete, or update a value in a 'TMap'. In short: -- @'lookup' k ('alter' f k m) = f ('lookup' k m)@. +{-# INLINE alter #-} alter :: TKey k => (Maybe a -> Maybe a) -> k -> TMap k a -> TMap k a -alter f k (TMap m) = TMap (alterM elemSize (fmap Elem . f . fmap getElem) (toRep k) m) - -extract :: (TKey k, MonadPlus m) => (k -> a -> m (x, Maybe a)) -> TMap k a -> m (x, TMap k a) -extract f m = unwrapMonad (extractA (WrapMonad .: f) m) - --- | Projects information out of, and modifies or deletes, an individual association pair, --- alternating over all associations in the map. --- --- If @assocs m == [(k1, a1), ..., (kn, an)]@, then --- --- > extract f m = let upd k (x, maybeA) = (x, alter (const maybeA) k m) in --- > (upd k1 <$> f kn an) <|> ... <|> (upd kn <$> f kn an) --- --- This generalizes a large number of operations, including --- --- > minViewWithKey == getFirst (extract (\ k a -> return ((k, a), Nothing))) --- > updateMaxWithKey f m == maybe m snd (getLast (extract (\ k a -> return ((), f k a)) m)) --- --- In addition, --- --- > getFirst (extract (\ k a -> if p k a then return ((k, a), Nothing) else mzero) m) --- --- finds and removes the first association pair satisfying the predicate |p|. -extractA :: (TKey k, Alternative f) => (k -> a -> f (x, Maybe a)) -> TMap k a -> f (x, TMap k a) -extractA f (TMap m) = fmap TMap <$> extractM elemSize (\ k (Elem a) -> fmap (fmap (fmap Elem)) (f (fromRep k) a)) m +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 = alter f' k where - f' = Just . maybe a (f 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 = alter (const Nothing) +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 = updateWithKey (Just .: f) +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 = alter (>>= f) +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 = update (f k) k - -fold :: TKey k => (a -> b -> b) -> b -> TMap k a -> b -fold = foldWithKey . const +updateWithKey f k m = case search k m of + (Nothing, _) -> m + (Just a, hole) -> fillHole (f k a) hole -foldWithKey, foldrWithKey :: TKey k => (k -> a -> b -> b) -> b -> TMap k a -> b -foldWithKey f z (TMap m) = foldWithKeyM (\ k (Elem a) -> f (fromRep k) a) m z -foldrWithKey = foldWithKey +-- | 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 elemSize (\ k (Elem a) -> Elem <$> f (fromRep k) a) m +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 = fmap +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 elemSize (\ k (Elem a) -> Elem (f (fromRep k) a)) m) +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 f m = fromList [(f k, a) | (k, a) <- assocs m] +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 elemSize f' m1 m2) where +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 elemSize f' m1 m2) where +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, (\\) :: TKey k => TMap k a -> TMap k b -> TMap k a +-- | 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 elemSize f' m1 m2) where +differenceWithKey f (TMap m1) (TMap m2) = TMap (diffM f' m1 m2) where f' k (Elem a) (Elem b) = Elem <$> f (fromRep k) a b -minView, maxView :: TKey k => TMap k a -> Maybe (a, TMap k a) -minView m = first snd <$> minViewWithKey m -maxView m = first snd <$> maxViewWithKey m +-- | 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 -findMin, findMax :: TKey k => TMap k a -> (k, a) +-- | 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 -deleteMin, deleteMax :: TKey k => TMap k a -> TMap k a +-- | 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) -updateMin, updateMax :: TKey k => (a -> Maybe a) -> TMap k a -> TMap k a +-- | 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 -updateMinWithKey, updateMaxWithKey :: TKey k => (k -> a -> Maybe a) -> TMap k a -> TMap k a -updateMinWithKey f m = maybe m snd (getFirst (extract (\ k a -> return ((), f k a)) m)) -updateMaxWithKey f m = maybe m snd (getLast (extract (\ k a -> return ((), f k a)) m)) +{-# 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) -deleteFindMin, deleteFindMax :: TKey k => TMap k a -> ((k, a), TMap k a) +-- | 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) -minViewWithKey, maxViewWithKey :: TKey k => TMap k a -> Maybe ((k, a), TMap k a) -minViewWithKey = getFirst . extract (\ k a -> return ((k, a), Nothing)) -maxViewWithKey = getLast . extract (\ k a -> return ((k, a), Nothing)) +{-# 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 = fmap snd . assocs +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 = fmap fst . assocs +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 -> foldWithKey (curry c) n m) +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 elemSize elemSize f' m of +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 elemSize (\ k (Elem a) -> Elem <$> f (fromRep k) a) m) +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 (TMap m) = case splitLookupM elemSize f (toRep k) m of - (# mL, x, mR #) -> (TMap mL, x, TMap mR) - where f (Elem x) = (# Nothing, Just x, Nothing #) +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 -fromList, fromAscList :: TKey k => [(k, a)] -> TMap k a +-- | 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 -fromListWith, fromAscListWith :: TKey k => (a -> a -> a) -> [(k, a)] -> TMap k a +-- | 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 -fromListWithKey, fromAscListWithKey :: TKey k => (k -> a -> a -> a) -> [(k, a)] -> TMap k a -fromListWithKey f xs = TMap (fromListM elemSize (\ k (Elem a) (Elem b) -> Elem (f (fromRep k) a b)) [(toRep k, Elem a) | (k, a) <- xs]) -fromAscListWithKey f xs = TMap (fromAscListM elemSize (\ k (Elem a) (Elem b) -> Elem (f (fromRep k) a b)) [(toRep k, Elem a) | (k, a) <- xs]) +-- | 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 elemSize [(toRep k, Elem a) | (k, a) <- xs]) +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) = sizeM elemSize m +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)) diff --git a/Data/TrieMap/Applicative.hs b/Data/TrieMap/Applicative.hs --- a/Data/TrieMap/Applicative.hs +++ b/Data/TrieMap/Applicative.hs @@ -7,8 +7,6 @@ import Data.Monoid hiding (Dual) -newtype Id a = Id {unId :: a} - instance Functor First where fmap f (First m) = First (fmap f m) @@ -23,13 +21,6 @@ return = Last . return Last m >>= k = Last (m >>= getLast . k) -instance Applicative Id where - pure = Id - Id f <*> Id x = Id (f x) - -instance Functor Id where - fmap f (Id x) = Id (f x) - instance MonadPlus First where mzero = mempty mplus = mappend @@ -63,15 +54,15 @@ empty = mempty (<|>) = mplus -newtype Dual f a = Dual {runDual :: f a} - -instance Functor f => Functor (Dual f) where - fmap f (Dual x) = Dual (fmap f x) +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 x = Dual (f <*> x) + Dual f <*> Dual a = Dual (a <**> f) + Dual f *> Dual g = Dual (g <* f) + Dual f <* Dual g = Dual (g *> f) -instance Alternative f => Alternative (Dual f) where - empty = Dual empty - Dual a <|> Dual b = Dual (b <|> a) +instance MonadPlus m => MonadPlus (DualPlus m) where + mzero = DualPlus mzero + DualPlus m `mplus` DualPlus k = DualPlus (k `mplus` m) diff --git a/Data/TrieMap/Class.hs b/Data/TrieMap/Class.hs --- a/Data/TrieMap/Class.hs +++ b/Data/TrieMap/Class.hs @@ -13,6 +13,7 @@ import Prelude hiding (foldr) newtype TMap k a = TMap {getTMap :: TrieMap (Rep k) (Elem a)} + newtype TSet a = TSet (TMap a ()) class (Repr k, TrieKey (Rep k)) => TKey k @@ -23,7 +24,7 @@ fmap = fmapDefault instance TKey k => Foldable (TMap k) where - foldr f z (TMap m) = foldWithKeyM (\ _ (Elem a) -> f a) m z + foldr f z (TMap m) = foldrWithKeyM (\ _ (Elem a) -> f a) m z instance TKey k => Traversable (TMap k) where - traverse f (TMap m) = TMap <$> traverseWithKeyM elemSize (\ _ (Elem a) -> Elem <$> f a) m + traverse f (TMap m) = TMap <$> traverseWithKeyM (\ _ (Elem a) -> Elem <$> f a) m diff --git a/Data/TrieMap/IntMap.hs b/Data/TrieMap/IntMap.hs --- a/Data/TrieMap/IntMap.hs +++ b/Data/TrieMap/IntMap.hs @@ -1,11 +1,12 @@ {-# 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 (Applicative(..), Alternative(..), (<$>)) +import Control.Applicative +import Control.Monad hiding (join) import Data.Bits import Data.Maybe hiding (mapMaybe) @@ -37,31 +38,86 @@ type Prefix = Word32 type Mask = Word32 type Key = Word32 -type Size = Int +type Size = Int# +data Path a = Root + | LeftBin !Prefix !Mask !(Path a) !(TrieMap Word32 a) + | RightBin !Prefix !Mask !(TrieMap Word32 a) !(Path a) + instance TrieKey Word32 where data TrieMap Word32 a = Nil - | Tip {-# UNPACK #-} !Size {-# UNPACK #-} !Key a - | Bin {-# UNPACK #-} !Size {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(TrieMap Word32 a) !(TrieMap Word32 a) + | Tip !Size !Key a + | Bin !Size !Prefix !Mask !(TrieMap Word32 a) !(TrieMap Word32 a) + data Hole Word32 a = Hole !Key !(Path a) emptyM = Nil singletonM = singleton nullM = null - sizeM _ = size + sizeM = size lookupM = lookup - alterM = alter - alterLookupM = alterLookup traverseWithKeyM = traverseWithKey - foldWithKeyM = foldr + foldrWithKeyM = foldr foldlWithKeyM = foldl + mapWithKeyM = mapWithKey mapMaybeM = mapMaybe mapEitherM = mapEither - splitLookupM = splitLookup unionM = unionWithKey isectM = intersectionWithKey diffM = differenceWithKey - extractM s f = extract s f +-- extractM f = extract f isSubmapM = isSubmapOfBy + + singleHoleM k = Hole k Root + keyM (Hole k _) = k + 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 = onUnboxed (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 _ = (# error err, error err, error err #) where + err = "Error: empty trie" + 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 (singleton 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 + + clearM (Hole _ path) = clear Nil path where + clear t Root = t + clear t (LeftBin p m path r) = clear (bin p m t r) path + clear t (RightBin p m l path) = clear (bin p m l t) path +branchHole :: Key -> Prefix -> Path a -> TrieMap Word32 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 :: Word32 -> Nat natFromInt = id @@ -79,11 +135,10 @@ shiftRL x i = shiftR x (fromIntegral i) -- #endif - -size :: TrieMap Word32 a -> Int -size Nil = 0 -size (Tip s _ _) = s -size (Bin s _ _ _ _) = s +size :: TrieMap Word32 a -> Int# +size Nil = 0# +size (Tip sz _ _) = sz +size (Bin sz _ _ _ _) = sz null :: TrieMap Word32 a -> Bool null Nil = True @@ -95,44 +150,17 @@ | k == kx = Just x lookup _ _ = Nothing -singleton :: Sized a -> Key -> a -> TrieMap Word32 a -singleton s k a = Tip (s a) k a - -singletonMaybe :: Sized a -> Key -> Maybe a -> TrieMap Word32 a -singletonMaybe s k = maybe Nil (singleton s k) - -alter :: Sized a -> (Maybe a -> Maybe a) -> Key -> TrieMap Word32 a -> TrieMap Word32 a -alter s f k t = case t of - Bin _ p m l r - | nomatch k p m -> join k (singletonMaybe s k (f Nothing)) p t - | zero k m -> bin p m (alter s f k l) r - | otherwise -> bin p m l (alter s f k r) - Tip _ ky y - | k == ky -> singletonMaybe s k (f (Just y)) - | Just x <- f Nothing - -> join k (Tip (s x) k x) ky t - | otherwise -> t - Nil -> singletonMaybe s k (f Nothing) +singleton :: Sized a => Key -> a -> TrieMap Word32 a +singleton k a = Tip (getSize# a) k a -alterLookup :: Sized a -> (Maybe a -> (# x, Maybe a #)) -> Key -> TrieMap Word32 a -> (# x, TrieMap Word32 a #) -alterLookup s f k t = case t of - Bin _ p m l r - | nomatch k p m - -> onUnboxed (\ v -> join k (singletonMaybe s k v) p t) f Nothing - | zero k m - -> onUnboxed (\ l' -> bin p m l' r) (alterLookup s f k) l - | otherwise - -> onUnboxed (\ r' -> bin p m l r') (alterLookup s f k) r - Tip _ ky y - | k == ky -> onUnboxed (singletonMaybe s k) f (Just y) - | otherwise -> onUnboxed (\ v -> join k (singletonMaybe s k v) ky t) f Nothing - Nil -> onUnboxed (singletonMaybe s k) f Nothing +singletonMaybe :: Sized a => Key -> Maybe a -> TrieMap Word32 a +singletonMaybe k = maybe Nil (singleton k) -traverseWithKey :: Applicative f => Sized b -> (Key -> a -> f b) -> TrieMap Word32 a -> f (TrieMap Word32 b) -traverseWithKey s f t = case t of +traverseWithKey :: (Applicative f, Sized b) => (Key -> a -> f b) -> TrieMap Word32 a -> f (TrieMap Word32 b) +traverseWithKey f t = case t of Nil -> pure Nil - Tip _ kx x -> singleton s kx <$> f kx x - Bin _ p m l r -> bin p m <$> traverseWithKey s f l <*> traverseWithKey s f r + Tip _ kx x -> singleton kx <$> f kx x + Bin _ p m l r -> bin p m <$> traverseWithKey f l <*> traverseWithKey f r foldr :: (Key -> a -> b -> b) -> TrieMap Word32 a -> b -> b foldr f t @@ -148,110 +176,84 @@ Tip _ k x -> flip (f k) x Nil -> id -mapMaybe :: Sized b -> (Key -> a -> Maybe b) -> TrieMap Word32 a -> TrieMap Word32 b -mapMaybe s f (Bin _ p m l r) = bin p m (mapMaybe s f l) (mapMaybe s f r) -mapMaybe s f (Tip _ kx x) = singletonMaybe s kx (f kx x) -mapMaybe _ _ _ = Nil +mapWithKey :: Sized b => (Key -> a -> b) -> TrieMap Word32 a -> TrieMap Word32 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 kx x) +mapWithKey _ _ = Nil -mapEither :: Sized b -> Sized c -> EitherMap Key a b c -> +mapMaybe :: Sized b => (Key -> a -> Maybe b) -> TrieMap Word32 a -> TrieMap Word32 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 kx x) +mapMaybe _ _ = Nil + +mapEither :: (Sized b, Sized c) => EitherMap Key a b c -> TrieMap Word32 a -> (# TrieMap Word32 b, TrieMap Word32 c #) -mapEither s1 s2 f (Bin _ p m l r) - | (# lL, lR #) <- mapEither s1 s2 f l, (# rL, rR #) <- mapEither s1 s2 f r +mapEither f (Bin _ p m l r) + | (# lL, lR #) <- mapEither f l, + (# rL, rR #) <- mapEither f r = (# bin p m lL rL, bin p m lR rR #) -mapEither s1 s2 f (Tip _ kx x) = both (singletonMaybe s1 kx) (singletonMaybe s2 kx) (f kx) x -mapEither _ _ _ _ = (# Nil, Nil #) - -splitLookup :: Sized a -> SplitMap a x -> Key -> TrieMap Word32 a -> (# TrieMap Word32 a ,Maybe x,TrieMap Word32 a #) -splitLookup s f k t@(Bin _ p m l r) - | nomatch k p m = if k>p then (# t,Nothing,Nil #) else (# Nil,Nothing,t #) - | zero k m, (# lt, found, gt #) <- splitLookup s f k l - = (# lt,found,union s gt r #) - | (# lt, found, gt #) <- splitLookup s f k r - = (# union s l lt,found,gt #) -splitLookup s f k t@(Tip _ ky y) - | k>ky = (# t,Nothing,Nil #) - | k TrieMap Word32 a -> TrieMap Word32 a -> TrieMap Word32 a -union _ Nil t = t -union _ t Nil = t -union s (Tip _ k x) t = alter s (const (Just x)) k t -union s t (Tip _ k x) = alter s (Just . fromMaybe x) k t -- right bias -union s 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 (union s l1 l2) (union s 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 (union s l1 t2) r1 - | otherwise = bin p1 m1 l1 (union s r1 t2) - - union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2 - | zero p1 m2 = bin p2 m2 (union s t1 l2) r2 - | otherwise = bin p2 m2 l2 (union s t1 r2) +mapEither f (Tip _ kx x) = both (singletonMaybe kx) (singletonMaybe kx) (f kx) x +mapEither _ _ = (# Nil, Nil #) -unionWithKey :: Sized a -> UnionFunc Key a -> TrieMap Word32 a -> TrieMap Word32 a -> TrieMap Word32 a -unionWithKey _ _ Nil t = t -unionWithKey _ _ t Nil = t -unionWithKey s f (Tip _ k x) t = alter s (maybe (Just x) (f k x)) k t -unionWithKey s f t (Tip _ k x) = alter s (maybe (Just x) (flip (f k) x)) k t -unionWithKey s f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2) +unionWithKey :: Sized a => UnionFunc Key a -> TrieMap Word32 a -> TrieMap Word32 a -> TrieMap Word32 a +unionWithKey _ Nil t = t +unionWithKey _ t Nil = t +unionWithKey f (Tip _ k x) t = alterM (maybe (Just x) (f k x)) k t +unionWithKey f t (Tip _ k x) = alterM (maybe (Just x) (flip (f k) x)) k t +unionWithKey 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 (unionWithKey s f l1 l2) (unionWithKey s f r1 r2) + | p1 == p2 = bin p1 m1 (unionWithKey f l1 l2) (unionWithKey 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 (unionWithKey s f l1 t2) r1 - | otherwise = bin p1 m1 l1 (unionWithKey s f r1 t2) + | zero p2 m1 = bin p1 m1 (unionWithKey f l1 t2) r1 + | otherwise = bin p1 m1 l1 (unionWithKey f r1 t2) union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2 - | zero p1 m2 = bin p2 m2 (unionWithKey s f t1 l2) r2 - | otherwise = bin p2 m2 l2 (unionWithKey s f t1 r2) + | zero p1 m2 = bin p2 m2 (unionWithKey f t1 l2) r2 + | otherwise = bin p2 m2 l2 (unionWithKey f t1 r2) -intersectionWithKey :: Sized c -> IsectFunc Key a b c -> TrieMap Word32 a -> TrieMap Word32 b -> TrieMap Word32 c -intersectionWithKey _ _ Nil _ = Nil -intersectionWithKey _ _ _ Nil = Nil -intersectionWithKey s f (Tip _ k x) t2 - = singletonMaybe s k (lookup (natFromInt k) t2 >>= f k x) -intersectionWithKey s f t1 (Tip _ k y) - = singletonMaybe s k (lookup (natFromInt k) t1 >>= flip (f k) y) -intersectionWithKey s f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2) +intersectionWithKey :: Sized c => IsectFunc Key a b c -> TrieMap Word32 a -> TrieMap Word32 b -> TrieMap Word32 c +intersectionWithKey _ Nil _ = Nil +intersectionWithKey _ _ Nil = Nil +intersectionWithKey f (Tip _ k x) t2 + = singletonMaybe k (lookup (natFromInt k) t2 >>= f k x) +intersectionWithKey f t1 (Tip _ k y) + = singletonMaybe k (lookup (natFromInt k) t1 >>= flip (f k) y) +intersectionWithKey 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 (intersectionWithKey s f l1 l2) (intersectionWithKey s f r1 r2) + | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2) | otherwise = Nil where intersection1 | nomatch p2 p1 m1 = Nil - | zero p2 m1 = intersectionWithKey s f l1 t2 - | otherwise = intersectionWithKey s f r1 t2 + | zero p2 m1 = intersectionWithKey f l1 t2 + | otherwise = intersectionWithKey f r1 t2 intersection2 | nomatch p1 p2 m2 = Nil - | zero p1 m2 = intersectionWithKey s f t1 l2 - | otherwise = intersectionWithKey s f t1 r2 + | zero p1 m2 = intersectionWithKey f t1 l2 + | otherwise = intersectionWithKey f t1 r2 -differenceWithKey :: Sized a -> (Key -> a -> b -> Maybe a) -> TrieMap Word32 a -> TrieMap Word32 b -> TrieMap Word32 a -differenceWithKey _ _ Nil _ = Nil -differenceWithKey _ _ t Nil = t -differenceWithKey s f t1@(Tip _ k x) t2 - = maybe t1 (singletonMaybe s k . f k x) (lookup (natFromInt k) t2) -differenceWithKey s f t (Tip _ k y) = alter s (>>= flip (f k) y) k t -differenceWithKey s f t1@(Bin _ p1 m1 l1 r1) t2@(Bin _ p2 m2 l2 r2) +differenceWithKey :: Sized a => (Key -> a -> b -> Maybe a) -> TrieMap Word32 a -> TrieMap Word32 b -> TrieMap Word32 a +differenceWithKey _ Nil _ = Nil +differenceWithKey _ t Nil = t +differenceWithKey f t1@(Tip _ k x) t2 + = maybe t1 (singletonMaybe k . f k x) (lookup (natFromInt k) t2) +differenceWithKey f t (Tip _ k y) = alterM (>>= flip (f k) y) k t +differenceWithKey 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 (differenceWithKey s f l1 l2) (differenceWithKey s f r1 r2) + | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2) | otherwise = t1 where difference1 | nomatch p2 p1 m1 = t1 - | zero p2 m1 = bin p1 m1 (differenceWithKey s f l1 t2) r1 - | otherwise = bin p1 m1 l1 (differenceWithKey s f r1 t2) + | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1 + | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2) difference2 | nomatch p1 p2 m2 = t1 - | zero p1 m2 = differenceWithKey s f t1 l2 - | otherwise = differenceWithKey s f t1 r2 + | zero p1 m2 = differenceWithKey f t1 l2 + | otherwise = differenceWithKey f t1 r2 isSubmapOfBy :: LEq a b -> LEq (TrieMap Word32 a) (TrieMap Word32 b) isSubmapOfBy (<=) t1@(Bin _ p1 m1 l1 r1) (Bin _ p2 m2 l2 r2) @@ -266,11 +268,11 @@ isSubmapOfBy _ Nil _ = True -extract :: Alternative f => Sized a -> (Key -> a -> f (x, Maybe a)) -> TrieMap Word32 a -> f (x, TrieMap Word32 a) -extract s f (Bin _ p m l r) = - fmap (\ l' -> bin p m l' r) <$> extract s f l <|> fmap (bin p m l) <$> extract s f r -extract s f (Tip _ k x) = fmap (singletonMaybe s k) <$> f k x -extract _ _ _ = empty +-- extract :: Alternative f => Sized a -> (Key -> a -> f (x, Maybe a)) -> TrieMap Word32 a -> f (x, TrieMap Word32 a) +-- extract f (Bin _ p m l r) = +-- fmap (\ l' -> bin p m l' r) <$> extract f l <|> fmap (bin p m l) <$> extract f r +-- extract f (Tip _ k x) = fmap (singletonMaybe k) <$> f k x +-- extract _ _ _ = empty mask :: Key -> Mask -> Prefix mask i m @@ -327,4 +329,4 @@ bin :: Prefix -> Mask -> TrieMap Word32 a -> TrieMap Word32 a -> TrieMap Word32 a bin _ _ l Nil = l bin _ _ Nil r = r -bin p m l r = Bin (size l + size r) p m l r +bin p m l r = Bin (size l +# size r) p m l r diff --git a/Data/TrieMap/Key.hs b/Data/TrieMap/Key.hs --- a/Data/TrieMap/Key.hs +++ b/Data/TrieMap/Key.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE TypeFamilies, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses, FlexibleContexts #-} +{-# LANGUAGE TypeFamilies, UnboxedTuples #-} module Data.TrieMap.Key (Key(..)) where @@ -10,20 +10,33 @@ instance TKey k => TrieKey (Key k) where newtype TrieMap (Key k) a = KeyMap (TrieMap (Rep k) a) + newtype Hole (Key k) a = KeyHole (Hole (Rep k) a) + emptyM = KeyMap emptyM - singletonM s (Key k) a = KeyMap (singletonM s (toRep k) a) + singletonM (Key k) a = KeyMap (singletonM (toRep k) a) nullM (KeyMap m) = nullM m + sizeM (KeyMap m) = sizeM m lookupM (Key k) (KeyMap m) = lookupM (toRep k) m - alterM s f (Key k) (KeyMap m) = KeyMap (alterM s f (toRep k) m) - alterLookupM s f (Key k) (KeyMap m) = onUnboxed KeyMap (alterLookupM s f (toRep k)) m - traverseWithKeyM s f (KeyMap m) = KeyMap <$> traverseWithKeyM s (f . Key . fromRep) m - foldWithKeyM f (KeyMap m) = foldWithKeyM (f . Key . fromRep) m + traverseWithKeyM f (KeyMap m) = KeyMap <$> traverseWithKeyM (f . Key . fromRep) m + foldrWithKeyM f (KeyMap m) = foldrWithKeyM (f . Key . fromRep) m foldlWithKeyM f (KeyMap m) = foldlWithKeyM (f . Key . fromRep) m - mapMaybeM s f (KeyMap m) = KeyMap (mapMaybeM s (f . Key . fromRep) m) - mapEitherM s1 s2 f (KeyMap m) = both KeyMap KeyMap (mapEitherM s1 s2 (f . Key . fromRep)) m - splitLookupM s f (Key k) (KeyMap m) = sides KeyMap (splitLookupM s f (toRep k)) m - unionM s f (KeyMap m1) (KeyMap m2) = KeyMap (unionM s (f . Key . fromRep) m1 m2) - isectM s f (KeyMap m1) (KeyMap m2) = KeyMap (isectM s (f . Key . fromRep) m1 m2) - diffM s f (KeyMap m1) (KeyMap m2) = KeyMap (diffM s (f . Key . fromRep) m1 m2) - extractM s f (KeyMap m) = fmap KeyMap <$> extractM s (f . Key . fromRep) m + mapWithKeyM f (KeyMap m) = KeyMap (mapWithKeyM (f . Key . fromRep) m) + mapMaybeM f (KeyMap m) = KeyMap (mapMaybeM (f . Key . fromRep) m) + mapEitherM f (KeyMap m) = both KeyMap KeyMap (mapEitherM (f . Key . fromRep)) m + unionM f (KeyMap m1) (KeyMap m2) = KeyMap (unionM (f . Key . fromRep) m1 m2) + isectM f (KeyMap m1) (KeyMap m2) = KeyMap (isectM (f . Key . fromRep) m1 m2) + diffM f (KeyMap m1) (KeyMap m2) = KeyMap (diffM (f . Key . fromRep) m1 m2) isSubmapM (<=) (KeyMap m1) (KeyMap m2) = isSubmapM (<=) m1 m2 + + singleHoleM (Key k) = KeyHole (singleHoleM (toRep k)) + keyM (KeyHole hole) = Key (fromRep (keyM hole)) + beforeM a (KeyHole hole) = KeyMap (beforeM a hole) + afterM a (KeyHole hole) = KeyMap (afterM a hole) + searchM (Key k) (KeyMap m) = onUnboxed 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) + clearM (KeyHole hole) = KeyMap (clearM hole) diff --git a/Data/TrieMap/OrdMap.hs b/Data/TrieMap/OrdMap.hs --- a/Data/TrieMap/OrdMap.hs +++ b/Data/TrieMap/OrdMap.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE UnboxedTuples, TypeFamilies, PatternGuards #-} +{-# LANGUAGE BangPatterns, UnboxedTuples, TypeFamilies, PatternGuards, MagicHash, CPP, TupleSections #-} module Data.TrieMap.OrdMap () where @@ -6,42 +6,94 @@ import Data.TrieMap.Sized import Data.TrieMap.Modifiers -import Control.Applicative (Applicative(..), Alternative(..), (<$>)) +import Control.Applicative import Control.Monad hiding (join) import Prelude hiding (lookup) +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 data TrieMap (Ordered k) a = Tip - | Bin {-# UNPACK #-} !Int k a !(OrdMap k a) !(OrdMap k a) + | 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 s (Ord k) = singleton s k + singletonM (Ord k) = singleton k nullM Tip = True nullM _ = False - sizeM _ = size + sizeM = size# lookupM (Ord k) = lookup k - alterM s f (Ord k) = alter s f k - alterLookupM s f (Ord k) = alterLookup s f k - traverseWithKeyM s f = traverseWithKey s (f . Ord) - foldWithKeyM f = foldrWithKey (f . Ord) + traverseWithKeyM f = traverseWithKey (f . Ord) + foldrWithKeyM f = foldrWithKey (f . Ord) foldlWithKeyM f = foldlWithKey (f . Ord) - mapMaybeM s f = mapMaybe s (f . Ord) - mapEitherM s1 s2 f = mapEither s1 s2 (f . Ord) - extractM s f = extract s (f . Ord) - splitLookupM s f (Ord k) = splitLookup s f k + mapWithKeyM f = mapWithKey (f . Ord) + mapMaybeM f = mapMaybe (f . Ord) + mapEitherM f = mapEither (f . Ord) isSubmapM = isSubmap - fromAscListM s f xs = fromAscList s (f . Ord) [(k, a) | (Ord k, a) <- xs] - fromDistAscListM s xs = fromDistinctAscList s [(k, a) | (Ord k, a) <- xs] - unionM _ _ Tip m2 = m2 - unionM _ _ m1 Tip = m1 - unionM s f m1 m2 = hedgeUnionWithKey s (f . Ord) (const LT) (const GT) m1 m2 - isectM s f = isect s (f . Ord) - diffM _ _ Tip _ = Tip - diffM _ _ m1 Tip = m1 - diffM s f m1 m2 = hedgeDiffWithKey s (f . Ord) (const LT) (const GT) m1 m2 + fromAscListM f xs = fromAscList (f . Ord) [(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 = hedgeUnionWithKey (f . Ord) (const LT) (const GT) m1 m2 + isectM f = isect (f . Ord) + diffM _ Tip _ = Tip + diffM _ m1 Tip = m1 + diffM f m1 m2 = hedgeDiffWithKey (f . Ord) (const LT) (const GT) m1 m2 + + singleHoleM (Ord k) = Empty k Root + keyM (Empty k _) = Ord k + keyM (Full k _ _ _) = Ord k + 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 _ _ _ = (# error err, error err, error err #) where + err = "Error: empty trie" + 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 (singleton k x) path + assignM x (Full k path l r) = rebuild (join k x l r) path + clearM (Empty _ path) = rebuild Tip path + clearM (Full _ path l r) = rebuild (merge l r) path +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 @@ -49,30 +101,12 @@ GT -> lookup k r lookup _ _ = Nothing -alter :: Ord k => Sized a -> (Maybe a -> Maybe a) -> k -> OrdMap k a -> OrdMap k a -alter s f k Tip = case f Nothing of - Nothing -> Tip - Just x -> singleton s k x -alter s f k (Bin _ kx x l r) = case compare k kx of - LT -> balance s kx x (alter s f k l) r - EQ -> case f (Just x) of - Nothing -> glue s l r - Just x' -> balance s k x' l r - GT -> balance s kx x l (alter s f k r) - -alterLookup :: Ord k => Sized a -> (Maybe a -> (# z, Maybe a #)) -> k -> OrdMap k a -> (# z, OrdMap k a #) -alterLookup s f k Tip = onUnboxed (maybe Tip (singleton s k)) f Nothing -alterLookup s f k (Bin _ kx x l r) = case compare k kx of - LT -> onUnboxed (\ l' -> balance s kx x l' r) (alterLookup s f k) l - EQ -> onUnboxed (maybe (glue s l r) (\ x' -> balance s k x' l r)) f (Just x) - GT -> onUnboxed (balance s kx x l) (alterLookup s f k) r - -singleton :: Sized a -> k -> a -> OrdMap k a -singleton s k a = Bin (s a) k a Tip Tip +singleton :: Sized a => k -> a -> OrdMap k a +singleton k a = Bin (getSize# a) k a Tip Tip -traverseWithKey :: Applicative f => Sized b -> (k -> a -> f b) -> OrdMap k a -> f (OrdMap k b) -traverseWithKey _ _ Tip = pure Tip -traverseWithKey s f (Bin _ k a l r) = balance s k <$> f k a <*> traverseWithKey s f l <*> traverseWithKey s f r +traverseWithKey :: (Applicative f, Sized b) => (k -> a -> f b) -> OrdMap k a -> f (OrdMap k b) +traverseWithKey _ Tip = pure Tip +traverseWithKey f (Bin _ k a l r) = balance k <$> f k a <*> traverseWithKey f l <*> traverseWithKey f r foldrWithKey :: (k -> a -> b -> b) -> OrdMap k a -> b -> b foldrWithKey _ Tip = id @@ -82,41 +116,45 @@ foldlWithKey _ Tip = id foldlWithKey f (Bin _ k a l r) = foldlWithKey f r . flip (f k) a . foldlWithKey f l -mapMaybe :: Ord k => Sized b -> (k -> a -> Maybe b) -> OrdMap k a -> OrdMap k b -mapMaybe _ _ Tip = Tip -mapMaybe s f (Bin _ k a l r) = joinMaybe s k (f k a) (mapMaybe s f l) (mapMaybe s f r) +mapWithKey :: (Ord k, Sized b) => (k -> a -> b) -> OrdMap k a -> OrdMap k b +mapWithKey f (Bin _ k a l r) = join k (f k a) (mapWithKey f l) (mapWithKey f r) +mapWithKey _ _ = Tip -mapEither :: Ord k => Sized b -> Sized c -> EitherMap k a b c -> +mapMaybe :: (Ord k, Sized b) => (k -> a -> Maybe b) -> OrdMap k a -> OrdMap k b +mapMaybe f (Bin _ k a l r) = joinMaybe k (f k a) (mapMaybe f l) (mapMaybe f r) +mapMaybe _ _ = Tip + +mapEither :: (Ord k, Sized b, Sized c) => EitherMap k a b c -> OrdMap k a -> (# OrdMap k b, OrdMap k c #) -mapEither _ _ _ Tip = (# Tip, Tip #) -mapEither s1 s2 f (Bin _ k a l r) +mapEither f (Bin _ k a l r) | (# aL, aR #) <- f k a, - (# lL, lR #) <- mapEither s1 s2 f l, - (# rL, rR #) <- mapEither s1 s2 f r - = (# joinMaybe s1 k aL lL rL, joinMaybe s2 k aR lR rR #) + (# lL, lR #) <- mapEither f l, + (# rL, rR #) <- mapEither f r + = (# joinMaybe k aL lL rL, joinMaybe k aR lR rR #) +mapEither _ _ = (# Tip, Tip #) -splitLookup :: Ord k => Sized a -> SplitMap a x -> k -> OrdMap k a -> (# OrdMap k a, Maybe x, OrdMap k a #) -splitLookup s f k m = case m of +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 -> case splitLookup s f k l of - (# lL, ans, lR #) -> (# lL, ans, join s kx x lR r #) + LT -> case splitLookup f k l of + (# lL, ans, lR #) -> (# lL, ans, join kx x lR r #) EQ -> case f x of - (# xL, ans, xR #) -> (# maybe l (\ xL -> insertMax s kx xL l) xL, ans, - maybe r (\ xR -> insertMin s kx xR r) xR #) - GT -> case splitLookup s f k r of - (# rL, ans, rR #) -> (# join s kx x l rL, ans, rR #) + (# xL, ans, xR #) -> (# maybe l (\ xL -> insertMax kx xL l) xL, ans, + maybe r (\ xR -> insertMin kx xR r) xR #) + GT -> case splitLookup f k r of + (# rL, ans, rR #) -> (# join kx x l rL, ans, rR #) -isSubmap :: Ord k => LEq a b -> LEq (OrdMap k a) (OrdMap k b) +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 splitLookup (const 1) (\ x -> (# Nothing, Just x, Nothing #)) kx t of +isSubmap (<=) (Bin _ kx x l r) t = case splitLookup (\ x -> (# Nothing, Just (Elem x), Nothing #)) kx t of (# lt, found, gt #) -> case found of Nothing -> False - Just y -> x <= y && isSubmap (<=) l lt && isSubmap (<=) r gt + Just (Elem y) -> x <= y && isSubmap (<=) l lt && isSubmap (<=) r gt -fromAscList :: Eq k => Sized a -> (k -> a -> a -> a) -> [(k, a)] -> OrdMap k a -fromAscList s f xs = fromDistinctAscList s (combineEq xs) where +fromAscList :: (Eq k, Sized a) => (k -> a -> a -> a) -> [(k, a)] -> OrdMap k a +fromAscList f xs = fromDistinctAscList (combineEq xs) where combineEq (x:xs) = combineEq' x xs combineEq [] = [] @@ -125,16 +163,16 @@ | kz == kx = combineEq' (kx, f kx xx zz) xs | otherwise = (kz,zz):combineEq' x xs -fromDistinctAscList :: Sized a -> [(k, a)] -> OrdMap k a -fromDistinctAscList s xs = build const (length xs) 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 s k4 x4 (bin s k2 x2 (singleton s k1 x1) (singleton s k3 x3)) (singleton s k5 x5)) xx - _ -> error "fromDistinctAscList build" + ((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 @@ -142,19 +180,19 @@ 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 s k x l r) zs + buildB l k x c r zs = c (bin k x l r) zs -hedgeUnionWithKey :: Ord k - => Sized a -> (k -> a -> a -> Maybe a) +hedgeUnionWithKey :: (Ord k, Sized a) + => (k -> a -> a -> Maybe a) -> (k -> Ordering) -> (k -> Ordering) -> OrdMap k a -> OrdMap k a -> OrdMap k a -hedgeUnionWithKey _ _ _ _ t1 Tip +hedgeUnionWithKey _ _ _ t1 Tip = t1 -hedgeUnionWithKey s _ cmplo cmphi Tip (Bin _ kx x l r) - = join s kx x (filterGt s cmplo l) (filterLt s cmphi r) -hedgeUnionWithKey s f cmplo cmphi (Bin _ kx x l r) t2 - = joinMaybe s kx newx (hedgeUnionWithKey s f cmplo cmpkx l lt) - (hedgeUnionWithKey s f cmpkx cmphi r gt) +hedgeUnionWithKey _ cmplo cmphi Tip (Bin _ kx x l r) + = join kx x (filterGt cmplo l) (filterLt cmphi r) +hedgeUnionWithKey f cmplo cmphi (Bin _ kx x l r) t2 + = joinMaybe kx newx (hedgeUnionWithKey f cmplo cmpkx l lt) + (hedgeUnionWithKey f cmpkx cmphi r gt) where cmpkx k = compare kx k lt = trim cmplo cmpkx t2 @@ -163,20 +201,20 @@ Nothing -> Just x Just (_,y) -> f kx x y -filterGt :: Ord k => Sized a -> (k -> Ordering) -> OrdMap k a -> OrdMap k a -filterGt _ _ Tip = Tip -filterGt s cmp (Bin _ kx x l r) +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 s kx x (filterGt s cmp l) r - GT -> filterGt s cmp r + 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 s cmp (Bin _ kx x l 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 s cmp l - GT -> join s kx x l (filterLt s cmp r) + 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 @@ -193,156 +231,165 @@ 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) + 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 c -> IsectFunc k a b c -> OrdMap k a -> OrdMap k b -> OrdMap k c -isect s f t1@Bin{} (Bin _ k2 x2 l2 r2) - | (# lt, found, gt #) <- splitLookup (const 1) (\ x -> (# Nothing, Just x, Nothing #)) k2 t1 - = let tl = isect s f lt l2 - tr = isect s f gt r2 - in joinMaybe s k2 (found >>= \ x1' -> f k2 x1' x2) tl tr -isect _ _ _ _ = Tip +isect :: (Ord k, Sized a, Sized b, Sized c) => IsectFunc k a b c -> OrdMap k a -> OrdMap k b -> OrdMap k c +isect f t1@Bin{} (Bin _ k2 x2 l2 r2) + | (# found, hole #) <- search k2 Root t1 + = let tl = isect f (beforeM Nothing hole) l2 + tr = isect f (afterM Nothing hole) r2 + in joinMaybe k2 (found >>= \ x1' -> f k2 x1' x2) tl tr +isect _ _ _ = Tip -hedgeDiffWithKey :: Ord k - => Sized a -> (k -> a -> b -> Maybe a) +hedgeDiffWithKey :: (Ord k, Sized a) + => (k -> a -> b -> Maybe a) -> (k -> Ordering) -> (k -> Ordering) -> OrdMap k a -> OrdMap k b -> OrdMap k a -hedgeDiffWithKey _ _ _ _ Tip _ +hedgeDiffWithKey _ _ _ Tip _ = Tip -hedgeDiffWithKey s _ cmplo cmphi (Bin _ kx x l r) Tip - = join s kx x (filterGt s cmplo l) (filterLt s cmphi r) -hedgeDiffWithKey s f cmplo cmphi t (Bin _ kx x l r) +hedgeDiffWithKey _ cmplo cmphi (Bin _ kx x l r) Tip + = join kx x (filterGt cmplo l) (filterLt cmphi r) +hedgeDiffWithKey f cmplo cmphi t (Bin _ kx x l r) = case found of - Nothing -> merge s tl tr + Nothing -> merge tl tr Just (ky,y) -> case f ky y x of - Nothing -> merge s tl tr - Just z -> join s ky z tl tr + 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 = hedgeDiffWithKey s f cmplo cmpkx lt l - tr = hedgeDiffWithKey s f cmpkx cmphi gt r - -joinMaybe :: Ord k => Sized a -> k -> Maybe a -> OrdMap k a -> OrdMap k a -> OrdMap k a -joinMaybe s kx = maybe (merge s) (join s kx) + tl = hedgeDiffWithKey f cmplo cmpkx lt l + tr = hedgeDiffWithKey f cmpkx cmphi gt r -join :: Ord k => Sized a -> k -> a -> OrdMap k a -> OrdMap k a -> OrdMap k a -join s kx x Tip r = insertMin s kx x r -join s kx x l Tip = insertMax s kx x l -join s kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz) - | delta*sizeL <= sizeR = balance s kz z (join s kx x l lz) rz - | delta*sizeR <= sizeL = balance s ky y ly (join s kx x ry r) - | otherwise = bin s kx x l 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 s kx x t +insertMax,insertMin :: Sized a => k -> a -> OrdMap k a -> OrdMap k a +insertMax kx x t = case t of - Tip -> singleton s kx x + Tip -> singleton kx x Bin _ ky y l r - -> balance s ky y l (insertMax s kx x r) + -> balance ky y l (insertMax kx x r) -insertMin s kx x t +insertMin kx x t = case t of - Tip -> singleton s kx x + Tip -> singleton kx x Bin _ ky y l r - -> balance s ky y (insertMin s kx x 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 s l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) - | delta*sizeL <= sizeR = balance s ky y (merge s l ly) ry - | delta*sizeR <= sizeL = balance s kx x lx (merge s rx r) - | otherwise = glue s l r +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 s l r - | size l > size r = let (f,l') = deleteFindMax s (\ k a -> (balance s k a, Nothing)) l in f l' r - | otherwise = let (f,r') = deleteFindMin s (\ k a -> (balance s k a, Nothing)) r in f l r' - -extract :: Alternative t => Sized a -> (k -> a -> t (z, Maybe a)) -> OrdMap k a -> t (z, OrdMap k a) -extract s f t = case t of - Bin _ k x l r -> - fmap (\ l' -> balance s k x l' r) <$> extract s f l <|> - fmap (maybe (glue s l r) (\ x' -> balance s k x' l r)) <$> f k x <|> - fmap (balance s k x l) <$> extract s f r - Tip -> empty +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 = case deleteFindMax (\ k a -> (# balance k a, Nothing #)) l of + (# f, l' #) -> f l' r + | otherwise = case deleteFindMin (\ k a -> (# balance k a, Nothing #)) r of + (# f, r' #) -> f l r' -deleteFindMin :: Sized a -> (k -> a -> (x, Maybe a)) -> OrdMap k a -> (x, OrdMap k a) -deleteFindMin s f t +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 -> let (ans, x') = f k x in (ans, maybe r (\ y' -> bin s k y' Tip r) x') - Bin _ k x l r -> let (km,l') = deleteFindMin s f l in (km,balance s k x l' r) - Tip -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip) + Bin _ k x Tip r -> onUnboxed (maybe r (\ y' -> bin k y' Tip r)) (f k) x + Bin _ k x l r -> onUnboxed (\ l' -> balance k x l' r) (deleteFindMin f) l + _ -> (# error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip #) -deleteFindMax :: Sized a -> (k -> a -> (x, Maybe a)) -> OrdMap k a -> (x, OrdMap k a) -deleteFindMax s f t +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 -> let (ans, x') = f k x in (ans, maybe l (\ y -> bin s k y l Tip) x') - Bin _ k x l r -> let (km,r') = deleteFindMax s f r in (km,balance s k x l r') - Tip -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip) - -delta,ratio :: Int -delta = 5 -ratio = 2 + Bin _ k x l Tip -> onUnboxed (maybe l (\ y -> bin k y l Tip)) (f k) x + Bin _ k x l r -> onUnboxed (balance k x l) (deleteFindMax f) r + Tip -> (# error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip #) -size :: OrdMap k a -> Int -size Tip = 0 -size (Bin s _ _ _ _) = s +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 s k x l r - | sizeL + sizeR <= 1 = Bin sizeX k x l r - | sizeR >= delta*sizeL = rotateL s k x l r - | sizeL >= delta*sizeR = rotateR s k x l r - | otherwise = Bin sizeX k x l r +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 - sizeL = size l - sizeR = size r - sizeX = sizeL + sizeR + s x + !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 s k x l r@(Bin _ _ _ ly ry) - | size ly < ratio*size ry = singleL s k x l r - | otherwise = doubleL s k x l r -rotateL _ _ _ _ Tip = error "rotateL Tip" +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 s k x l@(Bin _ _ _ ly ry) r - | size ry < ratio*size ly = singleR s k x l r - | otherwise = doubleR s k x l r -rotateR _ _ _ Tip _ = error "rotateR 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 s k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin s k2 x2 (bin s k1 x1 t1 t2) t3 -singleL s k1 x1 t1 Tip = bin s k1 x1 t1 Tip -singleR s k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin s k2 x2 t1 (bin s k1 x1 t2 t3) -singleR s k1 x1 Tip t2 = bin s k1 x1 Tip t2 +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 s k1 x1 t1 (Bin _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = bin s k3 x3 (bin s k1 x1 t1 t2) (bin s k2 x2 t3 t4) -doubleL s k1 x1 t1 t2 = singleL s k1 x1 t1 t2 -doubleR s k1 x1 (Bin _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = bin s k3 x3 (bin s k2 x2 t1 t2) (bin s k1 x1 t3 t4) -doubleR s k1 x1 t1 t2 = singleR s k1 x1 t1 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 s k x l r - = Bin (size l + size r + s x) k x l r +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 diff --git a/Data/TrieMap/ProdMap.hs b/Data/TrieMap/ProdMap.hs --- a/Data/TrieMap/ProdMap.hs +++ b/Data/TrieMap/ProdMap.hs @@ -2,12 +2,12 @@ module Data.TrieMap.ProdMap () where +import Data.TrieMap.Sized import Data.TrieMap.TrieKey import Data.TrieMap.Applicative import Control.Applicative -import Data.Maybe import Data.Foldable import Data.Sequence ((|>)) @@ -15,40 +15,54 @@ instance (TrieKey k1, TrieKey k2) => TrieKey (k1, k2) where 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 s (k1, k2) a = PMap (singletonM (sizeM s) k1 (singletonM s k2 a)) + singletonM (k1, k2) a = PMap (singletonM k1 (singletonM k2 a)) nullM (PMap m) = nullM m - sizeM s (PMap m) = sizeM (sizeM s) m + sizeM (PMap m) = sizeM m lookupM (k1, k2) (PMap m) = lookupM k1 m >>= lookupM k2 - alterM s f (a, b) (PMap m) = PMap (alterM (sizeM s) g a m) where - g = guardNullM . alterM s f b . fromMaybe emptyM - alterLookupM s f (a, b) (PMap m) = onUnboxed PMap (alterLookupM (sizeM s) g a) m where - g (Just m) = onUnboxed guardNullM (alterLookupM s f b) m - g _ = onUnboxed guardNullM (alterLookupM s f b) emptyM - traverseWithKeyM s f (PMap m) = PMap <$> traverseWithKeyM (sizeM s) (\ a -> traverseWithKeyM s (f . (a,))) m - foldWithKeyM f (PMap m) = foldWithKeyM (\ a -> foldWithKeyM (f . (a,))) m + traverseWithKeyM f (PMap m) = PMap <$> traverseWithKeyM (\ a -> traverseWithKeyM (f . (a,))) m + foldrWithKeyM f (PMap m) = foldrWithKeyM (\ a -> foldrWithKeyM (f . (a,))) m foldlWithKeyM f (PMap m) = foldlWithKeyM (\ a -> flip (foldlWithKeyM (f . (a,)))) m - mapMaybeM s f (PMap m) = PMap (mapMaybeM (sizeM s) g m) where - g a = guardNullM . mapMaybeM s (f . (a,)) - mapEitherM s1 s2 f (PMap m) = both PMap PMap (mapEitherM (sizeM s1) (sizeM s2) g) m where - g a m = both guardNullM guardNullM (mapEitherM s1 s2 (f . (a,))) m - splitLookupM s f (a, b) (PMap m) = sides PMap (splitLookupM (sizeM s) g a) m where - g = sides guardNullM (splitLookupM s f b) + mapWithKeyM f (PMap m) = PMap (mapWithKeyM (\ a -> mapWithKeyM (f . (a,))) m) + mapMaybeM f (PMap m) = PMap (mapMaybeM g m) where + g a = guardNullM . mapMaybeM (f . (a,)) + mapEitherM f (PMap m) = both PMap PMap (mapEitherM g) m where + g a m = both guardNullM guardNullM (mapEitherM (f . (a,))) m isSubmapM (<=) (PMap m1) (PMap m2) = isSubmapM (isSubmapM (<=)) m1 m2 - unionM s f (PMap m1) (PMap m2) = PMap (unionM (sizeM s) (\ a -> guardNullM .: unionM s (f . (a,))) m1 m2) - isectM s f (PMap m1) (PMap m2) = PMap (isectM (sizeM s) (\ a -> guardNullM .: isectM s (f . (a,))) m1 m2) - diffM s f (PMap m1) (PMap m2) = PMap (diffM (sizeM s) (\ a -> guardNullM .: diffM s (f . (a,))) m1 m2) - extractM s f (PMap m) = fmap PMap <$> extractM (sizeM s) g m where - g a = fmap guardNullM <.> extractM s (f . (a,)) - fromListM s f xs = PMap (mapWithKeyM (sizeM s) (\ a -> fromListM s (f . (a,))) - (fromListM (const 1) (const (++)) (breakFst xs))) - fromAscListM s f xs = PMap (fromDistAscListM (sizeM s) - [(a, fromAscListM s (f . (a,)) ys) | (a, ys) <- breakFst xs]) + unionM f (PMap m1) (PMap m2) = PMap (unionM (\ a -> guardNullM .: unionM (f . (a,))) m1 m2) + isectM f (PMap m1) (PMap m2) = PMap (isectM (\ a -> guardNullM .: isectM (f . (a,))) m1 m2) + diffM f (PMap m1) (PMap m2) = PMap (diffM (\ a -> guardNullM .: diffM (f . (a,))) m1 m2) + fromListM f xs = PMap (mapWithKeyM (\ a (Elem xs) -> fromListM (f . (a,)) xs) + (fromListM (\ _ (Elem xs) (Elem ys) -> Elem (xs ++ ys)) (breakFst xs))) + fromAscListM f xs = PMap (fromDistAscListM + [(a, fromAscListM (f . (a,)) ys) | (a, Elem ys) <- breakFst xs]) -breakFst :: Eq k1 => [((k1, k2), a)] -> [(k1, [(k2, a)])] + singleHoleM (k1, k2) = PHole (singleHoleM k1) (singleHoleM k2) + keyM (PHole hole1 hole2) = (keyM hole1, keyM hole2) + assignM v (PHole hole1 hole2) = PMap (assignM (assignM v hole2) hole1) + clearM (PHole hole1 hole2) = PMap (fillHoleM (guardNullM (clearM hole2)) hole1) + beforeM a (PHole hole1 hole2) + = PMap (beforeM (guardNullM (beforeM a hole2)) hole1) + afterM a (PHole hole1 hole2) + = PMap (afterM (guardNullM (afterM a hole2)) hole1) + searchM (k1, k2) (PMap m) = case searchM k1 m of + (# Nothing, hole1 #) -> (# Nothing, PHole hole1 (singleHoleM k2) #) + (# Just m', hole1 #) -> onUnboxed (PHole hole1) (searchM k2) m' + indexM i (PMap m) + | (# i', m', hole1 #) <- indexM i m, + (# i'', v, hole2 #) <- indexM i' m' + = (# i'', v, PHole hole1 hole2 #) + extractHoleM (PMap m) = do + (m', hole1) <- extractHoleM m + (v, hole2) <- extractHoleM m' + return (v, PHole hole1 hole2) + +breakFst :: Eq 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, toList vs):breakFst' a' (Seq.singleton (b', v')) xs - breakFst' a vs [] = [(a, toList vs)] + | otherwise = (a, Elem $ toList vs):breakFst' a' (Seq.singleton (b', v')) xs + breakFst' a vs [] = [(a, Elem $ toList vs)] diff --git a/Data/TrieMap/RadixTrie.hs b/Data/TrieMap/RadixTrie.hs --- a/Data/TrieMap/RadixTrie.hs +++ b/Data/TrieMap/RadixTrie.hs @@ -1,217 +1,267 @@ -{-# LANGUAGE BangPatterns, UnboxedTuples, TupleSections, TypeFamilies, PatternGuards, UnboxedTuples #-} +{-# LANGUAGE BangPatterns, UnboxedTuples, TupleSections, TypeFamilies, PatternGuards, MagicHash #-} module Data.TrieMap.RadixTrie () where import Data.TrieMap.TrieKey import Data.TrieMap.Sized -import Data.TrieMap.Applicative +-- import Data.TrieMap.Applicative import Control.Applicative import Control.Monad import Data.Maybe -import Data.Foldable +import Data.Foldable (foldr, foldl) import Data.Traversable +import GHC.Exts + import Prelude hiding (lookup, foldr, foldl) -data Edge k m a = Edge {-# UNPACK #-} !Int [k] (Maybe a) (m (Edge k m a)) -type Edge' k a = Edge k (TrieMap k) a -type MEdge' k a = Maybe (Edge' k a) +data Assoc k a = Empty | Assoc [k] a +data Edge k a = Edge Int# [k] (Assoc k a) (TrieMap k (Edge k a)) +type MEdge k a = Maybe (Edge k a) -edgeSize :: Edge k m a -> Int -edgeSize (Edge sz _ _ _) = sz +instance Sized (Edge k a) where + getSize# (Edge sz _ _ _) = sz +instance Sized a => Sized (Assoc k a) where + getSize# (Assoc _ a) = getSize# a + getSize# _ = 0# + +data Path k a = Root + | Deep (Path k a) [k] (Assoc k a) (Hole k (Edge k a)) + instance TrieKey k => TrieKey [k] where - newtype TrieMap [k] a = Radix (MEdge' k a) + newtype TrieMap [k] a = Radix (MEdge k a) + data Hole [k] a = Hole [k] [k] (TrieMap k (Edge k a)) (Path k a) + emptyM = Radix Nothing - singletonM s ks a = Radix (Just (Edge (s a) ks (Just a) emptyM)) + singletonM ks a = Radix (Just (Edge (getSize# a) ks (Assoc ks a) emptyM)) nullM (Radix m) = isNothing m - sizeM _ (Radix m) = maybe 0 edgeSize m + sizeM (Radix (Just e)) = getSize# e + sizeM _ = 0# lookupM ks (Radix m) = m >>= lookup ks - alterM s f ks (Radix m) = Radix (alter s f ks m) - alterLookupM s f ks (Radix m) = onUnboxed Radix (alterLookupE s f ks) m - traverseWithKeyM s f (Radix m) = Radix <$> traverse (traverseE s f) m - extractM s f (Radix m) = maybe empty (fmap Radix <.> extractE s f) m - foldWithKeyM f (Radix m) z = foldr (foldE f) z m + traverseWithKeyM f (Radix m) = Radix <$> traverse (traverseE f) m + foldrWithKeyM f (Radix m) z = foldr (foldrE f) z m foldlWithKeyM f (Radix m) z = foldl (foldlE f) z m - mapMaybeM s f (Radix m) = Radix (m >>= mapMaybeE s f) - mapEitherM _ _ _ (Radix Nothing) = (# emptyM, emptyM #) - mapEitherM s1 s2 f (Radix (Just m)) = both Radix Radix (mapEitherE s1 s2 f) m - unionM s f (Radix m1) (Radix m2) = Radix (unionMaybe (unionE s f) m1 m2) - isectM s f (Radix m1) (Radix m2) = Radix (isectMaybe (isectE s f) m1 m2) - diffM s f (Radix m1) (Radix m2) = Radix (diffMaybe (diffE s f) m1 m2) --- lookupIxM s ks (Radix m) = maybe (empty, empty, empty) (lookupIxE s 0 ks) m + mapWithKeyM f (Radix m) = Radix (mapWithKeyE f <$> m) + mapMaybeM f (Radix m) = Radix (m >>= mapMaybeE f) + mapEitherM _ (Radix Nothing) = (# emptyM, emptyM #) + mapEitherM f (Radix (Just m)) = both Radix Radix (mapEitherE f) m + unionM f (Radix m1) (Radix m2) = Radix (unionMaybe (unionE f) m1 m2) + isectM f (Radix m1) (Radix m2) = Radix (isectMaybe (isectE f) m1 m2) + diffM f (Radix m1) (Radix m2) = Radix (diffMaybe (diffE f) m1 m2) isSubmapM (<=) (Radix m1) (Radix m2) = subMaybe (isSubmapE (<=)) m1 m2 - splitLookupM _ _ _ (Radix Nothing) = (# emptyM, Nothing, emptyM #) - splitLookupM s f ks (Radix (Just e)) = sides Radix (splitLookupE s f ks) e --- assocAtM s i (Radix m) = maybe (empty, empty, empty) (assocAtE s 0 i) m - -cat :: [k] -> Edge' k a -> Edge' k a + + singleHoleM ks = Hole ks ks emptyM Root + keyM (Hole ks _ _ _) = ks + beforeM a (Hole ks0 ks ts path) = before (compact (edge ks v ts)) path where + v = case a of + Nothing -> Empty + Just a -> Assoc ks0 a + before t Root = Radix t + before e (Deep path ks v tHole) = + before (compact $ edge ks v $ beforeM e tHole) path + afterM a (Hole ks0 ks ts path) = after (compact (edge ks v ts)) path where + v = case a of + Nothing -> Empty + Just a -> Assoc ks0 a + after t Root = Radix t + after e (Deep path ks v tHole) = + after (compact $ edge ks v $ afterM e tHole) path + + searchM ks (Radix Nothing) = (# Nothing, singleHoleM ks #) + searchM ks (Radix (Just e)) = case searchE ks e Root of + (# v, holer #) -> (# v, holer ks #) + + indexM _ (Radix Nothing) = (# error err, error err, error err #) + where err = "Error: trie map is empty" + indexM i# (Radix (Just e)) = indexE i# e Root + + extractHoleM (Radix Nothing) = mzero + extractHoleM (Radix (Just e)) = extractHoleE Root e + + assignM a (Hole ks0 ks ts path) = Radix $ rebuild (compact (edge ks (Assoc ks0 a) ts)) path + + clearM (Hole _ ks ts path) = Radix $ rebuild (compact (edge ks Empty ts)) path + +rebuild :: (TrieKey k, Sized a) => MEdge k a -> Path k a -> MEdge k a +rebuild e (Deep path ks v tHole) = + rebuild (compact (edge ks v (fillHoleM e tHole))) path +rebuild e _ = e + +cat :: [k] -> Edge k a -> Edge k a ks `cat` Edge sz ls v ts = Edge sz (ks ++ ls) v ts -cons :: k -> Edge' k a -> Edge' k a +cons :: k -> Edge k a -> Edge k a k `cons` Edge sz ks v ts = Edge sz (k:ks) v ts -edge :: TrieKey k => Sized a -> [k] -> Maybe a -> TrieMap k (Edge' k a) -> Edge' k a -edge s ks v ts = Edge (maybe 0 s v + sizeM edgeSize ts) ks v ts - -singleMaybe :: TrieKey k => Sized a -> [k] -> Maybe a -> MEdge' k a -singleMaybe s ks v = do v <- v - return (edge s ks (Just v) emptyM) +edge :: (TrieKey k, Sized a) => [k] -> Assoc k a -> TrieMap k (Edge k a) -> Edge k a +edge ks v ts = Edge (getSize# v +# getSize# ts) ks v ts -compact :: TrieKey k => Edge' k a -> MEdge' k a -compact e@(Edge _ ks Nothing ts) = case assocsM ts of +compact :: TrieKey k => Edge k a -> MEdge k a +compact e@(Edge _ ks Empty ts) = case assocsM ts of [] -> Nothing [(l, e')] -> compact (ks `cat` (l `cons` e')) _ -> Just e compact e = Just e -lookup :: (Eq k, TrieKey k) => [k] -> Edge' k a -> Maybe a +lookup :: (Eq k, TrieKey k) => [k] -> Edge k a -> Maybe a lookup ks (Edge _ ls v ts) = match ks ls where match (k:ks) (l:ls) | k == l = match ks ls match (k:ks) [] = lookupM k ts >>= lookup ks - match [] [] = v + match [] [] = case v of + Assoc _ a -> Just a + _ -> Nothing match _ _ = Nothing -alter :: TrieKey k => Sized a -> (Maybe a -> Maybe a) -> [k] -> MEdge' k a -> MEdge' k a -alter s f ks0 Nothing = singleMaybe s ks0 (f Nothing) -alter s f ks0 (Just e@(Edge sz ls0 v ts)) = match 0 ks0 ls0 where - match !i (k:ks) (l:ls) = case compare k l of - LT | Just v' <- f Nothing - -> Just $ let sv = s v' in Edge (sv + sz) (take i ls0) Nothing (fromDistAscListM edgeSize - [(k, Edge sv ks (Just v') emptyM), (l, Edge sz ls v ts)]) - EQ -> match (i+1) ks ls - GT | Just v' <- f Nothing - -> Just $ let sv = s v' in Edge (sv + sz) (take i ls0) Nothing (fromDistAscListM edgeSize - [(l, Edge sz ls v ts), (k, Edge sv ks (Just v') emptyM)]) - _ -> Just e - match _ (k:ks) [] = compact $ edge s ls0 v (alterM edgeSize g k ts) where - g = alter s f ks - match _ [] (l:ls) - | Just v' <- f Nothing - = Just (Edge (s v' + sz) ks0 (Just v') (singletonM edgeSize l (Edge sz ls v ts))) - match _ [] [] - = compact (edge s ls0 (f v) ts) - match _ _ _ = Just e +traverseA :: Applicative f => ([k] -> a -> f b) -> Assoc k a -> f (Assoc k b) +traverseA f (Assoc ks a) = Assoc ks <$> f ks a +traverseA _ _ = pure Empty -alterLookupE :: TrieKey k => Sized a -> (Maybe a -> (# z, Maybe a #)) -> [k] -> MEdge' k a -> (# z, MEdge' k a #) -alterLookupE s f ks Nothing = onUnboxed (singleMaybe s ks) f Nothing -alterLookupE s f ks0 (Just e@(Edge sz ls0 v0 ts0)) = match 0 ks0 ls0 where - match !i (k:ks) (l:ls) = case compare k l of - LT -> onUnboxed (Just . maybe e (\ v' -> let sv = s v' in Edge (sz + sv) (take i ls0) Nothing $ - fromDistAscListM edgeSize [(k, Edge sv ks (Just v') emptyM), (l, Edge sz ls v0 ts0)])) - f Nothing - GT -> onUnboxed (Just . maybe e (\ v' -> let sv = s v' in Edge (sz + sv) (take i ls0) Nothing $ - fromDistAscListM edgeSize [(l, Edge sz ls v0 ts0), (k, Edge sv ks (Just v') emptyM)])) - f Nothing - EQ -> match (i+1) ks ls - match _ (k:ks) [] = onUnboxed (compact . edge s ls0 v0) (alterLookupM edgeSize g k) ts0 where - g = alterLookupE s f ks - match _ [] (l:ls) = onUnboxed (Just . maybe e (\ v' -> let sv = s v' in - Edge (sv + sz) ks0 (Just v') (singletonM edgeSize l (Edge sz ls v0 ts0)))) - f Nothing - match _ [] [] = onUnboxed (\ v' -> compact $ edge s ls0 v' ts0) f v0 +traverseE :: (Applicative f, TrieKey k, Sized b) => ([k] -> a -> f b) -> Edge k a -> f (Edge k b) +traverseE f (Edge _ ks v ts) + = edge ks <$> traverseA f v <*> traverseM (traverseE f) ts -traverseE :: (Applicative f, TrieKey k) => Sized b -> ([k] -> a -> f b) -> Edge' k a -> f (Edge' k b) -traverseE s f (Edge _ ks v ts) - = edge s ks <$> traverse (f ks) v <*> traverseWithKeyM edgeSize g ts - where g l = traverseE s (\ ls -> f (ks ++ l:ls)) +foldrA :: ([k] -> a -> b -> b) -> Assoc k a -> b -> b +foldrA f (Assoc ks a) = f ks a +foldrA _ _ = id -extractE :: (Alternative f, TrieKey k) => Sized a -> ([k] -> a -> f (x, Maybe a)) -> Edge' k a -> f (x, MEdge' k a) -extractE s f (Edge _ ks v ts) = case v of - Nothing -> rest - Just v -> fmap (\ v' -> compact (edge s ks v' ts)) <$> f ks v <|> rest - where rest = fmap (compact . edge s ks v) <$> extractM edgeSize g ts - g l = extractE s (\ ls -> f (ks ++ l:ls)) +foldlA :: ([k] -> b -> a -> b) -> b -> Assoc k a -> b +foldlA f z (Assoc ks a) = f ks z a +foldlA _ z _ = z -foldE :: TrieKey k => ([k] -> a -> b -> b) -> Edge' k a -> b -> b -foldE f (Edge _ ks v ts) z = foldr (f ks) (foldWithKeyM g ts z) v where - g l = foldE (\ ls -> f (ks ++ l:ls)) +foldrE :: TrieKey k => ([k] -> a -> b -> b) -> Edge k a -> b -> b +foldrE f (Edge _ _ v ts) z = foldrA f v (foldr (foldrE f) z ts) -foldlE :: TrieKey k => ([k] -> b -> a -> b) -> b -> Edge' k a -> b -foldlE f z (Edge _ ks v ts) = foldlWithKeyM g ts (foldl (f ks) z v) where - g l = foldlE (\ ls -> f (ks ++ l:ls)) +foldlE :: TrieKey k => ([k] -> b -> a -> b) -> b -> Edge k a -> b +foldlE f z (Edge _ _ v ts) = foldl (foldlE f) (foldlA f z v) ts -mapMaybeE :: TrieKey k => Sized b -> ([k] -> a -> Maybe b) -> Edge' k a -> MEdge' k b -mapMaybeE s f (Edge _ ks v ts) = compact (edge s ks (v >>= f ks) - (mapMaybeM edgeSize (\ l -> mapMaybeE s (\ ls -> f (ks ++ l:ls))) ts)) +mapWithKeyA :: ([k] -> a -> b) -> Assoc k a -> Assoc k b +mapWithKeyA f (Assoc ks a) = Assoc ks (f ks a) +mapWithKeyA _ _ = Empty -mapEitherE :: TrieKey k => Sized b -> Sized c -> ([k] -> a -> (# Maybe b, Maybe c #)) -> Edge' k a -> - (# MEdge' k b, MEdge' k c #) -mapEitherE s1 s2 f (Edge _ ks v ts) = case mapEitherM edgeSize edgeSize (\ l -> mapEitherE s1 s2 (\ ls -> f (ks ++ l:ls))) ts of - (# tsL, tsR #) -> case v of - Nothing -> (# compact (edge s1 ks Nothing tsL), compact (edge s2 ks Nothing tsR) #) - Just v -> case f ks v of - (# vL, vR #) -> (# compact (edge s1 ks vL tsL), compact (edge s2 ks vR tsR) #) +mapWithKeyE :: (TrieKey k, Sized b) => ([k] -> a -> b) -> Edge k a -> Edge k b +mapWithKeyE f (Edge _ ks v ts) = edge ks (mapWithKeyA f v) (fmapM (mapWithKeyE f) ts) -unionE :: TrieKey k => Sized a -> ([k] -> a -> a -> Maybe a) -> Edge' k a -> Edge' k a -> MEdge' k a -unionE s f (Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match 0 ks0 ls0 where +mapMaybeA :: ([k] -> a -> Maybe b) -> Assoc k a -> Assoc k b +mapMaybeA f (Assoc ks a) = maybe Empty (Assoc ks) (f ks a) +mapMaybeA _ _ = Empty + +mapMaybeE :: (TrieKey k, Sized b) => ([k] -> a -> Maybe b) -> Edge k a -> MEdge k b +mapMaybeE f (Edge _ ks v ts) = compact (edge ks (mapMaybeA f v) + (mapMaybeM (const $ mapMaybeE f) ts)) + +mapEitherA :: ([k] -> a -> (# Maybe b, Maybe c #)) -> Assoc k a -> (# Assoc k b, Assoc k c #) +mapEitherA f (Assoc ks a) = case f ks a of + (# vL, vR #) -> (# maybe Empty (Assoc ks) vL, maybe Empty (Assoc ks) vR #) +mapEitherA _ _ = (# Empty, Empty #) + +mapEitherE :: (TrieKey k, Sized b, Sized c) => ([k] -> a -> (# Maybe b, Maybe c #)) -> Edge k a -> + (# MEdge k b, MEdge k c #) +mapEitherE f (Edge _ ks v ts) = case mapEitherA f v of + (# vL, vR #) -> case mapEitherM (\ _ -> mapEitherE f) ts of + (# tsL, tsR #) -> (# compact (edge ks vL tsL), compact (edge ks vR tsR) #) + +unionE :: (TrieKey k, Sized a) => ([k] -> a -> a -> Maybe a) -> Edge k a -> Edge k a -> MEdge k a +unionE f (Edge szK# ks0 vK tsK) (Edge szL# ls0 vL tsL) = match 0 ks0 ls0 where match !i (k:ks) (l:ls) = case compare k l of EQ -> match (i+1) ks ls - LT -> Just $ Edge (szK + szL) (take i ks0) Nothing (fromDistAscListM edgeSize - [(k, Edge szK ks vK tsK), (l, Edge szL ls vL tsL)]) - GT -> Just $ Edge (szK + szL) (take i ks0) Nothing (fromDistAscListM edgeSize - [(l, Edge szL ls vL tsL), (k, Edge szK ks vK tsK)]) - match _ [] (l:ls) = compact (edge s ks0 vK (alterM edgeSize g l tsK)) where - g (Just eK') = unionE s (\ ls' -> f (ks0 ++ l:ls')) eK' (Edge szL ls vL tsL) - g Nothing = Just (Edge szL ls vL tsL) - match _ (k:ks) [] = compact (edge s ls0 vL (alterM edgeSize g k tsL)) where - g Nothing = Just (Edge szK ks vK tsK) - g (Just eL') = unionE s (\ ks' -> f (ls0 ++ k:ks')) (Edge szK ks vK tsK) eL' - match _ [] [] = compact (edge s ls0 (unionMaybe (f ls0) vK vL) (unionM edgeSize g tsK tsL)) where - g x = unionE s (\ xs -> f (ls0 ++ x:xs)) + LT -> Just $ Edge (szK# +# szL#) (take i ks0) Empty (fromDistAscListM + [(k, Edge szK# ks vK tsK), (l, Edge szL# ls vL tsL)]) + GT -> Just $ Edge (szK# +# szL#) (take i ks0) Empty (fromDistAscListM + [(l, Edge szL# ls vL tsL), (k, Edge szK# ks vK tsK)]) + match _ [] (l:ls) = compact (edge ks0 vK (alterM g l tsK)) where + g (Just eK') = unionE f eK' (Edge szL# ls vL tsL) + g Nothing = Just (Edge szL# ls vL tsL) + match _ (k:ks) [] = compact (edge ls0 vL (alterM g k tsL)) where + g Nothing = Just (Edge szK# ks vK tsK) + g (Just eL') = unionE f (Edge szK# ks vK tsK) eL' + match _ [] [] = compact (edge ls0 (unionA f vK vL) (unionM (const $ unionE f) tsK tsL)) -isectE :: TrieKey k => Sized c -> ([k] -> a -> b -> Maybe c) -> Edge' k a -> Edge' k b -> MEdge' k c -isectE s f (Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match ks0 ls0 where +unionA :: ([k] -> a -> a -> Maybe a) -> Assoc k a -> Assoc k a -> Assoc k a +unionA f (Assoc ks v1) (Assoc _ v2) = maybe Empty (Assoc ks) (f ks v1 v2) +unionA _ Empty v = v +unionA _ v Empty = v + +isectE :: (TrieKey k, Sized c) => ([k] -> a -> b -> Maybe c) -> Edge k a -> Edge k b -> MEdge k c +isectE f (Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match ks0 ls0 where match (k:ks) (l:ls) | k == l = match ks ls match (k:ks) [] = do eL' <- lookupM k tsL - cat ls0 <$> cons k <$> isectE s (\ ks' -> f (ls0 ++ k:ks')) (Edge szK ks vK tsK) eL' + cat ls0 <$> cons k <$> isectE f (Edge szK ks vK tsK) eL' match [] (l:ls) = do eK' <- lookupM l tsK - cat ks0 <$> cons l <$> isectE s (\ ls' -> f (ks0 ++ l:ls')) eK' (Edge szL ls vL tsL) - match [] [] = compact (edge s ks0 (isectMaybe (f ks0) vK vL) (isectM edgeSize g tsK tsL)) where - g x = isectE s (\ xs -> f (ks0 ++ x:xs)) + cat ks0 <$> cons l <$> isectE f eK' (Edge szL ls vL tsL) + match [] [] = compact (edge ks0 (isectA f vK vL) (isectM (const $ isectE f) tsK tsL)) match _ _ = Nothing -diffE :: TrieKey k => Sized a -> ([k] -> a -> b -> Maybe a) -> Edge' k a -> Edge' k b -> MEdge' k a -diffE s f eK@(Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match ks0 ls0 where +isectA :: ([k] -> a -> b -> Maybe c) -> Assoc k a -> Assoc k b -> Assoc k c +isectA f (Assoc ks a) (Assoc _ b) = maybe Empty (Assoc ks) (f ks a b) +isectA _ _ _ = Empty + +diffE :: (TrieKey k, Sized a) => ([k] -> a -> b -> Maybe a) -> Edge k a -> Edge k b -> MEdge k a +diffE f eK@(Edge szK ks0 vK tsK) (Edge szL ls0 vL tsL) = match ks0 ls0 where match (k:ks) (l:ls) | k == l = match ks ls match (k:ks) [] | Just eL' <- lookupM k tsL - = cat ls0 . cons k <$> diffE s (\ ks' -> f (ls0 ++ k:ks')) (Edge szK ks vK tsK) eL' + = cat ls0 . cons k <$> diffE f (Edge szK ks vK tsK) eL' match [] (l:ls) - = compact (edge s ks0 vK (alterM edgeSize (>>= g) l tsK)) - where g eK' = diffE s (\ ls' -> f (ks0 ++ l:ls')) eK' (Edge szL ls vL tsL) - match [] [] = compact (edge s ks0 (diffMaybe (f ks0) vK vL) (diffM edgeSize g tsK tsL)) where - g x = diffE s (\ xs -> f (ks0 ++ x:xs)) + = compact (edge ks0 vK (alterM (>>= g) l tsK)) + where g eK' = diffE f eK' (Edge szL ls vL tsL) + match [] [] = compact (edge ks0 (diffA f vK vL) (diffM (const $ diffE f) tsK tsL)) match _ _ = Just eK +diffA :: ([k] -> a -> b -> Maybe a) -> Assoc k a -> Assoc k b -> Assoc k a +diffA f (Assoc ks a) (Assoc _ b) = maybe Empty (Assoc ks) (f ks a b) +diffA _ a@Assoc{} Empty = a +diffA _ Empty _ = Empty -isSubmapE :: TrieKey k => LEq a b -> LEq (Edge' k a) (Edge' k b) +isSubmapE :: TrieKey k => LEq a b -> LEq (Edge k a) (Edge k b) isSubmapE (<=) (Edge szK ks vK tsK) (Edge _ ls vL tsL) = match ks ls where match (k:ks) (l:ls) | k == l = match ks ls match (k:ks) [] | Just eL' <- lookupM k tsL = isSubmapE (<=) (Edge szK ks vK tsK) eL' - match [] [] = subMaybe (<=) vK vL && isSubmapM (isSubmapE (<=)) tsK tsL + match [] [] = subA (<=) vK vL && isSubmapM (isSubmapE (<=)) tsK tsL match _ _ = False -splitLookupE :: TrieKey k => Sized a -> (a -> (# Maybe a, Maybe x, Maybe a #)) -> [k] -> Edge' k a -> - (# MEdge' k a, Maybe x, MEdge' k a #) -splitLookupE s f ks e@(Edge _ ls v ts) = match ks ls where - match (k:ks) (l:ls) = case compare k l of - LT -> (# Nothing, Nothing, Just e #) - GT -> (# Just e, Nothing, Nothing #) - EQ -> match ks ls - match (k:ks) [] = case splitLookupM edgeSize g k ts of - (# tsL, x, tsR #) -> (# compact (edge s ls v tsL), x, compact (edge s ls Nothing tsR) #) - where g = splitLookupE s f ks - match [] (_:_) = (# Nothing, Nothing, Just e #) - match [] [] = case v of - Nothing -> (# Nothing, Nothing, compact (edge s ls Nothing ts) #) - Just v -> case f v of - (# vL, x, vR #) -> (# singleMaybe s ls vL, x, compact (edge s ls vR ts) #) +subA :: LEq a b -> LEq (Assoc k a) (Assoc k b) +subA (<=) (Assoc _ a) (Assoc _ b) = a <= b +subA _ Assoc{} Empty = False +subA _ Empty _ = True + +searchE :: TrieKey k => [k] -> Edge k a -> Path k a -> (# Maybe a, [k] -> Hole [k] a #) +searchE ks0 (Edge sz ls0 v ts) path = match 0 ks0 ls0 where + match !_ [] [] = (# assocToMaybe v, \ k0 -> Hole k0 ls0 ts path #) + match _ (k:ks) [] = case searchM k ts of + (# Just e', tHole #) -> searchE ks e' (Deep path ls0 v tHole) + (# Nothing, tHole #) -> (# Nothing, \ k0 -> Hole k0 ks emptyM (Deep path ls0 v tHole) #) + match i [] (l:ls) = (# Nothing, \ k0 -> Hole k0 (take i ls0) (singletonM l (Edge sz ls v ts)) path #) + match i (k:ks) (l:ls) + | k == l = match (i+1) ks ls + | (# _, kHole #) <- searchM k (singletonM l (Edge sz ls v ts)) + = (# Nothing, \ k0 -> Hole k0 ks emptyM (Deep path (take i ls0) Empty kHole) #) + +assocToMaybe :: Assoc k a -> Maybe a +assocToMaybe (Assoc _ a) = Just a +assocToMaybe _ = Nothing + +indexE :: (TrieKey k, Sized a) => Int# -> Edge k a -> Path k a -> (# Int#, a, Hole [k] a #) +indexE i# (Edge _ ks Empty ts) path + | (# i'#, e, tHole #) <- indexM i# ts + = indexE i'# e (Deep path ks Empty tHole) +indexE i# (Edge _ ks v@(Assoc ks0 a) ts) path + | i# <# sa# = (# i#, a, Hole ks0 ks ts path #) + | (# i'#, e, tHole #) <- indexM (i# -# sa#) ts + = indexE i'# e (Deep path ks v tHole) + where !sa# = getSize# a + +extractHoleE :: (TrieKey k, MonadPlus m) => Path k a -> Edge k a -> m (a, Hole [k] a) +extractHoleE path (Edge _ ks v ts) = case v of + Empty -> tsHoles + Assoc ks0 a -> return (a, Hole ks0 ks ts path) `mplus` tsHoles + where tsHoles = do (e, tHole) <- extractHoleM ts + extractHoleE (Deep path ks v tHole) e diff --git a/Data/TrieMap/Representation.hs b/Data/TrieMap/Representation.hs --- a/Data/TrieMap/Representation.hs +++ b/Data/TrieMap/Representation.hs @@ -15,8 +15,8 @@ instance (TKey k, Repr a) => Repr (TMap k a) where type Rep (TMap k a) = [(Rep k, Rep a)] - toRep (TMap m) = foldWithKeyM (\ k (Elem a) xs -> (k, toRep a):xs) m [] - fromRep xs = TMap (fromDistAscListM (const 1) [(k, Elem (fromRep a)) | (k, a) <- xs]) + toRep (TMap m) = foldrWithKeyM (\ k (Elem a) xs -> (k, toRep a):xs) m [] + fromRep xs = TMap (fromDistAscListM [(k, Elem (fromRep a)) | (k, a) <- xs]) genOrdRepr ''Float genOrdRepr ''Double diff --git a/Data/TrieMap/Representation/TH.hs b/Data/TrieMap/Representation/TH.hs --- a/Data/TrieMap/Representation/TH.hs +++ b/Data/TrieMap/Representation/TH.hs @@ -39,10 +39,7 @@ [VarE x]]) --- | Given the name of a type constructor, automatically generates an efficient 'Repr' instance. --- /Warning/: Generalized tries do not work for "infinitely complicated types," for example, a --- type-system construction of the natural numbers. In these cases, a context reduction stack --- overflow will occur at compile time when you use the 'TKey' instance for that type. +-- | Given the name of a type constructor, automatically generates an efficient 'Repr' instance. genRepr :: Name -> Q [Dec] genRepr tycon = do TyConI dec <- reify tycon diff --git a/Data/TrieMap/ReverseMap.hs b/Data/TrieMap/ReverseMap.hs --- a/Data/TrieMap/ReverseMap.hs +++ b/Data/TrieMap/ReverseMap.hs @@ -1,8 +1,9 @@ -{-# LANGUAGE UnboxedTuples, TypeFamilies #-} +{-# LANGUAGE UnboxedTuples, TypeFamilies, BangPatterns, MagicHash #-} module Data.TrieMap.ReverseMap (reverse, unreverse) where import Data.TrieMap.TrieKey +import Data.TrieMap.Sized import Data.TrieMap.Modifiers import Data.TrieMap.Applicative @@ -11,31 +12,44 @@ import Prelude hiding (reverse) import qualified Data.List as L +import GHC.Exts + instance TrieKey k => TrieKey (Rev k) where newtype TrieMap (Rev k) a = RMap (TrieMap k a) + newtype Hole (Rev k) a = RHole (Hole k a) emptyM = RMap emptyM - singletonM s (Rev k) a = RMap (singletonM s k a) + singletonM (Rev k) a = RMap (singletonM k a) nullM (RMap m) = nullM m - sizeM s (RMap m) = sizeM s m + sizeM (RMap m) = sizeM m lookupM (Rev k) (RMap m) = lookupM k m - traverseWithKeyM s f (RMap m) = RMap <$> runDual (traverseWithKeyM s (\ k a -> Dual (f (Rev k) a)) m) - alterM s f (Rev k) (RMap m) = RMap (alterM s f k m) - alterLookupM s f (Rev k) (RMap m) = onUnboxed RMap (alterLookupM s f k) m - splitLookupM s f (Rev k) (RMap m) = sides RMap (splitLookupM s f' k) m - where f' x = case f x of - (# xL, ans, xR #) -> (# xR, ans, xL #) - mapMaybeM s f (RMap m) = RMap (mapMaybeM s (f . Rev) m) - mapEitherM s1 s2 f (RMap m) = both RMap RMap (mapEitherM s1 s2 (f . Rev)) m - foldWithKeyM f (RMap m) = foldlWithKeyM (flip . f . Rev) m - foldlWithKeyM f (RMap m) = foldWithKeyM (flip . f . Rev) m - unionM s f (RMap m1) (RMap m2) = RMap (unionM s (f . Rev) m1 m2) - isectM s f (RMap m1) (RMap m2) = RMap (isectM s (f . Rev) m1 m2) - diffM s f (RMap m1) (RMap m2) = RMap (diffM s (f . Rev) m1 m2) - extractM s f (RMap m) = fmap RMap <$> runDual (extractM s (\ k a -> Dual (f (Rev k) a)) m) + mapWithKeyM f (RMap m) = RMap (mapWithKeyM (f . Rev) m) + traverseWithKeyM f (RMap m) = RMap <$> runDual (traverseWithKeyM g m) + where g k a = Dual (f (Rev k) a) + mapMaybeM f (RMap m) = RMap (mapMaybeM (f . Rev) m) + mapEitherM f (RMap m) = both RMap RMap (mapEitherM (f . Rev)) m + foldrWithKeyM f (RMap m) = foldlWithKeyM (flip . f . Rev) m + foldlWithKeyM f (RMap m) = foldrWithKeyM (flip . f . Rev) m + unionM f (RMap m1) (RMap m2) = RMap (unionM (f . Rev) m1 m2) + isectM f (RMap m1) (RMap m2) = RMap (isectM (f . Rev) m1 m2) + diffM f (RMap m1) (RMap m2) = RMap (diffM (f . Rev) m1 m2) isSubmapM (<=) (RMap m1) (RMap m2) = isSubmapM (<=) m1 m2 - fromListM s f xs = RMap (fromListM s (f . Rev) [(k, a) | (Rev k, a) <- xs]) - fromAscListM s f xs = RMap (fromAscListM s (\ k -> flip (f (Rev k))) [(k, a) | (Rev k, a) <- L.reverse xs]) - fromDistAscListM s xs = RMap (fromDistAscListM s [(k, a) | (Rev k, a) <- L.reverse xs]) + fromListM f xs = RMap (fromListM (f . Rev) [(k, a) | (Rev k, a) <- xs]) + fromAscListM f xs = RMap (fromAscListM (\ k a1 a2 -> f (Rev k) a2 a1) [(k, a) | (Rev k, a) <- L.reverse xs]) + fromDistAscListM xs = RMap (fromDistAscListM [(k, a) | (Rev k, a) <- L.reverse xs]) + + singleHoleM (Rev k) = RHole (singleHoleM k) + keyM (RHole hole) = Rev (keyM hole) + beforeM a (RHole hole) = RMap (afterM a hole) + afterM a (RHole hole) = RMap (beforeM a hole) + searchM (Rev k) (RMap m) = onUnboxed RHole (searchM k) m + indexM i# (RMap m) = case indexM (sm# -# 1# -# i#) m of + (# i'#, v, hole #) -> (# getSize# v -# 1# -# i'#, v, RHole hole #) + where !sm# = sizeM m + extractHoleM (RMap m) = do + (v, hole) <- runDualPlus (extractHoleM m) + return (v, RHole hole) + assignM x (RHole hole) = RMap (assignM x hole) + clearM (RHole hole) = RMap (clearM hole) reverse :: TrieMap k a -> TrieMap (Rev k) a reverse = RMap diff --git a/Data/TrieMap/Sized.hs b/Data/TrieMap/Sized.hs --- a/Data/TrieMap/Sized.hs +++ b/Data/TrieMap/Sized.hs @@ -1,18 +1,19 @@ -{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE MagicHash #-} module Data.TrieMap.Sized where --- class Sized f where --- getSize :: f a -> Int --- --- newtype Elem a = Elem {getElem :: a} --- --- instance Sized Elem where --- getSize _ = 1 +import GHC.Exts -type Sized a = a -> Int +class Sized a where + getSize# :: a -> Int# newtype Elem a = Elem {getElem :: a} -elemSize :: Sized (Elem a) -elemSize _ = 1 +instance Sized (Elem a) where + getSize# _ = 1# + +getSize :: Sized a => a -> Int +getSize a = I# (getSize# a) + +unbox :: Int -> Int# +unbox (I# i#) = i# diff --git a/Data/TrieMap/TrieKey.hs b/Data/TrieMap/TrieKey.hs --- a/Data/TrieMap/TrieKey.hs +++ b/Data/TrieMap/TrieKey.hs @@ -1,105 +1,119 @@ -{-# LANGUAGE TupleSections, TypeFamilies, UnboxedTuples #-} +{-# LANGUAGE TupleSections, TypeFamilies, UnboxedTuples, MagicHash #-} module Data.TrieMap.TrieKey where -import Data.TrieMap.Applicative import Data.TrieMap.Sized import Control.Applicative -import Control.Arrow +import Control.Monad import Data.Monoid +import Data.Foldable +import Prelude hiding (foldr, foldl) + + +import GHC.Exts + type EitherMap k a b c = k -> a -> (# Maybe b, Maybe c #) type SplitMap a x = a -> (# Maybe a, Maybe x, Maybe a #) type UnionFunc k a = k -> a -> a -> Maybe a type IsectFunc k a b c = k -> a -> b -> Maybe c type DiffFunc k a b = k -> a -> b -> Maybe a -type ExtractFunc f m k a x = (k -> a -> f (x, Maybe a)) -> m -> f (x, m) type LEq a b = a -> b -> Bool -data Assoc k a = Asc {-# UNPACK #-} !Int k a -type IndexPos k a = (# Last (Assoc k a), Maybe (Assoc k a), First (Assoc k a) #) - -onIndexA :: (Int -> Int) -> Assoc k a -> Assoc k a -onIndexA f (Asc i k a) = Asc (f i) k a - -onKeyA :: (k -> k') -> Assoc k a -> Assoc k' a -onKeyA = onValueA . first - -onValA :: (a -> a') -> Assoc k a -> Assoc k a' -onValA = onValueA . second - -{-# INLINE onValueA #-} -onValueA :: ((k, a) -> (k', a')) -> Assoc k a -> Assoc k' a' -onValueA f (Asc i k a) = uncurry (Asc i) (f (k, a)) - onUnboxed :: (c -> d) -> (a -> (# b, c #)) -> a -> (# b, d #) onUnboxed g f a = case f a of - (# b, c #) -> (# b, g c #) + (# b, c #) -> (# b, g c #) +instance TrieKey k => Foldable (TrieMap k) where + foldr f z m = foldrWithKeyM (const f) m z + foldl f z m = foldlWithKeyM (const f) m z + class Ord k => TrieKey k where data TrieMap k :: * -> * emptyM :: TrieMap k a - singletonM :: Sized a -> k -> a -> TrieMap k a + singletonM :: Sized a => k -> a -> TrieMap k a nullM :: TrieMap k a -> Bool - sizeM :: Sized a -> TrieMap k a -> Int + sizeM :: Sized a => TrieMap k a -> Int# lookupM :: k -> TrieMap k a -> Maybe a - alterM :: Sized a -> (Maybe (a) -> Maybe (a)) -> k -> TrieMap k a -> TrieMap k a - alterLookupM :: Sized a -> (Maybe a -> (# x, Maybe a #)) -> k -> TrieMap k a -> (# x, TrieMap k a #) - {-# SPECIALIZE traverseWithKeyM :: (k -> a -> Id (b)) -> TrieMap k a -> Id (TrieMap k b) #-} - traverseWithKeyM :: (TrieMap k ~ m, Applicative f) => Sized b -> - (k -> a -> f (b)) -> TrieMap k a -> f (TrieMap k b) - foldWithKeyM :: (k -> a -> b -> b) -> TrieMap k a -> b -> b + mapWithKeyM :: Sized b => (k -> a -> b) -> TrieMap k a -> TrieMap k b + traverseWithKeyM :: (Applicative f, Sized b) => + (k -> a -> f b) -> TrieMap k a -> f (TrieMap k b) + foldrWithKeyM :: (k -> a -> b -> b) -> TrieMap k a -> b -> b foldlWithKeyM :: (k -> b -> a -> b) -> TrieMap k a -> b -> b - mapMaybeM :: Sized b -> (k -> a -> Maybe b) -> TrieMap k a -> TrieMap k b - mapEitherM :: Sized b -> Sized c -> EitherMap k (a) (b) (c) -> TrieMap k a -> (# TrieMap k b, TrieMap k c #) - splitLookupM :: Sized a -> SplitMap a x -> k -> TrieMap k a -> (# TrieMap k a, Maybe x, TrieMap k a #) - unionM :: Sized a -> UnionFunc k (a) -> TrieMap k a -> TrieMap k a -> TrieMap k a - isectM :: Sized c -> IsectFunc k (a) (b) (c) -> TrieMap k a -> TrieMap k b -> TrieMap k c - diffM :: Sized a -> DiffFunc k (a) (b) -> TrieMap k a -> TrieMap k b -> TrieMap k a - extractM :: (Alternative f) => Sized a -> ExtractFunc f (TrieMap k a) k a x - isSubmapM :: LEq (a) (b) -> LEq (TrieMap k a) (TrieMap k b) - fromListM, fromAscListM :: Sized a -> (k -> a -> a -> a) -> [(k, a)] -> TrieMap k a - fromDistAscListM :: Sized a -> [(k, a)] -> TrieMap k a + mapMaybeM :: Sized b => (k -> a -> Maybe b) -> TrieMap k a -> TrieMap k b + mapEitherM :: (Sized b, Sized c) => EitherMap k a b c -> TrieMap k a -> (# TrieMap k b, TrieMap k c #) + unionM :: Sized a => UnionFunc k a -> TrieMap k a -> TrieMap k a -> TrieMap k a + isectM :: (Sized a, Sized b, Sized c) => IsectFunc k a b c -> TrieMap k a -> TrieMap k b -> TrieMap k c + diffM :: Sized a => DiffFunc k a b -> TrieMap k a -> TrieMap k b -> TrieMap k a + isSubmapM :: (Sized a, Sized b) => LEq a b -> LEq (TrieMap k a) (TrieMap k b) + fromListM, fromAscListM :: Sized a => (k -> a -> a -> a) -> [(k, a)] -> TrieMap k a + fromDistAscListM :: Sized a => [(k, a)] -> TrieMap k a - sizeM s m = foldWithKeyM (\ _ a n -> s a + n) m 0 - fromListM s f = foldr (uncurry (insertWithKeyM s f)) emptyM + data Hole k :: * -> * + singleHoleM :: k -> Hole k a + keyM :: Hole k a -> k + beforeM :: Sized a => Maybe a -> Hole k a -> TrieMap k a + afterM :: Sized a => Maybe a -> Hole k a -> TrieMap k a + searchM :: k -> TrieMap k a -> (# Maybe a, Hole k a #) + indexM :: Sized a => Int# -> TrieMap k a -> (# Int#, a, Hole k a #) + {-# SPECIALIZE extractHoleM :: Sized a => TrieMap k a -> First (a, Hole k a) #-} + {-# SPECIALIZE extractHoleM :: Sized a => TrieMap k a -> Last (a, Hole k a) #-} + extractHoleM :: MonadPlus m => Sized a => TrieMap k a -> m (a, Hole k a) + assignM :: Sized a => a -> Hole k a -> TrieMap k a + clearM :: Sized a => Hole k a -> TrieMap k a + + singletonM k a = assignM a (singleHoleM k) + lookupM k m = case searchM k m of + (# a, _ #) -> a + foldrWithKeyM f = appEndo . getConst . traverseWithKeyM (endofy f) where + endofy :: (k -> a -> b -> b) -> k -> a -> Const (Endo b) (Elem ()) + endofy f k a = Const (Endo (f k a)) + foldlWithKeyM f m = foldrWithKeyM (\ k a g z -> g (f k z a)) m id + fromListM f = foldr (uncurry (insertWithKeyM f)) emptyM fromAscListM = fromListM - fromDistAscListM s = fromAscListM s (const const) + fromDistAscListM = fromAscListM (const const) +instance (TrieKey k, Sized a) => Sized (TrieMap k a) where + getSize# = sizeM + +{-# INLINE alterM #-} +alterM :: (TrieKey k, Sized a) => (Maybe a -> Maybe a) -> k -> TrieMap k a -> TrieMap k a +alterM f k m = case searchM k m of + (# Nothing, hole #) -> maybe m (\ a -> assignM a hole) (f Nothing) + (# a, hole #) -> fillHoleM (f a) hole + +traverseM :: (Applicative f, TrieKey k, Sized b) => (a -> f b) -> TrieMap k a -> f (TrieMap k b) +traverseM f = traverseWithKeyM (const f) + guardNullM :: TrieKey k => TrieMap k a -> Maybe (TrieMap k a) guardNullM m | nullM m = Nothing | otherwise = Just m +fillHoleM :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> TrieMap k a +fillHoleM Nothing hole = clearM hole +fillHoleM (Just a) hole = assignM a hole + sides :: (b -> d) -> (a -> (# b, c, b #)) -> a -> (# d, c, d #) sides g f a = case f a of - (# x, y, z #) -> (# g x, y, g z #) + (# x, y, z #) -> (# g x, y, g z #) both :: (b -> b') -> (c -> c') -> (a -> (# b, c #)) -> a -> (# b', c' #) both g1 g2 f a = case f a of - (# x, y #) -> (# g1 x, g2 y #) - -{-# INLINE [1] mapWithKeyM #-} -mapWithKeyM :: TrieKey k => Sized b -> (k -> a -> b) -> TrieMap k a -> TrieMap k b -mapWithKeyM s f = unId . traverseWithKeyM s (Id .: f) + (# x, y #) -> (# g1 x, g2 y #) -mapM :: TrieKey k => Sized b -> (a -> b) -> TrieMap k a -> TrieMap k b -mapM s = mapWithKeyM s . const +fmapM :: (TrieKey k, Sized b) => (a -> b) -> TrieMap k a -> TrieMap k b +fmapM = mapWithKeyM . const assocsM :: TrieKey k => TrieMap k a -> [(k, a)] -assocsM m = foldWithKeyM (\ k a xs -> (k, a):xs) m [] - -insertM :: TrieKey k => Sized a -> k -> a -> TrieMap k a -> TrieMap k a -insertM s = insertWithKeyM s (const const) - -insertWithKeyM :: TrieKey k => Sized a -> (k -> a -> a -> a) -> k -> a -> TrieMap k a -> TrieMap k a -insertWithKeyM s f k a = alterM s f' k where - f' = Just . maybe a (f k a) +assocsM m = build (\ f z -> foldrWithKeyM (\ k a xs -> (k, a) `f` xs) m z) -fromListM' :: TrieKey k => Sized a -> [(k, a)] -> TrieMap k a -fromListM' s = fromListM s (const const) --xs = foldr (uncurry insertM) emptyM xs +insertWithKeyM :: (TrieKey k, Sized a) => (k -> a -> a -> a) -> k -> a -> TrieMap k a -> TrieMap k a +insertWithKeyM f k a m = case searchM k m of + (# Nothing, hole #) -> assignM a hole + (# Just a', hole #) -> assignM (f k a a') hole unionMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a unionMaybe _ Nothing y = y @@ -119,17 +133,3 @@ subMaybe _ Nothing _ = True subMaybe (<=) (Just a) (Just b) = a <= b subMaybe _ _ _ = False - -aboutM :: (TrieKey k, Alternative t) => (k -> a -> t z) -> TrieMap k a -> t z -aboutM f = fst <.> extractM (const 0) (\ k a -> fmap (, Nothing) (f k a)) - -{-# RULES --- "lookupM/emptyM" forall k . lookupM k emptyM = Nothing; --- "sizeM/emptyM" forall s . sizeM s emptyM = 0; --- "traverseWithKeyM/emptyM" forall s f . traverseWithKeyM s f emptyM = pure emptyM; --- "extractM/emptyM" forall s f . extractM s f emptyM = empty; --- "foldWithKeyM/emptyM" forall f . foldWithKeyM f emptyM z = z; --- "foldlWithKeyM/emptyM" forall f . foldlWithKeyM f emptyM z = z; --- "lookupIxM/emptyM" forall s k . lookupIxM s k emptyM = (empty, empty, empty); --- "mapEitherM/emptyM" forall s1 s2 f . mapEitherM s1 s2 f emptyM = (emptyM, emptyM); - #-} diff --git a/Data/TrieMap/UnionMap.hs b/Data/TrieMap/UnionMap.hs --- a/Data/TrieMap/UnionMap.hs +++ b/Data/TrieMap/UnionMap.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE PatternGuards, UnboxedTuples, TypeFamilies, PatternGuards, ViewPatterns #-} +{-# LANGUAGE PatternGuards, UnboxedTuples, TypeFamilies, PatternGuards, ViewPatterns, MagicHash #-} {-# OPTIONS -funbox-strict-fields #-} module Data.TrieMap.UnionMap () where @@ -6,98 +6,125 @@ import Data.TrieMap.Sized import Control.Applicative +import Control.Monad -union :: (TrieKey k1, TrieKey k2) => Sized a -> TrieMap k1 a -> TrieMap k2 a -> TrieMap (Either k1 k2) a -union _ (nullM -> True) (nullM -> True) = Empty -union s m1@(sizeM s -> s1) m2@(sizeM s -> s2) = Union (s1 + s2) m1 m2 +import GHC.Exts -singletonMaybe :: (TrieKey k1, TrieKey k2) => Sized a -> Either k1 k2 -> Maybe a -> TrieMap (Either k1 k2) a -singletonMaybe s k a = maybe Empty (singletonM s k) a +(&) :: (TrieKey k1, TrieKey k2, Sized a) => TrieMap k1 a -> TrieMap k2 a -> TrieMap (Either k1 k2) a +m1 & m2 + | nullM m1, nullM m2 = Empty + | otherwise = Union (getSize# m1 +# getSize# m2) m1 m2 -singletonL :: (TrieKey k1, TrieKey k2) => Sized a -> k1 -> a -> TrieMap (Either k1 k2) a -singletonL s k a = Union (s a) (singletonM s k a) emptyM +singletonL :: (TrieKey k1, TrieKey k2, Sized a) => k1 -> a -> TrieMap (Either k1 k2) a +singletonL k a = Union (getSize# a) (singletonM k a) emptyM -singletonR :: (TrieKey k1, TrieKey k2) => Sized a -> k2 -> a -> TrieMap (Either k1 k2) a -singletonR s k a = Union (s a) emptyM (singletonM s k a) +singletonR :: (TrieKey k1, TrieKey k2, Sized a) => k2 -> a -> TrieMap (Either k1 k2) a +singletonR k a = Union (getSize# a) emptyM (singletonM k a) instance (TrieKey k1, TrieKey k2) => TrieKey (Either k1 k2) where - data TrieMap (Either k1 k2) a = Empty | Union !Int (TrieMap k1 a) (TrieMap k2 a) + data TrieMap (Either k1 k2) a = Empty | Union Int# (TrieMap k1 a) (TrieMap k2 a) + data Hole (Either k1 k2) a = + LHole (Hole k1 a) (TrieMap k2 a) + | RHole (TrieMap k1 a) (Hole k2 a) emptyM = Empty - singletonM s = either (singletonL s) (singletonR s) + singletonM = either singletonL singletonR nullM Empty = True nullM _ = False - sizeM _ Empty = 0 - sizeM _ (Union s _ _) = s + sizeM Empty = 0# + sizeM (Union s _ _) = s lookupM k (Union _ m1 m2) = either (`lookupM` m1) (`lookupM` m2) k lookupM _ _ = Nothing - - alterM s f k (Union _ m1 m2) = case k of - Left k -> union s (alterM s f k m1) m2 - Right k -> union s m1 (alterM s f k m2) - alterM s f k _ = singletonMaybe s k (f Nothing) - alterLookupM s f k Empty = onUnboxed (singletonMaybe s k) f Nothing - alterLookupM s f (Left k) (Union _ m1 m2) = onUnboxed (flip (union s) m2) (alterLookupM s f k) m1 - alterLookupM s f (Right k) (Union _ m1 m2) = onUnboxed (union s m1) (alterLookupM s f k) m2 - - traverseWithKeyM s f (Union _ m1 m2) = union s <$> traverseWithKeyM s (f . Left) m1 <*> traverseWithKeyM s (f . Right) m2 - traverseWithKeyM _ _ _ = pure Empty + traverseWithKeyM f (Union _ m1 m2) = (&) <$> traverseWithKeyM (f . Left) m1 <*> traverseWithKeyM (f . Right) m2 + traverseWithKeyM _ _ = pure Empty - foldWithKeyM f (Union _ m1 m2) = foldWithKeyM (f . Left) m1 . foldWithKeyM (f . Right) m2 - foldWithKeyM _ _ = id + foldrWithKeyM f (Union _ m1 m2) = foldrWithKeyM (f . Left) m1 . foldrWithKeyM (f . Right) m2 + foldrWithKeyM _ _ = id foldlWithKeyM f (Union _ m1 m2) = foldlWithKeyM (f . Right) m2 . foldlWithKeyM (f . Left) m1 foldlWithKeyM _ _ = id - mapMaybeM s f (Union _ m1 m2) = union s (mapMaybeM s (f . Left) m1) (mapMaybeM s (f . Right) m2) - mapMaybeM _ _ _ = Empty - - mapEitherM s1 s2 f (Union _ m1 m2) - | (# m1L, m1R #) <- mapEitherM s1 s2 (f . Left) m1, - (# m2L, m2R #) <- mapEitherM s1 s2 (f . Right) m2 - = (# union s1 m1L m2L, union s2 m1R m2R #) - mapEitherM _ _ _ _ = (# Empty, Empty #) + mapWithKeyM f (Union _ m1 m2) = mapWithKeyM (f . Left) m1 & mapWithKeyM (f . Right) m2 + mapWithKeyM _ _ = Empty - extractM s f (Union _ m1 m2) = let (&) = union s in fmap (& m2) <$> extractM s (f . Left) m1 <|> - fmap (m1 &) <$> extractM s (f . Right) m2 - extractM _ _ _ = empty + mapMaybeM f (Union _ m1 m2) = mapMaybeM (f . Left) m1 & mapMaybeM (f . Right) m2 + mapMaybeM _ _ = Empty - splitLookupM s f k (Union _ m1 m2) = let (&) = union s in case k of - Left k | (# m1L, x, m1R #) <- splitLookupM s f k m1 - -> (# m1L & emptyM, x, m1R & m2 #) - Right k | (# m2L, x, m2R #) <- splitLookupM s f k m2 - -> (# m1 & m2L, x, emptyM & m2R #) - splitLookupM _ _ _ _ = (# emptyM, Nothing, emptyM #) + mapEitherM f (Union _ m1 m2) + | (# m1L, m1R #) <- mapEitherM (f . Left) m1, + (# m2L, m2R #) <- mapEitherM (f . Right) m2 + = (# m1L & m2L, m1R & m2R #) + mapEitherM _ _ = (# Empty, Empty #) - unionM s f (Union _ m11 m12) (Union _ m21 m22) - = union s (unionM s (f . Left) m11 m21) (unionM s (f . Right) m12 m22) - unionM _ _ Empty m2 = m2 - unionM _ _ m1 Empty = m1 + unionM f (Union _ m11 m12) (Union _ m21 m22) + = unionM (f . Left) m11 m21 & unionM (f . Right) m12 m22 + unionM _ Empty m2 = m2 + unionM _ m1 Empty = m1 - isectM _ _ Empty _ = Empty - isectM _ _ _ Empty = Empty - isectM s f (Union _ m11 m12) (Union _ m21 m22) - = union s (isectM s (f . Left) m11 m21) (isectM s (f . Right) m12 m22) + isectM _ Empty _ = Empty + isectM _ _ Empty = Empty + isectM f (Union _ m11 m12) (Union _ m21 m22) + = isectM (f . Left) m11 m21 & isectM (f . Right) m12 m22 - diffM _ _ Empty _ = Empty - diffM _ _ m1 Empty = m1 - diffM s f (Union _ m11 m12) (Union _ m21 m22) - = union s (diffM s (f . Left) m11 m21) (diffM s (f . Right) m12 m22) + diffM _ Empty _ = Empty + diffM _ m1 Empty = m1 + diffM f (Union _ m11 m12) (Union _ m21 m22) + = diffM (f . Left) m11 m21 & diffM (f . Right) m12 m22 isSubmapM _ Empty _ = True isSubmapM (<=) (Union _ m11 m12) (Union _ m21 m22) = isSubmapM (<=) m11 m21 && isSubmapM (<=) m12 m22 isSubmapM _ Union{} Empty = False - fromListM s f = onPair (union s) (fromListM s (f . Left)) (fromListM s (f . Right)) . partEithers + fromListM f = onPair (&) (fromListM (f . Left)) (fromListM (f . Right)) . partEithers - fromAscListM s f = onPair (union s) (fromAscListM s (f . Left)) (fromAscListM s (f . Right)) . partEithers + fromAscListM f = onPair (&) (fromAscListM (f . Left)) (fromAscListM (f . Right)) . partEithers - fromDistAscListM s = onPair (union s) (fromDistAscListM s) (fromDistAscListM s) . partEithers + fromDistAscListM = onPair (&) fromDistAscListM fromDistAscListM . partEithers + + singleHoleM (Left k) = LHole (singleHoleM k) emptyM + singleHoleM (Right k) = RHole emptyM (singleHoleM k) + + keyM (LHole holeL _) = Left (keyM holeL) + keyM (RHole _ holeR) = Right (keyM holeR) + + beforeM a (LHole holeL _) = let mL = beforeM a holeL in + if nullM mL then Empty else Union (getSize# mL) mL emptyM + beforeM a (RHole mL holeR) = mL & beforeM a holeR + + afterM a (LHole holeL mR) = afterM a holeL & mR + afterM a (RHole _ holeR) = let mR = afterM a holeR in + if nullM mR then Empty else Union (getSize# mR) emptyM mR + + searchM k Empty = (# Nothing, singleHoleM k #) + searchM (Left k) (Union _ mL mR) = onUnboxed (`LHole` mR) (searchM k) mL + searchM (Right k) (Union _ mL mR) = onUnboxed (RHole mL) (searchM k) mR + + indexM i# (Union _ mL mR) + | i# <# sL#, (# i'#, v, holeL #) <- indexM i# mL + = (# i'#, v, LHole holeL mR #) + | (# i'#, v, holeR #) <- indexM (i# -# sL#) mR + = (# i'#, v, RHole mL holeR #) + where !sL# = getSize# mL + indexM _ _ = (# error err, error err, error err #) where + err = "Error: empty trie" + + extractHoleM (Union _ mL mR) = (do + (v, holeL) <- extractHoleM mL + return (v, LHole holeL mR)) `mplus` (do + (v, holeR) <- extractHoleM mR + return (v, RHole mL holeR)) + extractHoleM _ = mzero + + assignM v (LHole holeL mR) = assignM v holeL & mR + assignM v (RHole mL holeR) = mL & assignM v holeR + + clearM (LHole holeL mR) = clearM holeL & mR + clearM (RHole mL holeR) = mL & clearM holeR onPair :: (c -> d -> e) -> (a -> c) -> (b -> d) -> (a, b) -> e onPair f g h (a, b) = f (g a) (h b) diff --git a/Data/TrieMap/UnitMap.hs b/Data/TrieMap/UnitMap.hs --- a/Data/TrieMap/UnitMap.hs +++ b/Data/TrieMap/UnitMap.hs @@ -1,10 +1,12 @@ -{-# LANGUAGE TypeFamilies, UnboxedTuples #-} +{-# LANGUAGE TypeFamilies, UnboxedTuples, MagicHash #-} module Data.TrieMap.UnitMap where import Data.TrieMap.TrieKey +import Data.TrieMap.Sized import Control.Applicative +import Control.Monad import Data.Foldable import Data.Traversable @@ -14,25 +16,40 @@ instance TrieKey () where newtype TrieMap () a = Unit {getUnit :: Maybe a} + data Hole () a = Hole + emptyM = Unit Nothing - singletonM _ _ = Unit . Just + singletonM _ = Unit . Just nullM = isNothing . getUnit - sizeM s = maybe 0 s . getUnit + sizeM (Unit (Just a)) = getSize# a + sizeM _ = 0# lookupM _ (Unit m) = m - traverseWithKeyM _ f (Unit m) = Unit <$> traverse (f ()) m - foldWithKeyM f (Unit m) z = foldr (f ()) z m + traverseWithKeyM f (Unit m) = Unit <$> traverse (f ()) m + foldrWithKeyM f (Unit m) z = foldr (f ()) z m foldlWithKeyM f (Unit m) z = foldl (f ()) z m - mapMaybeM _ f (Unit m) = Unit (m >>= f ()) - mapEitherM _ _ f (Unit (Just a)) = both Unit Unit (f ()) a - mapEitherM _ _ _ _ = (# emptyM, emptyM #) - splitLookupM _ f _ (Unit (Just a)) = sides Unit f a - splitLookupM _ _ _ _ = (# emptyM, Nothing, emptyM #) - alterM _ f _ (Unit m) = Unit (f m) - alterLookupM _ f _ (Unit m) = onUnboxed Unit f m - unionM _ f (Unit m1) (Unit m2) = Unit (unionMaybe (f ()) m1 m2) - isectM _ f (Unit m1) (Unit m2) = Unit (isectMaybe (f ()) m1 m2) - diffM _ f (Unit m1) (Unit m2) = Unit (diffMaybe (f ()) m1 m2) - extractM _ f (Unit m) = maybe empty (fmap (fmap Unit) . f ()) m + mapWithKeyM f (Unit m) = Unit (f () <$> m) + mapMaybeM f (Unit m) = Unit (m >>= f ()) + mapEitherM f (Unit (Just a)) = both Unit Unit (f ()) a + mapEitherM _ _ = (# emptyM, emptyM #) + unionM f (Unit m1) (Unit m2) = Unit (unionMaybe (f ()) m1 m2) + isectM f (Unit m1) (Unit m2) = Unit (isectMaybe (f ()) m1 m2) + diffM f (Unit m1) (Unit m2) = Unit (diffMaybe (f ()) m1 m2) isSubmapM (<=) (Unit m1) (Unit m2) = subMaybe (<=) m1 m2 - fromListM _ _ [] = Unit Nothing - fromListM _ f ((_, v):xs) = Unit $ Just (foldl (\ v' -> f () v' . snd) v xs) + fromListM _ [] = Unit Nothing + fromListM f ((_, v):xs) = Unit $ Just (foldl (\ v' -> f () v' . snd) v xs) + + singleHoleM _ = Hole + keyM _ = () + beforeM a _ = Unit a + afterM a _ = Unit a + searchM _ (Unit m) = (# m, Hole #) + + indexM i (Unit (Just v)) = (# i, v, Hole #) + indexM _ _ = (# error err, error err, error err #) where + err = "Error: empty trie" + + extractHoleM (Unit (Just v)) = return (v, Hole) + extractHoleM _ = mzero + + assignM v _ = Unit (Just v) + clearM _ = emptyM diff --git a/Tests.hs b/Tests.hs --- a/Tests.hs +++ b/Tests.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE TemplateHaskell, TypeFamilies, GADTs, ExistentialQuantification, CPP #-} +{-# LANGUAGE TemplateHaskell, TypeFamilies, GADTs, ExistentialQuantification, CPP, ViewPatterns #-} -- module Tests where import Control.Monad @@ -7,10 +7,10 @@ import Test.QuickCheck import Prelude hiding (null, lookup) -type Key = [String] -type Val = [String] +type Key = [Int] +type Val = [Int] -main = quickCheck (verify M.empty T.empty) +main = quickCheckWith stdArgs{maxSize = 800, maxSuccess = 800} (verify M.empty T.empty) instance Arbitrary Op where arbitrary = oneof [ @@ -22,12 +22,20 @@ liftM (Op . Delete) arbitrary, return (Op MinView), return (Op MaxView), - return (Op MapMaybe)] + return (Op MapMaybe), + liftM Op (liftM Union recurse), + liftM Op (liftM Isect recurse), + liftM (Op . ElemAt) (arbitrary `suchThat` (>= 0)), + liftM (Op . DeleteAt) (arbitrary `suchThat` (>= 0))] shrink (Op (Insert k v)) = [Op (Insert k' v') | k' <- shrink k, v' <- shrink v] shrink (Op (Lookup k)) = map (Op . Lookup) (shrink k) shrink (Op (Delete k)) = map (Op . Delete) (shrink k) + shrink (Op (Union ops)) = map (Op . Union) (shrink ops) shrink _ = [] +recurse :: Gen [Op] +recurse = sized (\ n -> resize (n `quot` 5) arbitrary) + data Op = forall r . Op (Operation r) instance Show Op where @@ -40,6 +48,9 @@ show (Op MinView) = "MinView" show (Op MaxView) = "MaxView" show (Op MapMaybe) = "MapMaybe" + show (Op (Union ops)) = "Union " ++ show ops + show (Op (DeleteAt i)) = "DeleteAt " ++ show i + show (Op (ElemAt i)) = "ElemAt " ++ show i data Operation r where Insert :: Key -> Val -> Operation () @@ -51,11 +62,29 @@ MinView :: Operation (Maybe (Key, Val)) MaxView :: Operation (Maybe (Key, Val)) MapMaybe :: Operation () + Union :: [Op] -> Operation () + Isect :: [Op] -> Operation () + DeleteAt :: Int -> Operation () + ElemAt :: Int -> Operation (Maybe (Key, Val)) +mapFunc :: Key -> Val -> Val +mapFunc = (++) + +mapMaybeFunc :: Key -> Val -> Maybe Val +mapMaybeFunc (k:ks) xs + | even k = Just (ks ++ xs) +mapMaybeFunc _ _ = Nothing + +isectFunc :: Key -> Val -> Val -> Val +isectFunc ks xs ys = ks ++ xs ++ ys + +generateMap :: M.Map Key Val -> [Op] -> M.Map Key Val +generateMap = foldl (\ mm (Op op) -> snd (operateMap mm op)) + operateMap :: M.Map Key Val -> Operation r -> (r, M.Map Key Val) operateMap m (Insert k v) = ((), M.insert k v m) operateMap m (Lookup k) = (M.lookup k m, m) -operateMap m Map = ((), M.mapWithKey (\ k a -> k ++ a) m) +operateMap m Map = ((), M.mapWithKey mapFunc m) operateMap m ToList = (M.assocs m, m) operateMap m Size = (M.size m, m) operateMap m (Delete k) = ((), M.delete k m) @@ -65,14 +94,20 @@ operateMap m MaxView = case M.maxViewWithKey m of Nothing -> (Nothing, m) Just (kv, m') -> (Just kv, m') -operateMap m MapMaybe = ((), M.mapMaybeWithKey f m) - where f ("":xs) ("":ys) = Just (xs ++ ys) - f _ _ = Nothing +operateMap m MapMaybe = ((), M.mapMaybeWithKey mapMaybeFunc m) +operateMap m (Union ops) = + let m' = generateMap M.empty ops in ((), M.union m m') +operateMap m (DeleteAt i) = if M.null m then ((), m) else ((), M.deleteAt (i `mod` M.size m) m) +operateMap m (ElemAt i) = if M.null m then (Nothing, m) else (Just $ M.elemAt (i `mod` M.size m) m, m) +operateMap m (Isect ops) = ((), M.intersectionWithKey isectFunc m (generateMap M.empty ops)) +generateTMap :: T.TMap Key Val -> [Op] -> T.TMap Key Val +generateTMap = foldl (\ m (Op op) -> snd (operateTMap m op)) + operateTMap :: T.TMap Key Val -> Operation r -> (r, T.TMap Key Val) operateTMap m (Insert k v) = ((), T.insert k v m) operateTMap m (Lookup k) = (T.lookup k m, m) -operateTMap m Map = ((), T.mapWithKey (\ k a -> k ++ a) m) +operateTMap m Map = ((), T.mapWithKey mapFunc m) operateTMap m ToList = (T.assocs m, m) operateTMap m Size = (T.size m, m) operateTMap m (Delete k) = ((), T.delete k m) @@ -82,9 +117,15 @@ operateTMap m MaxView = case T.maxViewWithKey m of Nothing -> (Nothing, m) Just (kv, m') -> (Just kv, m') -operateTMap m MapMaybe = ((), T.mapMaybeWithKey f m) - where f ("":xs) ("":ys) = Just (xs ++ ys) - f _ _ = Nothing +operateTMap m MapMaybe = ((), T.mapMaybeWithKey mapMaybeFunc m) +operateTMap m (Union ops) = ((), T.union m $ generateTMap T.empty ops) +operateTMap m (Isect ops) = ((), T.intersectionWithKey isectFunc m (generateTMap T.empty ops)) +operateTMap m (DeleteAt i) + | T.null m = ((), m) + | otherwise = ((), T.deleteAt (i `mod` T.size m) m) +operateTMap m (ElemAt i) + | T.null m = (Nothing, m) + | otherwise = (Just $ T.elemAt (i `mod` T.size m) m, m) #define VERIFYOP(operation) verifyOp op@operation{} m tm = \ case (operateMap m op, operateTMap tm op) of \ @@ -100,6 +141,10 @@ VERIFYOP(MinView) VERIFYOP(MaxView) VERIFYOP(MapMaybe) +VERIFYOP(Union) +VERIFYOP(DeleteAt) +VERIFYOP(ElemAt) +VERIFYOP(Isect) verify :: M.Map Key Val -> T.TMap Key Val -> [Op] -> Bool verify m tm (Op op:ops) = case verifyOp op m tm of diff --git a/TrieMap.cabal b/TrieMap.cabal --- a/TrieMap.cabal +++ b/TrieMap.cabal @@ -1,9 +1,10 @@ name: TrieMap -version: 1.0.0 +version: 1.5.0 tested-with: GHC category: Algorithms -synopsis: Automatic type inference of generalized tries. -description: Builds on the multirec library to create a system capable of automatic or simple generalized trie type inference. +synopsis: Automatic type inference of generalized tries with Template Haskell. +description: Provides a efficient and compact implementation of generalized tries, and Template Haskell tools to generate + the necessary translation code. This is meant as a drop-in replacement for Data.Map. license: BSD3 license-file: LICENSE author: Louis Wasserman @@ -15,11 +16,11 @@ exposed-modules: Data.TrieMap, Data.TrieSet, - Data.TrieMap.Class, Data.TrieMap.Representation, Data.TrieMap.Representation.TH, Data.TrieMap.Modifiers other-modules: + Data.TrieMap.Class, Data.TrieMap.Class.Instances, Data.TrieMap.Key, Data.TrieMap.TrieKey,