diff --git a/Data/IntMap.hs b/Data/IntMap.hs --- a/Data/IntMap.hs +++ b/Data/IntMap.hs @@ -39,7 +39,7 @@ -- -- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\", -- Workshop on ML, September 1998, pages 77-86, --- +-- -- -- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve -- Information Coded In Alphanumeric/\", Journal of the ACM, 15(4), diff --git a/Data/IntMap/Base.hs b/Data/IntMap/Base.hs --- a/Data/IntMap/Base.hs +++ b/Data/IntMap/Base.hs @@ -1,16 +1,18 @@ {-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} #if __GLASGOW_HASKELL__ {-# LANGUAGE MagicHash, DeriveDataTypeable, StandaloneDeriving #-} +{-# LANGUAGE ScopedTypeVariables #-} #endif #if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 {-# LANGUAGE Trustworthy #-} #endif -{-# LANGUAGE ScopedTypeVariables #-} #if __GLASGOW_HASKELL__ >= 708 {-# LANGUAGE TypeFamilies #-} #endif #include "containers.h" +{-# OPTIONS_HADDOCK hide #-} ----------------------------------------------------------------------------- -- | @@ -22,6 +24,20 @@ -- Stability : provisional -- Portability : portable -- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. +-- +-- = Description +-- -- This defines the data structures and core (hidden) manipulations -- on representations. ----------------------------------------------------------------------------- @@ -89,6 +105,7 @@ , updateWithKey , updateLookupWithKey , alter + , alterF -- * Combine @@ -162,6 +179,8 @@ -- * Filter , filter , filterWithKey + , restrictKeys + , withoutKeys , partition , partitionWithKey @@ -206,6 +225,8 @@ , intFromNat , link , bin + , binCheckLeft + , binCheckRight , zero , nomatch , match @@ -244,6 +265,9 @@ import Data.Data (Data(..), Constr, mkConstr, constrIndex, Fixity(Prefix), DataType, mkDataType) import GHC.Exts (build) +#if !MIN_VERSION_base(4,8,0) +import Data.Functor ((<$)) +#endif #if __GLASGOW_HASKELL__ >= 708 import qualified GHC.Exts as GHCExts #endif @@ -294,7 +318,7 @@ -- > fromList [(5,'a'), (3,'b')] ! 5 == 'a' (!) :: IntMap a -> Key -> a -m ! k = find k m +(!) m k = find k m -- | Same as 'difference'. (\\) :: IntMap a -> IntMap b -> IntMap a @@ -349,8 +373,7 @@ toList = elems -- NB: Foldable.toList /= IntMap.toList {-# INLINE toList #-} elem = go - where STRICT_1_OF_2(go) - go _ Nil = False + where go !_ Nil = False go x (Tip _ y) = x == y go x (Bin _ _ l r) = go x l || go x r {-# INLINABLE elem #-} @@ -359,8 +382,7 @@ start (Tip _ y) = y start (Bin _ _ l r) = go (start l) r - STRICT_1_OF_2(go) - go m Nil = m + go !m Nil = m go m (Tip _ y) = max m y go m (Bin _ _ l r) = go (go m l) r {-# INLINABLE maximum #-} @@ -369,8 +391,7 @@ start (Tip _ y) = y start (Bin _ _ l r) = go (start l) r - STRICT_1_OF_2(go) - go m Nil = m + go !m Nil = m go m (Tip _ y) = min m y go m (Bin _ _ l r) = go (go m l) r {-# INLINABLE minimum #-} @@ -434,11 +455,9 @@ -- > size (singleton 1 'a') == 1 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3 size :: IntMap a -> Int -size t - = case t of - Bin _ _ l r -> size l + size r - Tip _ _ -> 1 - Nil -> 0 +size (Bin _ _ l r) = size l + size r +size (Tip _ _) = 1 +size Nil = 0 -- | /O(min(n,W))/. Is the key a member of the map? -- @@ -447,7 +466,7 @@ -- See Note: Local 'go' functions and capturing] member :: Key -> IntMap a -> Bool -member k = k `seq` go +member !k = go where go (Bin p m l r) | nomatch k p m = False | zero k m = go l @@ -467,7 +486,7 @@ -- See Note: Local 'go' functions and capturing] lookup :: Key -> IntMap a -> Maybe a -lookup k = k `seq` go +lookup !k = go where go (Bin p m l r) | nomatch k p m = Nothing | zero k m = go l @@ -479,7 +498,7 @@ -- See Note: Local 'go' functions and capturing] find :: Key -> IntMap a -> a -find k = k `seq` go +find !k = go where go (Bin p m l r) | nomatch k p m = not_found | zero k m = go l @@ -499,7 +518,7 @@ -- See Note: Local 'go' functions and capturing] findWithDefault :: a -> Key -> IntMap a -> a -findWithDefault def k = k `seq` go +findWithDefault def !k = go where go (Bin p m l r) | nomatch k p m = def | zero k m = go l @@ -516,7 +535,7 @@ -- See Note: Local 'go' functions and capturing. lookupLT :: Key -> IntMap a -> Maybe (Key, a) -lookupLT k t = k `seq` case t of +lookupLT !k t = case t of Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r _ -> go Nil t where @@ -535,7 +554,7 @@ -- See Note: Local 'go' functions and capturing. lookupGT :: Key -> IntMap a -> Maybe (Key, a) -lookupGT k t = k `seq` case t of +lookupGT !k t = case t of Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r _ -> go Nil t where @@ -555,7 +574,7 @@ -- See Note: Local 'go' functions and capturing. lookupLE :: Key -> IntMap a -> Maybe (Key, a) -lookupLE k t = k `seq` case t of +lookupLE !k t = case t of Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r _ -> go Nil t where @@ -575,7 +594,7 @@ -- See Note: Local 'go' functions and capturing. lookupGE :: Key -> IntMap a -> Maybe (Key, a) -lookupGE k t = k `seq` case t of +lookupGE !k t = case t of Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r _ -> go Nil t where @@ -637,16 +656,14 @@ -- > insert 5 'x' empty == singleton 5 'x' insert :: Key -> a -> IntMap a -> IntMap a -insert k x t = k `seq` - case t of - Bin p m l r - | nomatch k p m -> link k (Tip k x) p t - | zero k m -> Bin p m (insert k x l) r - | otherwise -> Bin p m l (insert k x r) - Tip ky _ - | k==ky -> Tip k x - | otherwise -> link k (Tip k x) ky t - Nil -> Tip k x +insert !k x t@(Bin p m l r) + | nomatch k p m = link k (Tip k x) p t + | zero k m = Bin p m (insert k x l) r + | otherwise = Bin p m l (insert k x r) +insert k x t@(Tip ky _) + | k==ky = Tip k x + | otherwise = link k (Tip k x) ky t +insert k x Nil = Tip k x -- right-biased insertion, used by 'union' -- | /O(min(n,W))/. Insert with a combining function. @@ -675,16 +692,14 @@ -- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a -insertWithKey f k x t = k `seq` - case t of - Bin p m l r - | nomatch k p m -> link k (Tip k x) p t - | zero k m -> Bin p m (insertWithKey f k x l) r - | otherwise -> Bin p m l (insertWithKey f k x r) - Tip ky y - | k==ky -> Tip k (f k x y) - | otherwise -> link k (Tip k x) ky t - Nil -> Tip k x +insertWithKey f !k x t@(Bin p m l r) + | nomatch k p m = link k (Tip k x) p t + | zero k m = Bin p m (insertWithKey f k x l) r + | otherwise = Bin p m l (insertWithKey f k x r) +insertWithKey f k x t@(Tip ky y) + | k == ky = Tip k (f k x y) + | otherwise = link k (Tip k x) ky t +insertWithKey _ k x Nil = Tip k x -- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@) -- is a pair where the first element is equal to (@'lookup' k map@) @@ -702,16 +717,14 @@ -- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a) -insertLookupWithKey f k x t = k `seq` - case t of - Bin p m l r - | nomatch k p m -> (Nothing,link k (Tip k x) p t) - | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r) - | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r') - Tip ky y - | k==ky -> (Just y,Tip k (f k x y)) - | otherwise -> (Nothing,link k (Tip k x) ky t) - Nil -> (Nothing,Tip k x) +insertLookupWithKey f !k x t@(Bin p m l r) + | nomatch k p m = (Nothing,link k (Tip k x) p t) + | zero k m = let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r) + | otherwise = let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r') +insertLookupWithKey f k x t@(Tip ky y) + | k == ky = (Just y,Tip k (f k x y)) + | otherwise = (Nothing,link k (Tip k x) ky t) +insertLookupWithKey _ k x Nil = (Nothing,Tip k x) {-------------------------------------------------------------------- @@ -725,16 +738,14 @@ -- > delete 5 empty == empty delete :: Key -> IntMap a -> IntMap a -delete k t = k `seq` - case t of - Bin p m l r - | nomatch k p m -> t - | zero k m -> bin p m (delete k l) r - | otherwise -> bin p m l (delete k r) - Tip ky _ - | k==ky -> Nil - | otherwise -> t - Nil -> Nil +delete !k t@(Bin p m l r) + | nomatch k p m = t + | zero k m = binCheckLeft p m (delete k l) r + | otherwise = binCheckRight p m l (delete k r) +delete k t@(Tip ky _) + | k == ky = Nil + | otherwise = t +delete _k Nil = Nil -- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not -- a member of the map, the original map is returned. @@ -756,9 +767,16 @@ -- > adjustWithKey f 7 empty == empty adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a -adjustWithKey f - = updateWithKey (\k' x -> Just (f k' x)) +adjustWithKey f !k t@(Bin p m l r) + | nomatch k p m = t + | zero k m = Bin p m (adjustWithKey f k l) r + | otherwise = Bin p m l (adjustWithKey f k r) +adjustWithKey f k t@(Tip ky y) + | k == ky = Tip ky (f k y) + | otherwise = t +adjustWithKey _ _ Nil = Nil + -- | /O(min(n,W))/. 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@. @@ -782,18 +800,16 @@ -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a -updateWithKey f k t = k `seq` - case t of - Bin p m l r - | nomatch k p m -> t - | zero k m -> bin p m (updateWithKey f k l) r - | otherwise -> bin p m l (updateWithKey f k r) - Tip ky y - | k==ky -> case (f k y) of +updateWithKey f !k t@(Bin p m l r) + | nomatch k p m = t + | zero k m = binCheckLeft p m (updateWithKey f k l) r + | otherwise = binCheckRight p m l (updateWithKey f k r) +updateWithKey f k t@(Tip ky y) + | k == ky = case (f k y) of Just y' -> Tip ky y' Nothing -> Nil - | otherwise -> t - Nil -> Nil + | otherwise = t +updateWithKey _ _ Nil = Nil -- | /O(min(n,W))/. Lookup and update. -- The function returns original value, if it is updated. @@ -806,18 +822,16 @@ -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a) -updateLookupWithKey f k t = k `seq` - case t of - Bin p m l r - | nomatch k p m -> (Nothing,t) - | zero k m -> let (found,l') = updateLookupWithKey f k l in (found,bin p m l' r) - | otherwise -> let (found,r') = updateLookupWithKey f k r in (found,bin p m l r') - Tip ky y - | k==ky -> case (f k y) of - Just y' -> (Just y,Tip ky y') - Nothing -> (Just y,Nil) - | otherwise -> (Nothing,t) - Nil -> (Nothing,Nil) +updateLookupWithKey f !k t@(Bin p m l r) + | nomatch k p m = (Nothing,t) + | zero k m = let !(found,l') = updateLookupWithKey f k l in (found,binCheckLeft p m l' r) + | otherwise = let !(found,r') = updateLookupWithKey f k r in (found,binCheckRight p m l r') +updateLookupWithKey f k t@(Tip ky y) + | k==ky = case (f k y) of + Just y' -> (Just y,Tip ky y') + Nothing -> (Just y,Nil) + | otherwise = (Nothing,t) +updateLookupWithKey _ _ Nil = (Nothing,Nil) @@ -825,26 +839,60 @@ -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'. -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a -alter f k t = k `seq` - case t of - Bin p m l r - | nomatch k p m -> case f Nothing of - Nothing -> t - Just x -> link k (Tip k x) p t - | zero k m -> bin p m (alter f k l) r - | otherwise -> bin p m l (alter f k r) - Tip ky y - | k==ky -> case f (Just y) of - Just x -> Tip ky x - Nothing -> Nil - | otherwise -> case f Nothing of - Just x -> link k (Tip k x) ky t - Nothing -> Tip ky y - Nil -> case f Nothing of - Just x -> Tip k x - Nothing -> Nil +alter f !k t@(Bin p m l r) + | nomatch k p m = case f Nothing of + Nothing -> t + Just x -> link k (Tip k x) p t + | zero k m = binCheckLeft p m (alter f k l) r + | otherwise = binCheckRight p m l (alter f k r) +alter f k t@(Tip ky y) + | k==ky = case f (Just y) of + Just x -> Tip ky x + Nothing -> Nil + | otherwise = case f Nothing of + Just x -> link k (Tip k x) ky t + Nothing -> Tip ky y +alter f k Nil = case f Nothing of + Just x -> Tip k x + Nothing -> Nil +-- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at +-- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete, +-- or update a value in an 'IntMap'. In short : @'lookup' k <$> 'alterF' f k m = f +-- ('lookup' k m)@. +-- +-- Example: +-- +-- @ +-- interactiveAlter :: Int -> IntMap String -> IO (IntMap String) +-- interactiveAlter k m = alterF f k m where +-- f Nothing -> do +-- putStrLn $ show k ++ +-- " was not found in the map. Would you like to add it?" +-- getUserResponse1 :: IO (Maybe String) +-- f (Just old) -> do +-- putStrLn "The key is currently bound to " ++ show old ++ +-- ". Would you like to change or delete it?" +-- getUserresponse2 :: IO (Maybe String) +-- @ +-- +-- 'alterF' is the most general operation for working with an individual +-- key that may or may not be in a given map. +-- +-- Note: 'alterF' is a flipped version of the 'at' combinator from +-- 'Control.Lens.At'. +-- +-- @since 0.5.8 +alterF :: Functor f + => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a) +-- This implementation was stolen from 'Control.Lens.At'. +alterF f k m = (<$> f mv) $ \fres -> + case fres of + Nothing -> maybe m (const (delete k m)) mv + Just v' -> insert k v' m + where mv = lookup k m + {-------------------------------------------------------------------- Union --------------------------------------------------------------------} @@ -929,7 +977,50 @@ differenceWithKey f m1 m2 = mergeWithKey f id (const Nil) m1 m2 +-- | Remove all the keys in a given set from a map. +-- +-- @ +-- m `withoutKeys` s = 'filterWithKey' (\k _ -> k `'IntSet.notMember'` s) m +-- @ +-- +-- @since 0.5.8 +withoutKeys :: IntMap a -> IntSet.IntSet -> IntMap a +withoutKeys = go + where + go t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2) + | shorter m1 m2 = merge1 + | shorter m2 m1 = merge2 + | p1 == p2 = bin p1 m1 (go l1 l2) (go r1 r2) + | otherwise = t1 + where + merge1 | nomatch p2 p1 m1 = t1 + | zero p2 m1 = binCheckLeft p1 m1 (go l1 t2) r1 + | otherwise = binCheckRight p1 m1 l1 (go r1 t2) + merge2 | nomatch p1 p2 m2 = t1 + | zero p1 m2 = bin p2 m2 (go t1 l2) Nil + | otherwise = bin p2 m2 Nil (go t1 r2) + go t1'@(Bin _ _ _ _) t2'@(IntSet.Tip k2' _) = merge t2' k2' t1' + where merge t2 k2 t1@(Bin p1 m1 l1 r1) | nomatch k2 p1 m1 = t1 + | zero k2 m1 = binCheckLeft p1 m1 (merge t2 k2 l1) r1 + | otherwise = binCheckRight p1 m1 l1 (merge t2 k2 r1) + merge _ k2 t1@(Tip k1 _) | k1 == k2 = Nil + | otherwise = t1 + merge _ _ Nil = Nil + + go t1@(Bin _ _ _ _) IntSet.Nil = t1 + + go t1'@(Tip k1' _) t2' = merge t1' k1' t2' + where merge t1 k1 (IntSet.Bin p2 m2 l2 r2) | nomatch k1 p2 m2 = t1 + | zero k1 m2 = bin p2 m2 (merge t1 k1 l2) Nil + | otherwise = bin p2 m2 Nil (merge t1 k1 r2) + merge t1 k1 (IntSet.Tip k2 _) | k1 == k2 = Nil + | otherwise = t1 + merge t1 _ IntSet.Nil = t1 + + go Nil _ = Nil + + {-------------------------------------------------------------------- Intersection --------------------------------------------------------------------} @@ -941,6 +1032,50 @@ intersection m1 m2 = mergeWithKey' bin const (const Nil) (const Nil) m1 m2 +-- | /O(n+m)/. The restriction of a map to the keys in a set. +-- +-- @ +-- m `restrictKeys` s = 'filterWithKey' (\k _ -> k `'IntSet.member'` s) m +-- @ +-- +-- @since 0.5.8 +restrictKeys :: IntMap a -> IntSet.IntSet -> IntMap a +restrictKeys = go + where + go t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2) + | shorter m1 m2 = merge1 + | shorter m2 m1 = merge2 + | p1 == p2 = bin p1 m1 (go l1 l2) (go r1 r2) + | otherwise = Nil + where + merge1 | nomatch p2 p1 m1 = Nil + | zero p2 m1 = bin p1 m1 (go l1 t2) Nil + | otherwise = bin p1 m1 Nil (go r1 t2) + merge2 | nomatch p1 p2 m2 = Nil + | zero p1 m2 = bin p2 m2 (go t1 l2) Nil + | otherwise = bin p2 m2 Nil (go t1 r2) + + go t1'@(Bin _ _ _ _) t2'@(IntSet.Tip k2' _) = merge t2' k2' t1' + where merge t2 k2 (Bin p1 m1 l1 r1) | nomatch k2 p1 m1 = Nil + | zero k2 m1 = bin p1 m1 (merge t2 k2 l1) Nil + | otherwise = bin p1 m1 Nil (merge t2 k2 r1) + merge _ k2 t1@(Tip k1 _) | k1 == k2 = t1 + | otherwise = Nil + merge _ _ Nil = Nil + + go (Bin _ _ _ _) IntSet.Nil = Nil + + go t1'@(Tip k1' _) t2' = merge t1' k1' t2' + where merge t1 k1 (IntSet.Bin p2 m2 l2 r2) + | nomatch k1 p2 m2 = Nil + | zero k1 m2 = bin p2 m2 (merge t1 k1 l2) Nil + | otherwise = bin p2 m2 Nil (merge t1 k1 r2) + merge t1 k1 (IntSet.Tip k2 _) | k1 == k2 = t1 + | otherwise = Nil + merge _ _ IntSet.Nil = Nil + + go Nil _ = Nil + -- | /O(n+m)/. The intersection with a combining function. -- -- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" @@ -1070,10 +1205,10 @@ updateMinWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a updateMinWithKey f t = - case t of Bin p m l r | m < 0 -> bin p m l (go f r) + case t of Bin p m l r | m < 0 -> binCheckRight p m l (go f r) _ -> go f t where - go f' (Bin p m l r) = bin p m (go f' l) r + go f' (Bin p m l r) = binCheckLeft p m (go f' l) r go f' (Tip k y) = case f' k y of Just y' -> Tip k y' Nothing -> Nil @@ -1086,10 +1221,10 @@ updateMaxWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a updateMaxWithKey f t = - case t of Bin p m l r | m < 0 -> bin p m (go f l) r + case t of Bin p m l r | m < 0 -> binCheckLeft p m (go f l) r _ -> go f t where - go f' (Bin p m l r) = bin p m l (go f' r) + go f' (Bin p m l r) = binCheckRight p m l (go f' r) go f' (Tip k y) = case f' k y of Just y' -> Tip k y' Nothing -> Nil @@ -1104,10 +1239,10 @@ maxViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a) maxViewWithKey t = case t of Nil -> Nothing - Bin p m l r | m < 0 -> case go l of (result, l') -> Just (result, bin p m l' r) + Bin p m l r | m < 0 -> case go l of (result, l') -> Just (result, binCheckLeft p m l' r) _ -> Just (go t) where - go (Bin p m l r) = case go r of (result, r') -> (result, bin p m l r') + go (Bin p m l r) = case go r of (result, r') -> (result, binCheckRight p m l r') go (Tip k y) = ((k, y), Nil) go Nil = error "maxViewWithKey Nil" @@ -1120,10 +1255,10 @@ minViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a) minViewWithKey t = case t of Nil -> Nothing - Bin p m l r | m < 0 -> case go r of (result, r') -> Just (result, bin p m l r') + Bin p m l r | m < 0 -> case go r of (result, r') -> Just (result, binCheckRight p m l r') _ -> Just (go t) where - go (Bin p m l r) = case go l of (result, l') -> (result, bin p m l' r) + go (Bin p m l r) = case go l of (result, l') -> (result, binCheckLeft p m l' r) go (Tip k y) = ((k, y), Nil) go Nil = error "minViewWithKey Nil" @@ -1302,11 +1437,11 @@ -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] map :: (a -> b) -> IntMap a -> IntMap b -map f t - = case t of - Bin p m l r -> Bin p m (map f l) (map f r) - Tip k x -> Tip k (f x) - Nil -> Nil +map f = go + where + go (Bin p m l r) = Bin p m (go l) (go r) + go (Tip k x) = Tip k (f x) + go Nil = Nil #ifdef __GLASGOW_HASKELL__ {-# NOINLINE [1] map #-} @@ -1579,10 +1714,10 @@ case t of Bin _ m l r | m < 0 -> if k >= 0 -- handle negative numbers. - then case go k l of (lt :*: gt) -> let lt' = union r lt - in lt' `seq` (lt', gt) - else case go k r of (lt :*: gt) -> let gt' = union gt l - in gt' `seq` (lt, gt') + then case go k l of (lt :*: gt) -> let !lt' = union r lt + in (lt', gt) + else case go k r of (lt :*: gt) -> let !gt' = union gt l + in (lt, gt') _ -> case go k t of (lt :*: gt) -> (lt, gt) where @@ -1609,19 +1744,19 @@ Bin _ m l r | m < 0 -> if k >= 0 -- handle negative numbers. then case go k l of - (lt, fnd, gt) -> let lt' = union r lt - in lt' `seq` (lt', fnd, gt) + (lt, fnd, gt) -> let !lt' = union r lt + in (lt', fnd, gt) else case go k r of - (lt, fnd, gt) -> let gt' = union gt l - in gt' `seq` (lt, fnd, gt') + (lt, fnd, gt) -> let !gt' = union gt l + in (lt, fnd, gt') _ -> go k t where go 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 = case go k' l of - (lt, fnd, gt) -> let gt' = union gt r in gt' `seq` (lt, fnd, gt') + (lt, fnd, gt) -> let !gt' = union gt r in (lt, fnd, gt') | otherwise = case go k' r of - (lt, fnd, gt) -> let lt' = union l lt in lt' `seq` (lt', fnd, gt) + (lt, fnd, gt) -> let !lt' = union l lt in (lt', fnd, gt) go k' t'@(Tip ky y) | k' > ky = (t', Nothing, Nil) | k' < ky = (Nil, Nothing, t') | otherwise = (Nil, Just y, Nil) @@ -1659,8 +1794,7 @@ | otherwise -> go (go z r) l _ -> go z t where - STRICT_1_OF_2(go) - go z' Nil = z' + go !z' Nil = z' go z' (Tip _ x) = f x z' go z' (Bin _ _ l r) = go (go z' r) l {-# INLINE foldr' #-} @@ -1694,8 +1828,7 @@ | otherwise -> go (go z l) r _ -> go z t where - STRICT_1_OF_2(go) - go z' Nil = z' + go !z' Nil = z' go z' (Tip _ x) = f z' x go z' (Bin _ _ l r) = go (go z' l) r {-# INLINE foldl' #-} @@ -1730,8 +1863,7 @@ | otherwise -> go (go z r) l _ -> go z t where - STRICT_1_OF_2(go) - go z' Nil = z' + go !z' Nil = z' go z' (Tip kx x) = f kx x z' go z' (Bin _ _ l r) = go (go z' r) l {-# INLINE foldrWithKey' #-} @@ -1766,8 +1898,7 @@ | otherwise -> go (go z l) r _ -> go z t where - STRICT_1_OF_2(go) - go z' Nil = z' + go !z' Nil = z' go z' (Tip kx x) = f z' kx x go z' (Bin _ _ l r) = go (go z' l) r {-# INLINE foldlWithKey' #-} @@ -1827,8 +1958,7 @@ keysSet (Bin p m l r) | m .&. IntSet.suffixBitMask == 0 = IntSet.Bin p m (keysSet l) (keysSet r) | otherwise = IntSet.Tip (p .&. IntSet.prefixBitMask) (computeBm (computeBm 0 l) r) - where STRICT_1_OF_2(computeBm) - computeBm acc (Bin _ _ l' r') = computeBm (computeBm acc l') r' + where computeBm !acc (Bin _ _ l' r') = computeBm (computeBm acc l') r' computeBm acc (Tip kx _) = acc .|. IntSet.bitmapOf kx computeBm _ Nil = error "Data.IntSet.keysSet: Nil" @@ -1849,7 +1979,7 @@ -- We split bmask into halves corresponding to left and right subtree. -- If they are both nonempty, we create a Bin node, otherwise exactly -- one of them is nonempty and we construct the IntMap from that half. - buildTree g prefix bmask bits = prefix `seq` bmask `seq` case bits of + buildTree g !prefix !bmask bits = case bits of 0 -> Tip prefix (g prefix) _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of bits2 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 -> @@ -2006,14 +2136,20 @@ -- -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +#if __GLASGOW_HASKELL__ fromDistinctAscList :: forall a. [(Key,a)] -> IntMap a +#else +fromDistinctAscList :: [(Key,a)] -> IntMap a +#endif fromDistinctAscList [] = Nil fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada where work (kx,vx) [] stk = finish kx (Tip kx vx) stk work (kx,vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk +#if __GLASGOW_HASKELL__ reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a +#endif reduce z zs _ px tx Nada = work z zs (Push px tx Nada) reduce z zs m px tx stk@(Push py ty stk') = let mxy = branchMask px py @@ -2067,6 +2203,12 @@ instance Functor IntMap where fmap = map +#ifdef __GLASGOW_HASKELL__ + a <$ Bin p m l r = Bin p m (a <$ l) (a <$ r) + a <$ Tip k _ = Tip k a + _ <$ Nil = Nil +#endif + {-------------------------------------------------------------------- Show --------------------------------------------------------------------} @@ -2097,7 +2239,7 @@ Typeable --------------------------------------------------------------------} -INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap") +INSTANCE_TYPEABLE1(IntMap) {-------------------------------------------------------------------- Helpers @@ -2123,6 +2265,17 @@ bin p m l r = Bin p m l r {-# INLINE bin #-} +-- binCheckLeft only checks that the left subtree is non-empty +binCheckLeft :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a +binCheckLeft _ _ Nil r = r +binCheckLeft p m l r = Bin p m l r +{-# INLINE binCheckLeft #-} + +-- binCheckRight only checks that the right subtree is non-empty +binCheckRight :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a +binCheckRight _ _ l Nil = l +binCheckRight p m l r = Bin p m l r +{-# INLINE binCheckRight #-} {-------------------------------------------------------------------- Endian independent bit twiddling diff --git a/Data/IntMap/Lazy.hs b/Data/IntMap/Lazy.hs --- a/Data/IntMap/Lazy.hs +++ b/Data/IntMap/Lazy.hs @@ -96,6 +96,7 @@ , updateWithKey , updateLookupWithKey , alter + , alterF -- * Combine @@ -168,6 +169,8 @@ -- * Filter , IM.filter , filterWithKey + , restrictKeys + , withoutKeys , partition , partitionWithKey diff --git a/Data/IntMap/Strict.hs b/Data/IntMap/Strict.hs --- a/Data/IntMap/Strict.hs +++ b/Data/IntMap/Strict.hs @@ -1,4 +1,5 @@ {-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} #if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 {-# LANGUAGE Trustworthy #-} #endif @@ -102,6 +103,7 @@ , updateWithKey , updateLookupWithKey , alter + , alterF -- * Combine @@ -174,6 +176,8 @@ -- * Filter , filter , filterWithKey + , restrictKeys + , withoutKeys , partition , partitionWithKey @@ -227,6 +231,7 @@ , updateWithKey , updateLookupWithKey , alter + , alterF , unionsWith , unionWith , unionWithKey @@ -266,6 +271,9 @@ #if __GLASGOW_HASKELL__ >= 709 import Data.Coerce #endif +#if !MIN_VERSION_base(4,8,0) +import Data.Functor((<$>)) +#endif -- $strictness -- @@ -298,7 +306,7 @@ -- See IntMap.Base.Note: Local 'go' functions and capturing] findWithDefault :: a -> Key -> IntMap a -> a -findWithDefault def k = k `seq` go +findWithDefault def !k = go where go (Bin p m l r) | nomatch k p m = def | zero k m = go l @@ -316,8 +324,8 @@ -- > size (singleton 1 'a') == 1 singleton :: Key -> a -> IntMap a -singleton k x - = x `seq` Tip k x +singleton k !x + = Tip k x {-# INLINE singleton #-} {-------------------------------------------------------------------- @@ -333,7 +341,7 @@ -- > insert 5 'x' empty == singleton 5 'x' insert :: Key -> a -> IntMap a -> IntMap a -insert k x t = k `seq` x `seq` +insert !k !x t = case t of Bin p m l r | nomatch k p m -> link k (Tip k x) p t @@ -374,7 +382,7 @@ -- in the result of @f@. insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a -insertWithKey f k x t = k `seq` +insertWithKey f !k x t = case t of Bin p m l r | nomatch k p m -> link k (singleton k x) p t @@ -401,7 +409,7 @@ -- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a) -insertLookupWithKey f0 k0 x0 t0 = k0 `seq` toPair $ go f0 k0 x0 t0 +insertLookupWithKey f0 !k0 x0 t0 = toPair $ go f0 k0 x0 t0 where go f k x t = case t of @@ -438,8 +446,16 @@ -- > adjustWithKey f 7 empty == empty adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a -adjustWithKey f - = updateWithKey (\k' x -> Just (f k' x)) +adjustWithKey f !k t = + case t of + Bin p m l r + | nomatch k p m -> t + | zero k m -> Bin p m (adjustWithKey f k l) r + | otherwise -> Bin p m l (adjustWithKey f k r) + Tip ky y + | k==ky -> Tip ky $! f k y + | otherwise -> t + Nil -> Nil -- | /O(min(n,W))/. 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 @@ -464,15 +480,15 @@ -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a -updateWithKey f k t = k `seq` +updateWithKey f !k t = case t of Bin p m l r | nomatch k p m -> t - | zero k m -> bin p m (updateWithKey f k l) r - | otherwise -> bin p m l (updateWithKey f k r) + | zero k m -> binCheckLeft p m (updateWithKey f k l) r + | otherwise -> binCheckRight p m l (updateWithKey f k r) Tip ky y | k==ky -> case f k y of - Just y' -> y' `seq` Tip ky y' + Just !y' -> Tip ky y' Nothing -> Nil | otherwise -> t Nil -> Nil @@ -488,18 +504,18 @@ -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a) -updateLookupWithKey f0 k0 t0 = k0 `seq` toPair $ go f0 k0 t0 +updateLookupWithKey f0 !k0 t0 = toPair $ go f0 k0 t0 where go f k t = case t of Bin p m l r | nomatch k p m -> (Nothing :*: t) - | zero k m -> let (found :*: l') = go f k l in (found :*: bin p m l' r) - | otherwise -> let (found :*: r') = go f k r in (found :*: bin p m l r') + | zero k m -> let (found :*: l') = go f k l in (found :*: binCheckLeft p m l' r) + | otherwise -> let (found :*: r') = go f k r in (found :*: binCheckRight p m l r') Tip ky y | k==ky -> case f k y of - Just y' -> y' `seq` (Just y :*: Tip ky y') - Nothing -> (Just y :*: Nil) + Just !y' -> (Just y :*: Tip ky y') + Nothing -> (Just y :*: Nil) | otherwise -> (Nothing :*: t) Nil -> (Nothing :*: Nil) @@ -509,26 +525,63 @@ -- 'alter' can be used to insert, delete, or update a value in an 'IntMap'. -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a -alter f k t = k `seq` +alter f !k t = case t of Bin p m l r | nomatch k p m -> case f Nothing of Nothing -> t - Just x -> x `seq` link k (Tip k x) p t - | zero k m -> bin p m (alter f k l) r - | otherwise -> bin p m l (alter f k r) + Just !x -> link k (Tip k x) p t + | zero k m -> binCheckLeft p m (alter f k l) r + | otherwise -> binCheckRight p m l (alter f k r) Tip ky y | k==ky -> case f (Just y) of - Just x -> x `seq` Tip ky x + Just !x -> Tip ky x Nothing -> Nil | otherwise -> case f Nothing of - Just x -> x `seq` link k (Tip k x) ky t + Just !x -> link k (Tip k x) ky t Nothing -> t Nil -> case f Nothing of - Just x -> x `seq` Tip k x + Just !x -> Tip k x Nothing -> Nil +-- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at +-- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete, +-- or update a value in an 'IntMap'. In short : @'lookup' k <$> 'alterF' f k m = f +-- ('lookup' k m)@. +-- +-- Example: +-- +-- @ +-- interactiveAlter :: Int -> IntMap String -> IO (IntMap String) +-- interactiveAlter k m = alterF f k m where +-- f Nothing -> do +-- putStrLn $ show k ++ +-- " was not found in the map. Would you like to add it?" +-- getUserResponse1 :: IO (Maybe String) +-- f (Just old) -> do +-- putStrLn "The key is currently bound to " ++ show old ++ +-- ". Would you like to change or delete it?" +-- getUserresponse2 :: IO (Maybe String) +-- @ +-- +-- 'alterF' is the most general operation for working with an individual +-- key that may or may not be in a given map. +-- Note: 'alterF' is a flipped version of the 'at' combinator from +-- 'Control.Lens.At'. +-- +-- @since 0.5.8 + +alterF :: Functor f + => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a) +-- This implementation was modified from 'Control.Lens.At'. +alterF f k m = (<$> f mv) $ \fres -> + case fres of + Nothing -> maybe m (const (delete k m)) mv + Just !v' -> insert k v' m + where mv = lookup k m + + {-------------------------------------------------------------------- Union --------------------------------------------------------------------} @@ -651,7 +704,7 @@ mergeWithKey f g1 g2 = mergeWithKey' bin combine g1 g2 where -- We use the lambda form to avoid non-exhaustive pattern matches warning. combine = \(Tip k1 x1) (Tip _k2 x2) -> case f k1 x1 x2 of Nothing -> Nil - Just x -> x `seq` Tip k1 x + Just !x -> Tip k1 x {-# INLINE combine #-} {-# INLINE mergeWithKey #-} @@ -666,12 +719,12 @@ updateMinWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a updateMinWithKey f t = - case t of Bin p m l r | m < 0 -> bin p m l (go f r) + case t of Bin p m l r | m < 0 -> binCheckRight p m l (go f r) _ -> go f t where - go f' (Bin p m l r) = bin p m (go f' l) r + go f' (Bin p m l r) = binCheckLeft p m (go f' l) r go f' (Tip k y) = case f' k y of - Just y' -> y' `seq` Tip k y' + Just !y' -> Tip k y' Nothing -> Nil go _ Nil = error "updateMinWithKey Nil" @@ -682,12 +735,12 @@ updateMaxWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a updateMaxWithKey f t = - case t of Bin p m l r | m < 0 -> bin p m (go f l) r + case t of Bin p m l r | m < 0 -> binCheckLeft p m (go f l) r _ -> go f t where - go f' (Bin p m l r) = bin p m l (go f' r) + go f' (Bin p m l r) = binCheckRight p m l (go f' r) go f' (Tip k y) = case f' k y of - Just y' -> y' `seq` Tip k y' + Just !y' -> Tip k y' Nothing -> Nil go _ Nil = error "updateMaxWithKey Nil" @@ -716,11 +769,11 @@ -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] map :: (a -> b) -> IntMap a -> IntMap b -map f t - = case t of - Bin p m l r -> Bin p m (map f l) (map f r) - Tip k x -> Tip k $! f x - Nil -> Nil +map f = go + where + go (Bin p m l r) = Bin p m (go l) (go r) + go (Tip k x) = Tip k $! f x + go Nil = Nil #ifdef __GLASGOW_HASKELL__ {-# NOINLINE [1] map #-} @@ -789,7 +842,7 @@ Bin p m l r -> let (a1 :*: l') = go f a l (a2 :*: r') = go f a1 r in (a2 :*: Bin p m l' r') - Tip k x -> let (a',x') = f a k x in x' `seq` (a' :*: Tip k x') + Tip k x -> let !(a',!x') = f a k x in (a' :*: Tip k x') Nil -> (a :*: Nil) -- | /O(n)/. The function @'mapAccumR'@ threads an accumulating @@ -802,7 +855,7 @@ Bin p m l r -> let (a1 :*: r') = go f a r (a2 :*: l') = go f a1 l in (a2 :*: Bin p m l' r') - Tip k x -> let (a',x') = f a k x in x' `seq` (a' :*: Tip k x') + Tip k x -> let !(a',!x') = f a k x in (a' :*: Tip k x') Nil -> (a :*: Nil) -- | /O(n*log n)/. @@ -838,7 +891,7 @@ mapMaybeWithKey f (Bin p m l r) = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r) mapMaybeWithKey f (Tip k x) = case f k x of - Just y -> y `seq` Tip k y + Just !y -> Tip k y Nothing -> Nil mapMaybeWithKey _ Nil = Nil @@ -873,8 +926,8 @@ (l1 :*: l2) = go f l (r1 :*: r2) = go f r go f (Tip k x) = case f k x of - Left y -> y `seq` (Tip k y :*: Nil) - Right z -> z `seq` (Nil :*: Tip k z) + Left !y -> (Tip k y :*: Nil) + Right !z -> (Nil :*: Tip k z) go _ Nil = (Nil :*: Nil) {-------------------------------------------------------------------- @@ -898,7 +951,7 @@ -- We split bmask into halves corresponding to left and right subtree. -- If they are both nonempty, we create a Bin node, otherwise exactly -- one of them is nonempty and we construct the IntMap from that half. - buildTree g prefix bmask bits = prefix `seq` bmask `seq` case bits of + buildTree g !prefix !bmask bits = case bits of 0 -> Tip prefix $! g prefix _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of bits2 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 -> @@ -976,7 +1029,7 @@ -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] combineEq z [] = [z] combineEq z@(kz,zz) (x@(kx,xx):xs) - | kx==kz = let yy = f kx xx zz in yy `seq` combineEq (kx,yy) xs + | kx==kz = let !yy = f kx xx zz in combineEq (kx,yy) xs | otherwise = z:combineEq x xs -- | /O(n)/. Build a map from a list of key\/value pairs where @@ -989,8 +1042,8 @@ fromDistinctAscList [] = Nil fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada where - work (kx,vx) [] stk = vx `seq` finish kx (Tip kx vx) stk - work (kx,vx) (z@(kz,_):zs) stk = vx `seq` reduce z zs (branchMask kx kz) kx (Tip kx vx) stk + work (kx,!vx) [] stk = finish kx (Tip kx vx) stk + work (kx,!vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a reduce z zs _ px tx Nada = work z zs (Push px tx Nada) diff --git a/Data/IntSet/Base.hs b/Data/IntSet/Base.hs --- a/Data/IntSet/Base.hs +++ b/Data/IntSet/Base.hs @@ -1,6 +1,7 @@ {-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} #if __GLASGOW_HASKELL__ -{-# LANGUAGE MagicHash, BangPatterns, DeriveDataTypeable, StandaloneDeriving #-} +{-# LANGUAGE MagicHash, DeriveDataTypeable, StandaloneDeriving #-} #endif #if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 {-# LANGUAGE Trustworthy #-} @@ -10,6 +11,7 @@ #endif #include "containers.h" +{-# OPTIONS_HADDOCK hide #-} ----------------------------------------------------------------------------- -- | @@ -21,6 +23,20 @@ -- Stability : provisional -- Portability : portable -- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. +-- +-- = Description +-- -- An efficient implementation of integer sets. -- -- These modules are intended to be imported qualified, to avoid name @@ -297,17 +313,15 @@ -- | /O(n)/. Cardinality of the set. size :: IntSet -> Int -size t - = case t of - Bin _ _ l r -> size l + size r - Tip _ bm -> bitcount 0 bm - Nil -> 0 +size (Bin _ _ l r) = size l + size r +size (Tip _ bm) = bitcount 0 bm +size Nil = 0 -- | /O(min(n,W))/. Is the value a member of the set? -- See Note: Local 'go' functions and capturing] member :: Key -> IntSet -> Bool -member x = x `seq` go +member !x = go where go (Bin p m l r) | nomatch x p m = False @@ -327,7 +341,7 @@ -- See Note: Local 'go' functions and capturing. lookupLT :: Key -> IntSet -> Maybe Key -lookupLT x t = x `seq` case t of +lookupLT !x t = case t of Bin _ m l r | m < 0 -> if x >= 0 then go r l else go Nil r _ -> go Nil t where @@ -348,7 +362,7 @@ -- See Note: Local 'go' functions and capturing. lookupGT :: Key -> IntSet -> Maybe Key -lookupGT x t = x `seq` case t of +lookupGT !x t = case t of Bin _ m l r | m < 0 -> if x >= 0 then go Nil l else go l r _ -> go Nil t where @@ -370,7 +384,7 @@ -- See Note: Local 'go' functions and capturing. lookupLE :: Key -> IntSet -> Maybe Key -lookupLE x t = x `seq` case t of +lookupLE !x t = case t of Bin _ m l r | m < 0 -> if x >= 0 then go r l else go Nil r _ -> go Nil t where @@ -392,7 +406,7 @@ -- See Note: Local 'go' functions and capturing. lookupGE :: Key -> IntSet -> Maybe Key -lookupGE x t = x `seq` case t of +lookupGE !x t = case t of Bin _ m l r | m < 0 -> if x >= 0 then go Nil l else go l r _ -> go Nil t where @@ -442,39 +456,35 @@ -- | /O(min(n,W))/. Add a value to the set. There is no left- or right bias for -- IntSets. insert :: Key -> IntSet -> IntSet -insert x = x `seq` insertBM (prefixOf x) (bitmapOf x) +insert !x = insertBM (prefixOf x) (bitmapOf x) -- Helper function for insert and union. insertBM :: Prefix -> BitMap -> IntSet -> IntSet -insertBM kx bm t = kx `seq` bm `seq` - case t of - Bin p m l r - | nomatch kx p m -> link kx (Tip kx bm) p t - | zero kx m -> Bin p m (insertBM kx bm l) r - | otherwise -> Bin p m l (insertBM kx bm r) - Tip kx' bm' - | kx' == kx -> Tip kx' (bm .|. bm') - | otherwise -> link kx (Tip kx bm) kx' t - Nil -> Tip kx bm +insertBM !kx !bm t@(Bin p m l r) + | nomatch kx p m = link kx (Tip kx bm) p t + | zero kx m = Bin p m (insertBM kx bm l) r + | otherwise = Bin p m l (insertBM kx bm r) +insertBM kx bm t@(Tip kx' bm') + | kx' == kx = Tip kx' (bm .|. bm') + | otherwise = link kx (Tip kx bm) kx' t +insertBM kx bm Nil = Tip kx bm -- | /O(min(n,W))/. Delete a value in the set. Returns the -- original set when the value was not present. delete :: Key -> IntSet -> IntSet -delete x = x `seq` deleteBM (prefixOf x) (bitmapOf x) +delete !x = deleteBM (prefixOf x) (bitmapOf x) -- Deletes all values mentioned in the BitMap from the set. -- Helper function for delete and difference. deleteBM :: Prefix -> BitMap -> IntSet -> IntSet -deleteBM kx bm t = kx `seq` bm `seq` - case t of - Bin p m l r - | nomatch kx p m -> t - | zero kx m -> bin p m (deleteBM kx bm l) r - | otherwise -> bin p m l (deleteBM kx bm r) - Tip kx' bm' - | kx' == kx -> tip kx (bm' .&. complement bm) - | otherwise -> t - Nil -> Nil +deleteBM !kx !bm t@(Bin p m l r) + | nomatch kx p m = t + | zero kx m = bin p m (deleteBM kx bm l) r + | otherwise = bin p m l (deleteBM kx bm r) +deleteBM kx bm t@(Tip kx' bm') + | kx' == kx = tip kx (bm' .&. complement bm) + | otherwise = t +deleteBM _ _ Nil = Nil {-------------------------------------------------------------------- @@ -687,10 +697,10 @@ case t of Bin _ m l r | m < 0 -> if x >= 0 -- handle negative numbers. - then case go x l of (lt :*: gt) -> let lt' = union lt r - in lt' `seq` (lt', gt) - else case go x r of (lt :*: gt) -> let gt' = union gt l - in gt' `seq` (lt, gt') + then case go x l of (lt :*: gt) -> let !lt' = union lt r + in (lt', gt) + else case go x r of (lt :*: gt) -> let !gt' = union gt l + in (lt, gt') _ -> case go x t of (lt :*: gt) -> (lt, gt) where @@ -718,11 +728,11 @@ case t of Bin _ m l r | m < 0 -> if x >= 0 then case go x l of - (lt, fnd, gt) -> let lt' = union lt r - in lt' `seq` (lt', fnd, gt) + (lt, fnd, gt) -> let !lt' = union lt r + in (lt', fnd, gt) else case go x r of - (lt, fnd, gt) -> let gt' = union gt l - in gt' `seq` (lt, fnd, gt') + (lt, fnd, gt) -> let !gt' = union gt l + in (lt, fnd, gt') _ -> go x t where go x' t'@(Bin p m l r) @@ -736,16 +746,15 @@ | kx' > x' = (Nil, False, t') -- equivalent to kx' > prefixOf x' | kx' < prefixOf x' = (t', False, Nil) - | otherwise = let lt = tip kx' (bm .&. lowerBitmap) - found = (bm .&. bitmapOfx') /= 0 - gt = tip kx' (bm .&. higherBitmap) - in lt `seq` found `seq` gt `seq` (lt, found, gt) + | otherwise = let !lt = tip kx' (bm .&. lowerBitmap) + !found = (bm .&. bitmapOfx') /= 0 + !gt = tip kx' (bm .&. higherBitmap) + in (lt, found, gt) where bitmapOfx' = bitmapOf x' lowerBitmap = bitmapOfx' - 1 higherBitmap = complement (lowerBitmap + bitmapOfx') go _ Nil = (Nil, False, Nil) - {---------------------------------------------------------------------- Min/Max ----------------------------------------------------------------------} @@ -875,8 +884,7 @@ | otherwise -> go (go z r) l _ -> go z t where - STRICT_1_OF_2(go) - go z' Nil = z' + go !z' Nil = z' go z' (Tip kx bm) = foldr'Bits kx f z' bm go z' (Bin _ _ l r) = go (go z' r) l {-# INLINE foldr' #-} @@ -907,8 +915,7 @@ | otherwise -> go (go z l) r _ -> go z t where - STRICT_1_OF_2(go) - go z' Nil = z' + go !z' Nil = z' go z' (Tip kx bm) = foldl'Bits kx f z' bm go z' (Bin _ _ l r) = go (go z' l) r {-# INLINE foldl' #-} @@ -1083,7 +1090,7 @@ Typeable --------------------------------------------------------------------} -INSTANCE_TYPEABLE0(IntSet,intSetTc,"IntSet") +INSTANCE_TYPEABLE0(IntSet) {-------------------------------------------------------------------- NFData @@ -1152,10 +1159,8 @@ | otherwise = id showsBars :: [String] -> ShowS -showsBars bars - = case bars of - [] -> id - _ -> showString (concat (reverse (tail bars))) . showString node +showsBars [] = id +showsBars bars = showString (concat (reverse (tail bars))) . showString node showsBitMap :: Word -> ShowS showsBitMap = showString . showBitMap @@ -1365,30 +1370,30 @@ highestBitSet x = indexOfTheOnlyBit (highestBitMask x) foldlBits prefix f z bitmap = go bitmap z - where go bm acc | bm == 0 = acc - | otherwise = case lowestBitMask bm of - bitmask -> bitmask `seq` case indexOfTheOnlyBit bitmask of - bi -> bi `seq` go (bm `xor` bitmask) ((f acc) $! (prefix+bi)) + where go 0 acc = acc + go bm acc = go (bm `xor` bitmask) ((f acc) $! (prefix+bi)) + where + !bitmask = lowestBitMask bm + !bi = indexOfTheOnlyBit bitmask foldl'Bits prefix f z bitmap = go bitmap z - where STRICT_2_OF_2(go) - go bm acc | bm == 0 = acc - | otherwise = case lowestBitMask bm of - bitmask -> bitmask `seq` case indexOfTheOnlyBit bitmask of - bi -> bi `seq` go (bm `xor` bitmask) ((f acc) $! (prefix+bi)) + where go 0 acc = acc + go bm !acc = go (bm `xor` bitmask) ((f acc) $! (prefix+bi)) + where !bitmask = lowestBitMask bm + !bi = indexOfTheOnlyBit bitmask foldrBits prefix f z bitmap = go (revNat bitmap) z - where go bm acc | bm == 0 = acc - | otherwise = case lowestBitMask bm of - bitmask -> bitmask `seq` case indexOfTheOnlyBit bitmask of - bi -> bi `seq` go (bm `xor` bitmask) ((f $! (prefix+(WORD_SIZE_IN_BITS-1)-bi)) acc) + where go 0 acc = acc + go bm acc = go (bm `xor` bitmask) ((f $! (prefix+(WORD_SIZE_IN_BITS-1)-bi)) acc) + where !bitmask = lowestBitMask bm + !bi = indexOfTheOnlyBit bitmask + foldr'Bits prefix f z bitmap = go (revNat bitmap) z - where STRICT_2_OF_2(go) - go bm acc | bm == 0 = acc - | otherwise = case lowestBitMask bm of - bitmask -> bitmask `seq` case indexOfTheOnlyBit bitmask of - bi -> bi `seq` go (bm `xor` bitmask) ((f $! (prefix+(WORD_SIZE_IN_BITS-1)-bi)) acc) + where go 0 acc = acc + go bm !acc = go (bm `xor` bitmask) ((f $! (prefix+(WORD_SIZE_IN_BITS-1)-bi)) acc) + where !bitmask = lowestBitMask bm + !bi = indexOfTheOnlyBit bitmask #else {---------------------------------------------------------------------- @@ -1418,30 +1423,26 @@ foldlBits prefix f z bm = let lb = lowestBitSet bm in go (prefix+lb) z (bm `shiftRL` lb) - where STRICT_1_OF_3(go) - go _ acc 0 = acc + where go !_ acc 0 = acc go bi acc n | n `testBit` 0 = go (bi + 1) (f acc bi) (n `shiftRL` 1) | otherwise = go (bi + 1) acc (n `shiftRL` 1) foldl'Bits prefix f z bm = let lb = lowestBitSet bm in go (prefix+lb) z (bm `shiftRL` lb) - where STRICT_1_OF_3(go) - STRICT_2_OF_3(go) - go _ acc 0 = acc + where go !_ !acc 0 = acc go bi acc n | n `testBit` 0 = go (bi + 1) (f acc bi) (n `shiftRL` 1) | otherwise = go (bi + 1) acc (n `shiftRL` 1) foldrBits prefix f z bm = let lb = lowestBitSet bm in go (prefix+lb) (bm `shiftRL` lb) - where STRICT_1_OF_2(go) - go _ 0 = z + where go !_ 0 = z go bi n | n `testBit` 0 = f bi (go (bi + 1) (n `shiftRL` 1)) | otherwise = go (bi + 1) (n `shiftRL` 1) foldr'Bits prefix f z bm = let lb = lowestBitSet bm in go (prefix+lb) (bm `shiftRL` lb) - where STRICT_1_OF_2(go) - go _ 0 = z + where + go !_ 0 = z go bi n | n `testBit` 0 = f bi $! go (bi + 1) (n `shiftRL` 1) | otherwise = go (bi + 1) (n `shiftRL` 1) @@ -1493,11 +1494,9 @@ -- splitting all the way to individual singleton sets -- it stops at some -- point. splitRoot :: IntSet -> [IntSet] -splitRoot orig = - case orig of - Nil -> [] - -- NOTE: we don't currently split below Tip, but we could. - x@(Tip _ _) -> [x] - Bin _ m l r | m < 0 -> [r, l] - | otherwise -> [l, r] +splitRoot Nil = [] +-- NOTE: we don't currently split below Tip, but we could. +splitRoot x@(Tip _ _) = [x] +splitRoot (Bin _ m l r) | m < 0 = [r, l] + | otherwise = [l, r] {-# INLINE splitRoot #-} diff --git a/Data/Map.hs b/Data/Map.hs --- a/Data/Map.hs +++ b/Data/Map.hs @@ -36,11 +36,17 @@ -- * Stephen Adams, \"/Efficient sets: a balancing act/\", -- Journal of Functional Programming 3(4):553-562, October 1993, -- . --- -- * J. Nievergelt and E.M. Reingold, -- \"/Binary search trees of bounded balance/\", -- SIAM journal of computing 2(1), March 1973. -- +-- Bounds for 'union', 'intersection', and 'difference' are as given +-- by +-- +-- * Guy Blelloch, Daniel Ferizovic, and Yihan Sun, +-- \"/Just Join for Parallel Ordered Sets/\", +-- . +-- -- Note that the implementation is /left-biased/ -- the elements of a -- first argument are always preferred to the second, for example in -- 'union' or 'insert'. @@ -66,71 +72,61 @@ import Data.Map.Lazy import qualified Data.Map.Strict as Strict --- | /Deprecated./ As of version 0.5, replaced by 'Data.Map.Strict.insertWith'. --- --- /O(log n)/. Same as 'insertWith', but the value being inserted to the map is +-- | /O(log n)/. Same as 'insertWith', but the value being inserted to the map is -- evaluated to WHNF beforehand. -- -- For example, to update a counter: -- -- > insertWith' (+) k 1 m -- - +{-# DEPRECATED insertWith' "As of version 0.5, replaced by 'Data.Map.Strict.insertWith'." #-} insertWith' :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a insertWith' = Strict.insertWith -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE insertWith' #-} #else {-# INLINE insertWith' #-} #endif --- | /Deprecated./ As of version 0.5, replaced by --- 'Data.Map.Strict.insertWithKey'. --- --- /O(log n)/. Same as 'insertWithKey', but the value being inserted to the map is +-- | /O(log n)/. Same as 'insertWithKey', but the value being inserted to the map is -- evaluated to WHNF beforehand. - +{-# DEPRECATED insertWithKey' "As of version 0.5, replaced by 'Data.Map.Strict.insertWithKey'." #-} insertWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a -- We do not reuse Data.Map.Strict.insertWithKey, because it is stricter -- it -- forces evaluation of the given value. insertWithKey' = Strict.insertWithKey -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE insertWithKey' #-} #else {-# INLINE insertWithKey' #-} #endif --- | /Deprecated./ As of version 0.5, replaced by --- 'Data.Map.Strict.insertLookupWithKey'. --- --- /O(log n)/. Same as 'insertLookupWithKey', but the value being inserted to +-- | /O(log n)/. Same as 'insertLookupWithKey', but the value being inserted to -- the map is evaluated to WHNF beforehand. - +{-# DEPRECATED insertLookupWithKey' "As of version 0.5, replaced by 'Data.Map.Strict.insertLookupWithKey'." #-} insertLookupWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a) -- We do not reuse Data.Map.Strict.insertLookupWithKey, because it is stricter -- it -- forces evaluation of the given value. insertLookupWithKey' = Strict.insertLookupWithKey -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE insertLookupWithKey' #-} #else {-# INLINE insertLookupWithKey' #-} #endif --- | /Deprecated./ As of version 0.5, replaced by 'foldr'. --- --- /O(n)/. Fold the values in the map using the given right-associative +-- | /O(n)/. Fold the values in the map using the given right-associative -- binary operator. This function is an equivalent of 'foldr' and is present -- for compatibility only. +{-# DEPRECATED fold "As of version 0.5, replaced by 'foldr'." #-} fold :: (a -> b -> b) -> b -> Map k a -> b fold = foldr {-# INLINE fold #-} --- | /Deprecated./ As of version 0.4, replaced by 'foldrWithKey'. --- --- /O(n)/. Fold the keys and values in the map using the given right-associative +-- | /O(n)/. Fold the keys and values in the map using the given right-associative -- binary operator. This function is an equivalent of 'foldrWithKey' and is present -- for compatibility only. +{-# DEPRECATED foldWithKey "As of version 0.4, replaced by 'foldrWithKey'." #-} foldWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b foldWithKey = foldrWithKey {-# INLINE foldWithKey #-} diff --git a/Data/Map/Base.hs b/Data/Map/Base.hs --- a/Data/Map/Base.hs +++ b/Data/Map/Base.hs @@ -1,2862 +1,4069 @@ {-# LANGUAGE CPP #-} -#if __GLASGOW_HASKELL__ -{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} -#endif -#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 -{-# LANGUAGE Trustworthy #-} -#endif -#if __GLASGOW_HASKELL__ >= 708 -{-# LANGUAGE RoleAnnotations #-} -{-# LANGUAGE TypeFamilies #-} -#endif - -#include "containers.h" - ------------------------------------------------------------------------------ --- | --- Module : Data.Map.Base --- Copyright : (c) Daan Leijen 2002 --- (c) Andriy Palamarchuk 2008 --- License : BSD-style --- Maintainer : libraries@haskell.org --- Stability : provisional --- Portability : portable --- --- An efficient implementation of maps from keys to values (dictionaries). --- --- Since many function names (but not the type name) clash with --- "Prelude" names, this module is usually imported @qualified@, e.g. --- --- > import Data.Map (Map) --- > import qualified Data.Map as Map --- --- The implementation of 'Map' is based on /size balanced/ binary trees (or --- trees of /bounded balance/) as described by: --- --- * Stephen Adams, \"/Efficient sets: a balancing act/\", --- Journal of Functional Programming 3(4):553-562, October 1993, --- . --- --- * J. Nievergelt and E.M. Reingold, --- \"/Binary search trees of bounded balance/\", --- SIAM journal of computing 2(1), March 1973. --- --- Note that the implementation is /left-biased/ -- the elements of a --- first argument are always preferred to the second, for example in --- 'union' or 'insert'. --- --- Operation comments contain the operation time complexity in --- the Big-O notation . ------------------------------------------------------------------------------ - --- [Note: Using INLINABLE] --- ~~~~~~~~~~~~~~~~~~~~~~~ --- It is crucial to the performance that the functions specialize on the Ord --- type when possible. GHC 7.0 and higher does this by itself when it sees th --- unfolding of a function -- that is why all public functions are marked --- INLINABLE (that exposes the unfolding). - - --- [Note: Using INLINE] --- ~~~~~~~~~~~~~~~~~~~~ --- For other compilers and GHC pre 7.0, we mark some of the functions INLINE. --- We mark the functions that just navigate down the tree (lookup, insert, --- delete and similar). That navigation code gets inlined and thus specialized --- when possible. There is a price to pay -- code growth. The code INLINED is --- therefore only the tree navigation, all the real work (rebalancing) is not --- INLINED by using a NOINLINE. --- --- All methods marked INLINE have to be nonrecursive -- a 'go' function doing --- the real work is provided. - - --- [Note: Type of local 'go' function] --- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ --- If the local 'go' function uses an Ord class, it sometimes heap-allocates --- the Ord dictionary when the 'go' function does not have explicit type. --- In that case we give 'go' explicit type. But this slightly decrease --- performance, as the resulting 'go' function can float out to top level. - - --- [Note: Local 'go' functions and capturing] --- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ --- As opposed to Map, when 'go' function captures an argument, increased --- heap-allocation can occur: sometimes in a polymorphic function, the 'go' --- floats out of its enclosing function and then it heap-allocates the --- dictionary and the argument. Maybe it floats out too late and strictness --- analyzer cannot see that these could be passed on stack. --- --- For example, change 'member' so that its local 'go' function is not passing --- argument k and then look at the resulting code for hedgeInt. - - --- [Note: Order of constructors] --- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ --- The order of constructors of Map matters when considering performance. --- Currently in GHC 7.0, when type has 2 constructors, a forward conditional --- jump is made when successfully matching second constructor. Successful match --- of first constructor results in the forward jump not taken. --- On GHC 7.0, reordering constructors from Tip | Bin to Bin | Tip --- improves the benchmark by up to 10% on x86. - -module Data.Map.Base ( - -- * Map type - Map(..) -- instance Eq,Show,Read - - -- * Operators - , (!), (\\) - - -- * Query - , null - , size - , member - , notMember - , lookup - , findWithDefault - , lookupLT - , lookupGT - , lookupLE - , lookupGE - - -- * Construction - , empty - , singleton - - -- ** Insertion - , insert - , insertWith - , insertWithKey - , insertLookupWithKey - - -- ** Delete\/Update - , delete - , adjust - , adjustWithKey - , update - , updateWithKey - , updateLookupWithKey - , alter - - -- * Combine - - -- ** Union - , union - , unionWith - , unionWithKey - , unions - , unionsWith - - -- ** Difference - , difference - , differenceWith - , differenceWithKey - - -- ** Intersection - , intersection - , intersectionWith - , intersectionWithKey - - -- ** Universal combining function - , mergeWithKey - - -- * Traversal - -- ** Map - , map - , mapWithKey - , traverseWithKey - , mapAccum - , mapAccumWithKey - , mapAccumRWithKey - , mapKeys - , mapKeysWith - , mapKeysMonotonic - - -- * Folds - , foldr - , foldl - , foldrWithKey - , foldlWithKey - , foldMapWithKey - - -- ** Strict folds - , foldr' - , foldl' - , foldrWithKey' - , foldlWithKey' - - -- * Conversion - , elems - , keys - , assocs - , keysSet - , fromSet - - -- ** Lists - , toList - , fromList - , fromListWith - , fromListWithKey - - -- ** Ordered lists - , toAscList - , toDescList - , fromAscList - , fromAscListWith - , fromAscListWithKey - , fromDistinctAscList - - -- * Filter - , filter - , filterWithKey - , partition - , partitionWithKey - - , mapMaybe - , mapMaybeWithKey - , mapEither - , mapEitherWithKey - - , split - , splitLookup - , splitRoot - - -- * Submap - , isSubmapOf, isSubmapOfBy - , isProperSubmapOf, isProperSubmapOfBy - - -- * Indexed - , lookupIndex - , findIndex - , elemAt - , updateAt - , deleteAt - - -- * Min\/Max - , findMin - , findMax - , deleteMin - , deleteMax - , deleteFindMin - , deleteFindMax - , updateMin - , updateMax - , updateMinWithKey - , updateMaxWithKey - , minView - , maxView - , minViewWithKey - , maxViewWithKey - - -- * Debugging - , showTree - , showTreeWith - , valid - - -- Used by the strict version - , bin - , balance - , balanced - , balanceL - , balanceR - , delta - , link - , insertMax - , merge - , glue - , trim - , trimLookupLo - , MaybeS(..) - , filterGt - , filterLt - ) where - -#if !(MIN_VERSION_base(4,8,0)) -import Control.Applicative (Applicative(..), (<$>)) -import Data.Monoid (Monoid(..)) -import Data.Traversable (Traversable(traverse)) -#endif -#if MIN_VERSION_base(4,9,0) -import Data.Semigroup (Semigroup((<>), stimes), stimesIdempotentMonoid) -#endif - -import Control.DeepSeq (NFData(rnf)) -import Data.Bits (shiftL, shiftR) -import qualified Data.Foldable as Foldable -import Data.Typeable -import Prelude hiding (lookup, map, filter, foldr, foldl, null) - -import qualified Data.Set.Base as Set -import Data.Utils.StrictFold -import Data.Utils.StrictPair - -#if __GLASGOW_HASKELL__ -import GHC.Exts ( build ) -#if __GLASGOW_HASKELL__ >= 708 -import qualified GHC.Exts as GHCExts -#endif -import Text.Read -import Data.Data -#endif -#if __GLASGOW_HASKELL__ >= 709 -import Data.Coerce -#endif - - -{-------------------------------------------------------------------- - Operators ---------------------------------------------------------------------} -infixl 9 !,\\ -- - --- | /O(log n)/. Find the value at a key. --- Calls 'error' when the element can not be found. --- --- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map --- > fromList [(5,'a'), (3,'b')] ! 5 == 'a' - -(!) :: Ord k => Map k a -> k -> a -m ! k = find k m -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE (!) #-} -#endif - --- | Same as 'difference'. -(\\) :: Ord k => Map k a -> Map k b -> Map k a -m1 \\ m2 = difference m1 m2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE (\\) #-} -#endif - -{-------------------------------------------------------------------- - Size balanced trees. ---------------------------------------------------------------------} --- | A Map from keys @k@ to values @a@. - --- See Note: Order of constructors -data Map k a = Bin {-# UNPACK #-} !Size !k a !(Map k a) !(Map k a) - | Tip - -type Size = Int - -#if __GLASGOW_HASKELL__ >= 708 -type role Map nominal representational -#endif - -instance (Ord k) => Monoid (Map k v) where - mempty = empty - mconcat = unions -#if !(MIN_VERSION_base(4,9,0)) - mappend = union -#else - mappend = (<>) - -instance (Ord k) => Semigroup (Map k v) where - (<>) = union - stimes = stimesIdempotentMonoid -#endif - -#if __GLASGOW_HASKELL__ - -{-------------------------------------------------------------------- - A Data instance ---------------------------------------------------------------------} - --- This instance preserves data abstraction at the cost of inefficiency. --- We provide limited reflection services for the sake of data abstraction. - -instance (Data k, Data a, Ord k) => Data (Map k a) where - gfoldl f z m = z fromList `f` toList m - toConstr _ = fromListConstr - gunfold k z c = case constrIndex c of - 1 -> k (z fromList) - _ -> error "gunfold" - dataTypeOf _ = mapDataType - dataCast2 f = gcast2 f - -fromListConstr :: Constr -fromListConstr = mkConstr mapDataType "fromList" [] Prefix - -mapDataType :: DataType -mapDataType = mkDataType "Data.Map.Base.Map" [fromListConstr] - -#endif - -{-------------------------------------------------------------------- - Query ---------------------------------------------------------------------} --- | /O(1)/. Is the map empty? --- --- > Data.Map.null (empty) == True --- > Data.Map.null (singleton 1 'a') == False - -null :: Map k a -> Bool -null Tip = True -null (Bin {}) = False -{-# INLINE null #-} - --- | /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 :: Map k a -> Int -size Tip = 0 -size (Bin sz _ _ _ _) = sz -{-# INLINE size #-} - - --- | /O(log n)/. Lookup the value at a key in the map. --- --- The function will return the corresponding value as @('Just' value)@, --- or 'Nothing' if the key isn't in the map. --- --- An example of using @lookup@: --- --- > import Prelude hiding (lookup) --- > import Data.Map --- > --- > employeeDept = fromList([("John","Sales"), ("Bob","IT")]) --- > deptCountry = fromList([("IT","USA"), ("Sales","France")]) --- > countryCurrency = fromList([("USA", "Dollar"), ("France", "Euro")]) --- > --- > employeeCurrency :: String -> Maybe String --- > employeeCurrency name = do --- > dept <- lookup name employeeDept --- > country <- lookup dept deptCountry --- > lookup country countryCurrency --- > --- > main = do --- > putStrLn $ "John's currency: " ++ (show (employeeCurrency "John")) --- > putStrLn $ "Pete's currency: " ++ (show (employeeCurrency "Pete")) --- --- The output of this program: --- --- > John's currency: Just "Euro" --- > Pete's currency: Nothing -lookup :: Ord k => k -> Map k a -> Maybe a -lookup = go - where - STRICT_1_OF_2(go) - go _ Tip = Nothing - go k (Bin _ kx x l r) = case compare k kx of - LT -> go k l - GT -> go k r - EQ -> Just x -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE lookup #-} -#else -{-# INLINE lookup #-} -#endif - --- | /O(log n)/. 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 -member :: Ord k => k -> Map k a -> Bool -member = go - where - STRICT_1_OF_2(go) - go _ Tip = False - go k (Bin _ kx _ l r) = case compare k kx of - LT -> go k l - GT -> go k r - EQ -> True -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE member #-} -#else -{-# INLINE member #-} -#endif - --- | /O(log n)/. 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 - -notMember :: Ord k => k -> Map k a -> Bool -notMember k m = not $ member k m -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE notMember #-} -#else -{-# INLINE notMember #-} -#endif - --- | /O(log n)/. Find the value at a key. --- Calls 'error' when the element can not be found. -find :: Ord k => k -> Map k a -> a -find = go - where - STRICT_1_OF_2(go) - go _ Tip = error "Map.!: given key is not an element in the map" - go k (Bin _ kx x l r) = case compare k kx of - LT -> go k l - GT -> go k r - EQ -> x -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE find #-} -#else -{-# INLINE find #-} -#endif - --- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns --- the value at key @k@ or returns default value @def@ --- when the key is not in the map. --- --- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' --- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' -findWithDefault :: Ord k => a -> k -> Map k a -> a -findWithDefault = go - where - STRICT_2_OF_3(go) - go def _ Tip = def - go def k (Bin _ kx x l r) = case compare k kx of - LT -> go def k l - GT -> go def k r - EQ -> x -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE findWithDefault #-} -#else -{-# INLINE findWithDefault #-} -#endif - --- | /O(log n)/. Find largest key smaller than the given one and return the --- corresponding (key, value) pair. --- --- > lookupLT 3 (fromList [(3,'a'), (5,'b')]) == Nothing --- > lookupLT 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') -lookupLT :: Ord k => k -> Map k v -> Maybe (k, v) -lookupLT = goNothing - where - STRICT_1_OF_2(goNothing) - goNothing _ Tip = Nothing - goNothing k (Bin _ kx x l r) | k <= kx = goNothing k l - | otherwise = goJust k kx x r - - STRICT_1_OF_4(goJust) - goJust _ kx' x' Tip = Just (kx', x') - goJust k kx' x' (Bin _ kx x l r) | k <= kx = goJust k kx' x' l - | otherwise = goJust k kx x r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE lookupLT #-} -#else -{-# INLINE lookupLT #-} -#endif - --- | /O(log n)/. Find smallest key greater than the given one and return the --- corresponding (key, value) pair. --- --- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') --- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing -lookupGT :: Ord k => k -> Map k v -> Maybe (k, v) -lookupGT = goNothing - where - STRICT_1_OF_2(goNothing) - goNothing _ Tip = Nothing - goNothing k (Bin _ kx x l r) | k < kx = goJust k kx x l - | otherwise = goNothing k r - - STRICT_1_OF_4(goJust) - goJust _ kx' x' Tip = Just (kx', x') - goJust k kx' x' (Bin _ kx x l r) | k < kx = goJust k kx x l - | otherwise = goJust k kx' x' r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE lookupGT #-} -#else -{-# INLINE lookupGT #-} -#endif - --- | /O(log n)/. Find largest key smaller or equal to the given one and return --- the corresponding (key, value) pair. --- --- > lookupLE 2 (fromList [(3,'a'), (5,'b')]) == Nothing --- > lookupLE 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') --- > lookupLE 5 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') -lookupLE :: Ord k => k -> Map k v -> Maybe (k, v) -lookupLE = goNothing - where - STRICT_1_OF_2(goNothing) - goNothing _ Tip = Nothing - goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goNothing k l - EQ -> Just (kx, x) - GT -> goJust k kx x r - - STRICT_1_OF_4(goJust) - goJust _ kx' x' Tip = Just (kx', x') - goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx' x' l - EQ -> Just (kx, x) - GT -> goJust k kx x r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE lookupLE #-} -#else -{-# INLINE lookupLE #-} -#endif - --- | /O(log n)/. Find smallest key greater or equal to the given one and return --- the corresponding (key, value) pair. --- --- > lookupGE 3 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') --- > lookupGE 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') --- > lookupGE 6 (fromList [(3,'a'), (5,'b')]) == Nothing -lookupGE :: Ord k => k -> Map k v -> Maybe (k, v) -lookupGE = goNothing - where - STRICT_1_OF_2(goNothing) - goNothing _ Tip = Nothing - goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l - EQ -> Just (kx, x) - GT -> goNothing k r - - STRICT_1_OF_4(goJust) - goJust _ kx' x' Tip = Just (kx', x') - goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l - EQ -> Just (kx, x) - GT -> goJust k kx' x' r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE lookupGE #-} -#else -{-# INLINE lookupGE #-} -#endif - -{-------------------------------------------------------------------- - Construction ---------------------------------------------------------------------} --- | /O(1)/. The empty map. --- --- > empty == fromList [] --- > size empty == 0 - -empty :: Map k a -empty = Tip -{-# INLINE empty #-} - --- | /O(1)/. A map with a single element. --- --- > singleton 1 'a' == fromList [(1, 'a')] --- > size (singleton 1 'a') == 1 - -singleton :: k -> a -> Map k a -singleton k x = Bin 1 k x Tip Tip -{-# INLINE singleton #-} - -{-------------------------------------------------------------------- - Insertion ---------------------------------------------------------------------} --- | /O(log n)/. 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' - --- See Note: Type of local 'go' function -insert :: Ord k => k -> a -> Map k a -> Map k a -insert = go - where - go :: Ord k => k -> a -> Map k a -> Map k a - STRICT_1_OF_3(go) - go kx x Tip = singleton kx x - go kx x (Bin sz ky y l r) = - case compare kx ky of - LT -> balanceL ky y (go kx x l) r - GT -> balanceR ky y l (go kx x r) - EQ -> Bin sz kx x l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insert #-} -#else -{-# INLINE insert #-} -#endif - --- Insert a new key and value in the map if it is not already present. --- Used by `union`. - --- See Note: Type of local 'go' function -insertR :: Ord k => k -> a -> Map k a -> Map k a -insertR = go - where - go :: Ord k => k -> a -> Map k a -> Map k a - STRICT_1_OF_3(go) - go kx x Tip = singleton kx x - go kx x t@(Bin _ ky y l r) = - case compare kx ky of - LT -> balanceL ky y (go kx x l) r - GT -> balanceR ky y l (go kx x r) - EQ -> t -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertR #-} -#else -{-# INLINE insertR #-} -#endif - --- | /O(log n)/. 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" - -insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWith f = insertWithKey (\_ x' y' -> f x' y') -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertWith #-} -#else -{-# INLINE insertWith #-} -#endif - --- | /O(log n)/. 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" - --- See Note: Type of local 'go' function -insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWithKey = go - where - go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a - STRICT_2_OF_4(go) - go _ kx x Tip = singleton kx x - go f kx x (Bin sy ky y l r) = - case compare kx ky of - LT -> balanceL ky y (go f kx x l) r - GT -> balanceR ky y l (go f kx x r) - EQ -> Bin sy kx (f kx x y) l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertWithKey #-} -#else -{-# INLINE insertWithKey #-} -#endif - --- | /O(log n)/. 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") --- --- This is how to define @insertLookup@ using @insertLookupWithKey@: --- --- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t --- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) --- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) - --- See Note: Type of local 'go' function -insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a - -> (Maybe a, Map k a) -insertLookupWithKey = go - where - go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a) - STRICT_2_OF_4(go) - go _ kx x Tip = (Nothing, singleton kx x) - go f kx x (Bin sy ky y l r) = - case compare kx ky of - LT -> let (found, l') = go f kx x l - in (found, balanceL ky y l' r) - GT -> let (found, r') = go f kx x r - in (found, balanceR ky y l r') - EQ -> (Just y, Bin sy kx (f kx x y) l r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertLookupWithKey #-} -#else -{-# INLINE insertLookupWithKey #-} -#endif - -{-------------------------------------------------------------------- - Deletion ---------------------------------------------------------------------} --- | /O(log n)/. 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 - --- See Note: Type of local 'go' function -delete :: Ord k => k -> Map k a -> Map k a -delete = go - where - go :: Ord k => k -> Map k a -> Map k a - STRICT_1_OF_2(go) - go _ Tip = Tip - go k (Bin _ kx x l r) = - case compare k kx of - LT -> balanceR kx x (go k l) r - GT -> balanceL kx x l (go k r) - EQ -> glue l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE delete #-} -#else -{-# INLINE delete #-} -#endif - --- | /O(log n)/. 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 - -adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a -adjust f = adjustWithKey (\_ x -> f x) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE adjust #-} -#else -{-# INLINE adjust #-} -#endif - --- | /O(log n)/. 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 - -adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a -adjustWithKey f = updateWithKey (\k' x' -> Just (f k' x')) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE adjustWithKey #-} -#else -{-# INLINE adjustWithKey #-} -#endif - --- | /O(log n)/. 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 :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a -update f = updateWithKey (\_ x -> f x) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE update #-} -#else -{-# INLINE update #-} -#endif - --- | /O(log n)/. 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" - --- See Note: Type of local 'go' function -updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a -updateWithKey = go - where - go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a - STRICT_2_OF_3(go) - go _ _ Tip = Tip - go f k(Bin sx kx x l r) = - case compare k kx of - LT -> balanceR kx x (go f k l) r - GT -> balanceL kx x l (go f k r) - EQ -> case f kx x of - Just x' -> Bin sx kx x' l r - Nothing -> glue l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE updateWithKey #-} -#else -{-# INLINE updateWithKey #-} -#endif - --- | /O(log n)/. Lookup and update. See also 'updateWithKey'. --- The function returns changed value, if it is updated. --- Returns the original key value if the map entry is deleted. --- --- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing --- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "5:new a", fromList [(3, "b"), (5, "5:new a")]) --- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) --- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") - --- See Note: Type of local 'go' function -updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) -updateLookupWithKey = go - where - go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) - STRICT_2_OF_3(go) - go _ _ Tip = (Nothing,Tip) - go f k (Bin sx kx x l r) = - case compare k kx of - LT -> let (found,l') = go f k l in (found,balanceR kx x l' r) - GT -> let (found,r') = go f k r in (found,balanceL kx x l r') - EQ -> case f kx x of - Just x' -> (Just x',Bin sx kx x' l r) - Nothing -> (Just x,glue l r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE updateLookupWithKey #-} -#else -{-# INLINE updateLookupWithKey #-} -#endif - --- | /O(log n)/. 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 'Map'. --- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. --- --- > let f _ = Nothing --- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] --- > alter f 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" --- > --- > let f _ = Just "c" --- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "c")] --- > alter f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "c")] - --- See Note: Type of local 'go' function -alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a -alter = go - where - go :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a - STRICT_2_OF_3(go) - go f k Tip = case f Nothing of - Nothing -> Tip - Just x -> singleton k x - - go f k (Bin sx kx x l r) = case compare k kx of - LT -> balance kx x (go f k l) r - GT -> balance kx x l (go f k r) - EQ -> case f (Just x) of - Just x' -> Bin sx kx x' l r - Nothing -> glue l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE alter #-} -#else -{-# INLINE alter #-} -#endif - -{-------------------------------------------------------------------- - Indexing ---------------------------------------------------------------------} --- | /O(log n)/. Return the /index/ of a key, which is its zero-based index in --- the sequence sorted by keys. 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 - --- See Note: Type of local 'go' function -findIndex :: Ord k => k -> Map k a -> Int -findIndex = go 0 - where - go :: Ord k => Int -> k -> Map k a -> Int - STRICT_1_OF_3(go) - STRICT_2_OF_3(go) - go _ _ Tip = error "Map.findIndex: element is not in the map" - go idx k (Bin _ kx _ l r) = case compare k kx of - LT -> go idx k l - GT -> go (idx + size l + 1) k r - EQ -> idx + size l -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE findIndex #-} -#endif - --- | /O(log n)/. Lookup the /index/ of a key, which is its zero-based index in --- the sequence sorted by keys. The index is a number from /0/ up to, but not --- including, the 'size' of the map. --- --- > isJust (lookupIndex 2 (fromList [(5,"a"), (3,"b")])) == False --- > fromJust (lookupIndex 3 (fromList [(5,"a"), (3,"b")])) == 0 --- > fromJust (lookupIndex 5 (fromList [(5,"a"), (3,"b")])) == 1 --- > isJust (lookupIndex 6 (fromList [(5,"a"), (3,"b")])) == False - --- See Note: Type of local 'go' function -lookupIndex :: Ord k => k -> Map k a -> Maybe Int -lookupIndex = go 0 - where - go :: Ord k => Int -> k -> Map k a -> Maybe Int - STRICT_1_OF_3(go) - STRICT_2_OF_3(go) - go _ _ Tip = Nothing - go idx k (Bin _ kx _ l r) = case compare k kx of - LT -> go idx k l - GT -> go (idx + size l + 1) k r - EQ -> Just $! idx + size l -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE lookupIndex #-} -#endif - --- | /O(log n)/. Retrieve an element by its /index/, i.e. by its zero-based --- index in the sequence sorted by keys. If the /index/ is out of range (less --- than zero, greater or equal to 'size' of the map), 'error' is called. --- --- > 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 - -elemAt :: Int -> Map k a -> (k,a) -STRICT_1_OF_2(elemAt) -elemAt _ Tip = error "Map.elemAt: index out of range" -elemAt i (Bin _ kx x l r) - = case compare i sizeL of - LT -> elemAt i l - GT -> elemAt (i-sizeL-1) r - EQ -> (kx,x) - where - sizeL = size l - --- | /O(log n)/. Update the element at /index/, i.e. by its zero-based index in --- the sequence sorted by keys. If the /index/ is out of range (less than zero, --- greater or equal to 'size' of the map), 'error' is called. --- --- > 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 - -updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a -updateAt f i t = i `seq` - case t of - Tip -> error "Map.updateAt: index out of range" - Bin sx kx x l r -> case compare i sizeL of - LT -> balanceR kx x (updateAt f i l) r - GT -> balanceL kx x l (updateAt f (i-sizeL-1) r) - EQ -> case f kx x of - Just x' -> Bin sx kx x' l r - Nothing -> glue l r - where - sizeL = size l - --- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based index in --- the sequence sorted by keys. If the /index/ is out of range (less than zero, --- greater or equal to 'size' of the map), 'error' is called. --- --- > 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 - -deleteAt :: Int -> Map k a -> Map k a -deleteAt i t = i `seq` - case t of - Tip -> error "Map.deleteAt: index out of range" - Bin _ kx x l r -> case compare i sizeL of - LT -> balanceR kx x (deleteAt i l) r - GT -> balanceL kx x l (deleteAt (i-sizeL-1) r) - EQ -> glue l r - where - sizeL = size l - - -{-------------------------------------------------------------------- - Minimal, Maximal ---------------------------------------------------------------------} --- | /O(log n)/. 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 - -findMin :: Map k a -> (k,a) -findMin (Bin _ kx x Tip _) = (kx,x) -findMin (Bin _ _ _ l _) = findMin l -findMin Tip = error "Map.findMin: empty map has no minimal element" - --- | /O(log n)/. 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 - -findMax :: Map k a -> (k,a) -findMax (Bin _ kx x _ Tip) = (kx,x) -findMax (Bin _ _ _ _ r) = findMax r -findMax Tip = error "Map.findMax: empty map has no maximal element" - --- | /O(log n)/. 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 - -deleteMin :: Map k a -> Map k a -deleteMin (Bin _ _ _ Tip r) = r -deleteMin (Bin _ kx x l r) = balanceR kx x (deleteMin l) r -deleteMin Tip = Tip - --- | /O(log n)/. 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 - -deleteMax :: Map k a -> Map k a -deleteMax (Bin _ _ _ l Tip) = l -deleteMax (Bin _ kx x l r) = balanceL kx x l (deleteMax r) -deleteMax Tip = Tip - --- | /O(log n)/. 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" - -updateMin :: (a -> Maybe a) -> Map k a -> Map k a -updateMin f m - = updateMinWithKey (\_ x -> f x) m - --- | /O(log n)/. 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" - -updateMax :: (a -> Maybe a) -> Map k a -> Map k a -updateMax f m - = updateMaxWithKey (\_ x -> f x) m - - --- | /O(log n)/. 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" - -updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a -updateMinWithKey _ Tip = Tip -updateMinWithKey f (Bin sx kx x Tip r) = case f kx x of - Nothing -> r - Just x' -> Bin sx kx x' Tip r -updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r - --- | /O(log n)/. 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" - -updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a -updateMaxWithKey _ Tip = Tip -updateMaxWithKey f (Bin sx kx x l Tip) = case f kx x of - Nothing -> l - Just x' -> Bin sx kx x' l Tip -updateMaxWithKey f (Bin _ kx x l r) = balanceL kx x l (updateMaxWithKey f r) - --- | /O(log n)/. 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 :: Map k a -> Maybe ((k,a), Map k a) -minViewWithKey Tip = Nothing -minViewWithKey x = Just (deleteFindMin x) - --- | /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 :: Map k a -> Maybe ((k,a), Map k a) -maxViewWithKey Tip = Nothing -maxViewWithKey x = Just (deleteFindMax x) - --- | /O(log n)/. 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 - -minView :: Map k a -> Maybe (a, Map k a) -minView Tip = Nothing -minView x = Just (first snd $ deleteFindMin x) - --- | /O(log n)/. Retrieves the value associated with maximal key of the --- map, and the map stripped of that element, or 'Nothing' if passed an --- empty map. --- --- > maxView (fromList [(5,"a"), (3,"b")]) == Just ("a", singleton 3 "b") --- > maxView empty == Nothing - -maxView :: Map k a -> Maybe (a, Map k a) -maxView Tip = Nothing -maxView x = Just (first snd $ deleteFindMax x) - --- Update the 1st component of a tuple (special case of Control.Arrow.first) -first :: (a -> b) -> (a,c) -> (b,c) -first f (x,y) = (f x, y) - -{-------------------------------------------------------------------- - Union. ---------------------------------------------------------------------} --- | The union of a list of maps: --- (@'unions' == 'Prelude.foldl' 'union' 'empty'@). --- --- > unions [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] --- > == fromList [(3, "b"), (5, "a"), (7, "C")] --- > unions [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])] --- > == fromList [(3, "B3"), (5, "A3"), (7, "C")] - -unions :: Ord k => [Map k a] -> Map k a -unions ts - = foldlStrict union empty ts -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE unions #-} -#endif - --- | The union of a list of maps, with a combining operation: --- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@). --- --- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] --- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] - -unionsWith :: Ord k => (a->a->a) -> [Map k a] -> Map k a -unionsWith f ts - = foldlStrict (unionWith f) empty ts -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE unionsWith #-} -#endif - --- | /O(n+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. --- --- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")] - -union :: Ord k => Map k a -> Map k a -> Map k a -union Tip t2 = t2 -union t1 Tip = t1 -union t1 t2 = hedgeUnion NothingS NothingS t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE union #-} -#endif - --- left-biased hedge union -hedgeUnion :: Ord a => MaybeS a -> MaybeS a -> Map a b -> Map a b -> Map a b -hedgeUnion _ _ t1 Tip = t1 -hedgeUnion blo bhi Tip (Bin _ kx x l r) = link kx x (filterGt blo l) (filterLt bhi r) -hedgeUnion _ _ t1 (Bin _ kx x Tip Tip) = insertR kx x t1 -- According to benchmarks, this special case increases - -- performance up to 30%. It does not help in difference or intersection. -hedgeUnion blo bhi (Bin _ kx x l r) t2 = link kx x (hedgeUnion blo bmi l (trim blo bmi t2)) - (hedgeUnion bmi bhi r (trim bmi bhi t2)) - where bmi = JustS kx -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE hedgeUnion #-} -#endif - -{-------------------------------------------------------------------- - Union with a combining function ---------------------------------------------------------------------} --- | /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")] - -unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a -unionWith f m1 m2 - = unionWithKey (\_ x y -> f x y) m1 m2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE unionWith #-} -#endif - --- | /O(n+m)/. --- Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm. --- --- > 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")] - -unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a -unionWithKey f t1 t2 = mergeWithKey (\k x1 x2 -> Just $ f k x1 x2) id id t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE unionWithKey #-} -#endif - -{-------------------------------------------------------------------- - Difference ---------------------------------------------------------------------} --- | /O(n+m)/. 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 :: Ord k => Map k a -> Map k b -> Map k a -difference Tip _ = Tip -difference t1 Tip = t1 -difference t1 t2 = hedgeDiff NothingS NothingS t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE difference #-} -#endif - -hedgeDiff :: Ord a => MaybeS a -> MaybeS a -> Map a b -> Map a c -> Map a b -hedgeDiff _ _ Tip _ = Tip -hedgeDiff blo bhi (Bin _ kx x l r) Tip = link kx x (filterGt blo l) (filterLt bhi r) -hedgeDiff blo bhi t (Bin _ kx _ l r) = merge (hedgeDiff blo bmi (trim blo bmi t) l) - (hedgeDiff bmi bhi (trim bmi bhi t) r) - where bmi = JustS kx -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE hedgeDiff #-} -#endif - --- | /O(n+m)/. 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" - -differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a -differenceWith f m1 m2 - = differenceWithKey (\_ x y -> f x y) m1 m2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE differenceWith #-} -#endif - --- | /O(n+m)/. 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" - -differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a -differenceWithKey f t1 t2 = mergeWithKey f id (const Tip) t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE differenceWithKey #-} -#endif - - -{-------------------------------------------------------------------- - Intersection ---------------------------------------------------------------------} --- | /O(n+m)/. Intersection of two maps. --- Return data in the first map for the keys existing in both maps. --- (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@). --- The implementation uses an efficient /hedge/ algorithm comparable with --- /hedge-union/. --- --- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a" - -intersection :: Ord k => Map k a -> Map k b -> Map k a -intersection Tip _ = Tip -intersection _ Tip = Tip -intersection t1 t2 = hedgeInt NothingS NothingS t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE intersection #-} -#endif - -hedgeInt :: Ord k => MaybeS k -> MaybeS k -> Map k a -> Map k b -> Map k a -hedgeInt _ _ _ Tip = Tip -hedgeInt _ _ Tip _ = Tip -hedgeInt blo bhi (Bin _ kx x l r) t2 = let l' = hedgeInt blo bmi l (trim blo bmi t2) - r' = hedgeInt bmi bhi r (trim bmi bhi t2) - in if kx `member` t2 then link kx x l' r' else merge l' r' - where bmi = JustS kx -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE hedgeInt #-} -#endif - --- | /O(n+m)/. Intersection with a combining function. The implementation uses --- an efficient /hedge/ algorithm comparable with /hedge-union/. --- --- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" - -intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c -intersectionWith f m1 m2 - = intersectionWithKey (\_ x y -> f x y) m1 m2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE intersectionWith #-} -#endif - --- | /O(n+m)/. Intersection with a combining function. The implementation uses --- an efficient /hedge/ algorithm comparable with /hedge-union/. --- --- > 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" - - -intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c -intersectionWithKey f t1 t2 = mergeWithKey (\k x1 x2 -> Just $ f k x1 x2) (const Tip) (const Tip) t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE intersectionWithKey #-} -#endif - - -{-------------------------------------------------------------------- - MergeWithKey ---------------------------------------------------------------------} - --- | /O(n+m)/. A high-performance universal combining function. This function --- is used to define 'unionWith', 'unionWithKey', 'differenceWith', --- 'differenceWithKey', 'intersectionWith', 'intersectionWithKey' and can be --- used to define other custom combine functions. --- --- Please make sure you know what is going on when using 'mergeWithKey', --- otherwise you can be surprised by unexpected code growth or even --- corruption of the data structure. --- --- When 'mergeWithKey' is given three arguments, it is inlined to the call --- site. You should therefore use 'mergeWithKey' only to define your custom --- combining functions. For example, you could define 'unionWithKey', --- 'differenceWithKey' and 'intersectionWithKey' as --- --- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2 --- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2 --- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2 --- --- When calling @'mergeWithKey' combine only1 only2@, a function combining two --- 'Map's is created, such that --- --- * if a key is present in both maps, it is passed with both corresponding --- values to the @combine@ function. Depending on the result, the key is either --- present in the result with specified value, or is left out; --- --- * a nonempty subtree present only in the first map is passed to @only1@ and --- the output is added to the result; --- --- * a nonempty subtree present only in the second map is passed to @only2@ and --- the output is added to the result. --- --- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/. --- The values can be modified arbitrarily. Most common variants of @only1@ and --- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or --- @'filterWithKey' f@ could be used for any @f@. - -mergeWithKey :: Ord k => (k -> a -> b -> Maybe c) -> (Map k a -> Map k c) -> (Map k b -> Map k c) - -> Map k a -> Map k b -> Map k c -mergeWithKey f g1 g2 = go - where - go Tip t2 = g2 t2 - go t1 Tip = g1 t1 - go t1 t2 = hedgeMerge NothingS NothingS t1 t2 - - hedgeMerge _ _ t1 Tip = g1 t1 - hedgeMerge blo bhi Tip (Bin _ kx x l r) = g2 $ link kx x (filterGt blo l) (filterLt bhi r) - hedgeMerge blo bhi (Bin _ kx x l r) t2 = let l' = hedgeMerge blo bmi l (trim blo bmi t2) - (found, trim_t2) = trimLookupLo kx bhi t2 - r' = hedgeMerge bmi bhi r trim_t2 - in case found of - Nothing -> case g1 (singleton kx x) of - Tip -> merge l' r' - (Bin _ _ x' Tip Tip) -> link kx x' l' r' - _ -> error "mergeWithKey: Given function only1 does not fulfil required conditions (see documentation)" - Just x2 -> case f kx x x2 of - Nothing -> merge l' r' - Just x' -> link kx x' l' r' - where bmi = JustS kx -{-# INLINE mergeWithKey #-} - -{-------------------------------------------------------------------- - Submap ---------------------------------------------------------------------} --- | /O(n+m)/. --- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). --- -isSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool -isSubmapOf m1 m2 = isSubmapOfBy (==) m1 m2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE isSubmapOf #-} -#endif - -{- | /O(n+m)/. - 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)]) - - --} -isSubmapOfBy :: Ord k => (a->b->Bool) -> Map k a -> Map k b -> Bool -isSubmapOfBy f t1 t2 - = (size t1 <= size t2) && (submap' f t1 t2) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE isSubmapOfBy #-} -#endif - -submap' :: Ord a => (b -> c -> Bool) -> Map a b -> Map a c -> Bool -submap' _ Tip _ = True -submap' _ _ Tip = False -submap' f (Bin _ kx x l r) t - = case found of - Nothing -> False - Just y -> f x y && submap' f l lt && submap' f r gt - where - (lt,found,gt) = splitLookup kx t -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE submap' #-} -#endif - --- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). --- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@). -isProperSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool -isProperSubmapOf m1 m2 - = isProperSubmapOfBy (==) m1 m2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE isProperSubmapOf #-} -#endif - -{- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). - The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when - @m1@ and @m2@ are not equal, - all keys in @m1@ are in @m2@, and when @f@ returns 'True' when - applied to their respective values. For example, the following - expressions are all 'True': - - > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) - > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) - - But the following are all 'False': - - > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) - > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)]) - > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) - - --} -isProperSubmapOfBy :: Ord k => (a -> b -> Bool) -> Map k a -> Map k b -> Bool -isProperSubmapOfBy f t1 t2 - = (size t1 < size t2) && (submap' f t1 t2) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE isProperSubmapOfBy #-} -#endif - -{-------------------------------------------------------------------- - Filter and partition ---------------------------------------------------------------------} --- | /O(n)/. 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 - -filter :: (a -> Bool) -> Map k a -> Map k a -filter p m - = filterWithKey (\_ x -> p x) m - --- | /O(n)/. Filter all keys\/values that satisfy the predicate. --- --- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -filterWithKey :: (k -> a -> Bool) -> Map k a -> Map k a -filterWithKey _ Tip = Tip -filterWithKey p (Bin _ kx x l r) - | p kx x = link kx x (filterWithKey p l) (filterWithKey p r) - | otherwise = merge (filterWithKey p l) (filterWithKey p r) - --- | /O(n)/. 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")]) - -partition :: (a -> Bool) -> Map k a -> (Map k a,Map k a) -partition p m - = partitionWithKey (\_ x -> p x) m - --- | /O(n)/. 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'. --- --- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b") --- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty) --- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")]) - -partitionWithKey :: (k -> a -> Bool) -> Map k a -> (Map k a,Map k a) -partitionWithKey p0 t0 = toPair $ go p0 t0 - where - go _ Tip = (Tip :*: Tip) - go p (Bin _ kx x l r) - | p kx x = link kx x l1 r1 :*: merge l2 r2 - | otherwise = merge l1 r1 :*: link kx x l2 r2 - where - (l1 :*: l2) = go p l - (r1 :*: r2) = go p r - --- | /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" - -mapMaybe :: (a -> Maybe b) -> Map k a -> Map k b -mapMaybe f = mapMaybeWithKey (\_ x -> f x) - --- | /O(n)/. 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" - -mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b -mapMaybeWithKey _ Tip = Tip -mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of - Just y -> link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r) - Nothing -> merge (mapMaybeWithKey f l) (mapMaybeWithKey f r) - --- | /O(n)/. 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")]) - -mapEither :: (a -> Either b c) -> Map k a -> (Map k b, Map k c) -mapEither f m - = mapEitherWithKey (\_ x -> f x) m - --- | /O(n)/. 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")]) - -mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) -mapEitherWithKey f0 t0 = toPair $ go f0 t0 - where - go _ Tip = (Tip :*: Tip) - go f (Bin _ kx x l r) = case f kx x of - Left y -> link kx y l1 r1 :*: merge l2 r2 - Right z -> merge l1 r1 :*: link kx z l2 r2 - where - (l1 :*: l2) = go f l - (r1 :*: r2) = go f r - -{-------------------------------------------------------------------- - Mapping ---------------------------------------------------------------------} --- | /O(n)/. Map a function over all values in the map. --- --- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] - -map :: (a -> b) -> Map k a -> Map k b -map _ Tip = Tip -map f (Bin sx kx x l r) = Bin sx kx (f x) (map f l) (map f r) -#ifdef __GLASGOW_HASKELL__ -{-# NOINLINE [1] map #-} -{-# RULES -"map/map" forall f g xs . map f (map g xs) = map (f . g) xs - #-} -#endif -#if __GLASGOW_HASKELL__ >= 709 --- Safe coercions were introduced in 7.8, but did not work well with RULES yet. -{-# RULES -"map/coerce" map coerce = coerce - #-} -#endif - --- | /O(n)/. 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")] - -mapWithKey :: (k -> a -> b) -> Map k a -> Map k b -mapWithKey _ Tip = Tip -mapWithKey f (Bin sx kx x l r) = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r) - -#ifdef __GLASGOW_HASKELL__ -{-# NOINLINE [1] mapWithKey #-} -{-# RULES -"mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) = - mapWithKey (\k a -> f k (g k a)) xs -"mapWithKey/map" forall f g xs . mapWithKey f (map g xs) = - mapWithKey (\k a -> f k (g a)) xs -"map/mapWithKey" forall f g xs . map f (mapWithKey g xs) = - mapWithKey (\k a -> f (g k a)) xs - #-} -#endif - --- | /O(n)/. --- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@ --- That is, behaves exactly like a regular 'traverse' except that the traversing --- function also has access to the key associated with a value. --- --- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')]) --- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')]) == Nothing -traverseWithKey :: Applicative t => (k -> a -> t b) -> Map k a -> t (Map k b) -traverseWithKey f = go - where - go Tip = pure Tip - go (Bin 1 k v _ _) = (\v' -> Bin 1 k v' Tip Tip) <$> f k v - go (Bin s k v l r) = flip (Bin s k) <$> go l <*> f k v <*> go r -{-# INLINE traverseWithKey #-} - --- | /O(n)/. The function 'mapAccum' threads an accumulating --- argument through the map in ascending order of keys. --- --- > let f a b = (a ++ b, b ++ "X") --- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) - -mapAccum :: (a -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) -mapAccum f a m - = mapAccumWithKey (\a' _ x' -> f a' x') a m - --- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating --- argument through the map in ascending order of keys. --- --- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") --- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) - -mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) -mapAccumWithKey f a t - = mapAccumL f a t - --- | /O(n)/. The function 'mapAccumL' threads an accumulating --- argument through the map in ascending order of keys. -mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) -mapAccumL _ a Tip = (a,Tip) -mapAccumL f a (Bin sx kx x l r) = - let (a1,l') = mapAccumL f a l - (a2,x') = f a1 kx x - (a3,r') = mapAccumL f a2 r - in (a3,Bin sx kx x' l' r') - --- | /O(n)/. The function 'mapAccumR' threads an accumulating --- argument through the map in descending order of keys. -mapAccumRWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) -mapAccumRWithKey _ a Tip = (a,Tip) -mapAccumRWithKey f a (Bin sx kx x l r) = - let (a1,r') = mapAccumRWithKey f a r - (a2,x') = f a1 kx x - (a3,l') = mapAccumRWithKey f a2 l - in (a3,Bin sx kx x' l' r') - --- | /O(n*log n)/. --- @'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 greatest of the --- original 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" - -mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a -mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) [] -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapKeys #-} -#endif - --- | /O(n*log n)/. --- @'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" - -mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a -mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapKeysWith #-} -#endif - - --- | /O(n)/. --- @'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")] --- > valid (mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")])) == True --- > valid (mapKeysMonotonic (\ _ -> 1) (fromList [(5,"a"), (3,"b")])) == False - -mapKeysMonotonic :: (k1->k2) -> Map k1 a -> Map k2 a -mapKeysMonotonic _ Tip = Tip -mapKeysMonotonic f (Bin sz k x l r) = - Bin sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r) - -{-------------------------------------------------------------------- - Folds ---------------------------------------------------------------------} - --- | /O(n)/. Fold the values in the map using the given right-associative --- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@. --- --- For example, --- --- > elems map = foldr (:) [] map --- --- > let f a len = len + (length a) --- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4 -foldr :: (a -> b -> b) -> b -> Map k a -> b -foldr f z = go z - where - go z' Tip = z' - go z' (Bin _ _ x l r) = go (f x (go z' r)) l -{-# INLINE foldr #-} - --- | /O(n)/. A strict version of 'foldr'. Each application of the operator is --- evaluated before using the result in the next application. This --- function is strict in the starting value. -foldr' :: (a -> b -> b) -> b -> Map k a -> b -foldr' f z = go z - where - STRICT_1_OF_2(go) - go z' Tip = z' - go z' (Bin _ _ x l r) = go (f x (go z' r)) l -{-# INLINE foldr' #-} - --- | /O(n)/. Fold the values in the map using the given left-associative --- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@. --- --- For example, --- --- > elems = reverse . foldl (flip (:)) [] --- --- > let f len a = len + (length a) --- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4 -foldl :: (a -> b -> a) -> a -> Map k b -> a -foldl f z = go z - where - go z' Tip = z' - go z' (Bin _ _ x l r) = go (f (go z' l) x) r -{-# INLINE foldl #-} - --- | /O(n)/. A strict version of 'foldl'. Each application of the operator is --- evaluated before using the result in the next application. This --- function is strict in the starting value. -foldl' :: (a -> b -> a) -> a -> Map k b -> a -foldl' f z = go z - where - STRICT_1_OF_2(go) - go z' Tip = z' - go z' (Bin _ _ x l r) = go (f (go z' l) x) r -{-# INLINE foldl' #-} - --- | /O(n)/. Fold the keys and values in the map using the given right-associative --- binary operator, such that --- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@. --- --- For example, --- --- > keys map = foldrWithKey (\k x ks -> k:ks) [] map --- --- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")" --- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)" -foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b -foldrWithKey f z = go z - where - go z' Tip = z' - go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l -{-# INLINE foldrWithKey #-} - --- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is --- evaluated before using the result in the next application. This --- function is strict in the starting value. -foldrWithKey' :: (k -> a -> b -> b) -> b -> Map k a -> b -foldrWithKey' f z = go z - where - STRICT_1_OF_2(go) - go z' Tip = z' - go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l -{-# INLINE foldrWithKey' #-} - --- | /O(n)/. Fold the keys and values in the map using the given left-associative --- binary operator, such that --- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@. --- --- For example, --- --- > keys = reverse . foldlWithKey (\ks k x -> k:ks) [] --- --- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")" --- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)" -foldlWithKey :: (a -> k -> b -> a) -> a -> Map k b -> a -foldlWithKey f z = go z - where - go z' Tip = z' - go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r -{-# INLINE foldlWithKey #-} - --- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is --- evaluated before using the result in the next application. This --- function is strict in the starting value. -foldlWithKey' :: (a -> k -> b -> a) -> a -> Map k b -> a -foldlWithKey' f z = go z - where - STRICT_1_OF_2(go) - go z' Tip = z' - go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r -{-# INLINE foldlWithKey' #-} - --- | /O(n)/. Fold the keys and values in the map using the given monoid, such that --- --- @'foldMapWithKey' f = 'Prelude.fold' . 'mapWithKey' f@ --- --- This can be an asymptotically faster than 'foldrWithKey' or 'foldlWithKey' for some monoids. -foldMapWithKey :: Monoid m => (k -> a -> m) -> Map k a -> m -foldMapWithKey f = go - where - go Tip = mempty - go (Bin 1 k v _ _) = f k v - go (Bin _ k v l r) = go l `mappend` (f k v `mappend` go r) -{-# INLINE foldMapWithKey #-} - -{-------------------------------------------------------------------- - List variations ---------------------------------------------------------------------} --- | /O(n)/. --- Return all elements of the map in the ascending order of their keys. --- Subject to list fusion. --- --- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"] --- > elems empty == [] - -elems :: Map k a -> [a] -elems = foldr (:) [] - --- | /O(n)/. Return all keys of the map in ascending order. Subject to list --- fusion. --- --- > keys (fromList [(5,"a"), (3,"b")]) == [3,5] --- > keys empty == [] - -keys :: Map k a -> [k] -keys = foldrWithKey (\k _ ks -> k : ks) [] - --- | /O(n)/. An alias for 'toAscList'. Return all key\/value pairs in the map --- in ascending key order. Subject to list fusion. --- --- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] --- > assocs empty == [] - -assocs :: Map k a -> [(k,a)] -assocs m - = toAscList m - --- | /O(n)/. The set of all keys of the map. --- --- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.Set.fromList [3,5] --- > keysSet empty == Data.Set.empty - -keysSet :: Map k a -> Set.Set k -keysSet Tip = Set.Tip -keysSet (Bin sz kx _ l r) = Set.Bin sz kx (keysSet l) (keysSet r) - --- | /O(n)/. Build a map from a set of keys and a function which for each key --- computes its value. --- --- > fromSet (\k -> replicate k 'a') (Data.Set.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")] --- > fromSet undefined Data.Set.empty == empty - -fromSet :: (k -> a) -> Set.Set k -> Map k a -fromSet _ Set.Tip = Tip -fromSet f (Set.Bin sz x l r) = Bin sz x (f x) (fromSet f l) (fromSet f r) - -{-------------------------------------------------------------------- - Lists - use [foldlStrict] to reduce demand on the control-stack ---------------------------------------------------------------------} -#if __GLASGOW_HASKELL__ >= 708 -instance (Ord k) => GHCExts.IsList (Map k v) where - type Item (Map k v) = (k,v) - fromList = fromList - toList = toList -#endif - --- | /O(n*log n)/. 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. --- --- If the keys of the list are ordered, linear-time implementation is used, --- with the performance equal to 'fromDistinctAscList'. --- --- > 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")] - --- For some reason, when 'singleton' is used in fromList or in --- create, it is not inlined, so we inline it manually. -fromList :: Ord k => [(k,a)] -> Map k a -fromList [] = Tip -fromList [(kx, x)] = Bin 1 kx x Tip Tip -fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip Tip) xs0 - | otherwise = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 - where - not_ordered _ [] = False - not_ordered kx ((ky,_) : _) = kx >= ky - {-# INLINE not_ordered #-} - - fromList' t0 xs = foldlStrict ins t0 xs - where ins t (k,x) = insert k x t - - STRICT_1_OF_3(go) - go _ t [] = t - go _ t [(kx, x)] = insertMax kx x t - go s l xs@((kx, x) : xss) | not_ordered kx xss = fromList' l xs - | otherwise = case create s xss of - (r, ys, []) -> go (s `shiftL` 1) (link kx x l r) ys - (r, _, ys) -> fromList' (link kx x l r) ys - - -- The create is returning a triple (tree, xs, ys). Both xs and ys - -- represent not yet processed elements and only one of them can be nonempty. - -- If ys is nonempty, the keys in ys are not ordered with respect to tree - -- and must be inserted using fromList'. Otherwise the keys have been - -- ordered so far. - STRICT_1_OF_2(create) - create _ [] = (Tip, [], []) - create s xs@(xp : xss) - | s == 1 = case xp of (kx, x) | not_ordered kx xss -> (Bin 1 kx x Tip Tip, [], xss) - | otherwise -> (Bin 1 kx x Tip Tip, xss, []) - | otherwise = case create (s `shiftR` 1) xs of - res@(_, [], _) -> res - (l, [(ky, y)], zs) -> (insertMax ky y l, [], zs) - (l, ys@((ky, y):yss), _) | not_ordered ky yss -> (l, [], ys) - | otherwise -> case create (s `shiftR` 1) yss of - (r, zs, ws) -> (link ky y l r, zs, ws) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromList #-} -#endif - --- | /O(n*log n)/. 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 - -fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a -fromListWith f xs - = fromListWithKey (\_ x y -> f x y) xs -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromListWith #-} -#endif - --- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. --- --- > let f k a1 a2 = (show k) ++ a1 ++ a2 --- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "3ab"), (5, "5a5ba")] --- > fromListWithKey f [] == empty - -fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a -fromListWithKey f xs - = foldlStrict ins empty xs - where - ins t (k,x) = insertWithKey f k x t -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromListWithKey #-} -#endif - --- | /O(n)/. Convert the map to a list of key\/value pairs. Subject to list fusion. --- --- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] --- > toList empty == [] - -toList :: Map k a -> [(k,a)] -toList = toAscList - --- | /O(n)/. Convert the map to a list of key\/value pairs where the keys are --- in ascending order. Subject to list fusion. --- --- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] - -toAscList :: Map k a -> [(k,a)] -toAscList = foldrWithKey (\k x xs -> (k,x):xs) [] - --- | /O(n)/. Convert the map to a list of key\/value pairs where the keys --- are in descending order. Subject to list fusion. --- --- > toDescList (fromList [(5,"a"), (3,"b")]) == [(5,"a"), (3,"b")] - -toDescList :: Map k a -> [(k,a)] -toDescList = foldlWithKey (\xs k x -> (k,x):xs) [] - --- List fusion for the list generating functions. -#if __GLASGOW_HASKELL__ --- The foldrFB and foldlFB are fold{r,l}WithKey equivalents, used for list fusion. --- They are important to convert unfused methods back, see mapFB in prelude. -foldrFB :: (k -> a -> b -> b) -> b -> Map k a -> b -foldrFB = foldrWithKey -{-# INLINE[0] foldrFB #-} -foldlFB :: (a -> k -> b -> a) -> a -> Map k b -> a -foldlFB = foldlWithKey -{-# INLINE[0] foldlFB #-} - --- Inline assocs and toList, so that we need to fuse only toAscList. -{-# INLINE assocs #-} -{-# INLINE toList #-} - --- The fusion is enabled up to phase 2 included. If it does not succeed, --- convert in phase 1 the expanded elems,keys,to{Asc,Desc}List calls back to --- elems,keys,to{Asc,Desc}List. In phase 0, we inline fold{lr}FB (which were --- used in a list fusion, otherwise it would go away in phase 1), and let compiler --- do whatever it wants with elems,keys,to{Asc,Desc}List -- it was forbidden to --- inline it before phase 0, otherwise the fusion rules would not fire at all. -{-# NOINLINE[0] elems #-} -{-# NOINLINE[0] keys #-} -{-# NOINLINE[0] toAscList #-} -{-# NOINLINE[0] toDescList #-} -{-# RULES "Map.elems" [~1] forall m . elems m = build (\c n -> foldrFB (\_ x xs -> c x xs) n m) #-} -{-# RULES "Map.elemsBack" [1] foldrFB (\_ x xs -> x : xs) [] = elems #-} -{-# RULES "Map.keys" [~1] forall m . keys m = build (\c n -> foldrFB (\k _ xs -> c k xs) n m) #-} -{-# RULES "Map.keysBack" [1] foldrFB (\k _ xs -> k : xs) [] = keys #-} -{-# RULES "Map.toAscList" [~1] forall m . toAscList m = build (\c n -> foldrFB (\k x xs -> c (k,x) xs) n m) #-} -{-# RULES "Map.toAscListBack" [1] foldrFB (\k x xs -> (k, x) : xs) [] = toAscList #-} -{-# RULES "Map.toDescList" [~1] forall m . toDescList m = build (\c n -> foldlFB (\xs k x -> c (k,x) xs) n m) #-} -{-# RULES "Map.toDescListBack" [1] foldlFB (\xs k x -> (k, x) : xs) [] = toDescList #-} -#endif - -{-------------------------------------------------------------------- - Building trees from ascending/descending lists can be done in linear time. - - Note that if [xs] is ascending that: - fromAscList xs == fromList xs - fromAscListWith f xs == fromListWith f xs ---------------------------------------------------------------------} --- | /O(n)/. 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")] --- > valid (fromAscList [(3,"b"), (5,"a"), (5,"b")]) == True --- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False - -fromAscList :: Eq k => [(k,a)] -> Map k a -fromAscList xs - = fromAscListWithKey (\_ x _ -> x) xs -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromAscList #-} -#endif - --- | /O(n)/. 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")] --- > valid (fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")]) == True --- > valid (fromAscListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False - -fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a -fromAscListWith f xs - = fromAscListWithKey (\_ x y -> f x y) xs -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromAscListWith #-} -#endif - --- | /O(n)/. 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./ --- --- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 --- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] == fromList [(3, "b"), (5, "5:b5:ba")] --- > valid (fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")]) == True --- > valid (fromAscListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False - -fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a -fromAscListWithKey f xs - = fromDistinctAscList (combineEq f xs) - where - -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] - combineEq _ xs' - = case xs' of - [] -> [] - [x] -> [x] - (x:xx) -> combineEq' x xx - - combineEq' z [] = [z] - combineEq' z@(kz,zz) (x@(kx,xx):xs') - | kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs' - | otherwise = z:combineEq' x xs' -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromAscListWithKey #-} -#endif - - --- | /O(n)/. 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")] --- > valid (fromDistinctAscList [(3,"b"), (5,"a")]) == True --- > valid (fromDistinctAscList [(3,"b"), (5,"a"), (5,"b")]) == False - --- For some reason, when 'singleton' is used in fromDistinctAscList or in --- create, it is not inlined, so we inline it manually. -fromDistinctAscList :: [(k,a)] -> Map k a -fromDistinctAscList [] = Tip -fromDistinctAscList ((kx0, x0) : xs0) = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 - where - STRICT_1_OF_3(go) - go _ t [] = t - go s l ((kx, x) : xs) = case create s xs of - (r, ys) -> go (s `shiftL` 1) (link kx x l r) ys - - STRICT_1_OF_2(create) - create _ [] = (Tip, []) - create s xs@(x' : xs') - | s == 1 = case x' of (kx, x) -> (Bin 1 kx x Tip Tip, xs') - | otherwise = case create (s `shiftR` 1) xs of - res@(_, []) -> res - (l, (ky, y):ys) -> case create (s `shiftR` 1) ys of - (r, zs) -> (link ky y l r, zs) - - -{-------------------------------------------------------------------- - Utility functions that return sub-ranges of the original - tree. Some functions take a `Maybe value` as an argument to - allow comparisons against infinite values. These are called `blow` - (Nothing is -\infty) and `bhigh` (here Nothing is +\infty). - We use MaybeS value, which is a Maybe strict in the Just case. - - [trim blow bhigh t] A tree that is either empty or where [x > blow] - and [x < bhigh] for the value [x] of the root. - [filterGt blow t] A tree where for all values [k]. [k > blow] - [filterLt bhigh t] A tree where for all values [k]. [k < bhigh] - - [split k t] Returns two trees [l] and [r] where all keys - in [l] are <[k] and all keys in [r] are >[k]. - [splitLookup k t] Just like [split] but also returns whether [k] - was found in the tree. ---------------------------------------------------------------------} - -data MaybeS a = NothingS | JustS !a - -{-------------------------------------------------------------------- - [trim blo bhi t] trims away all subtrees that surely contain no - values between the range [blo] to [bhi]. The returned tree is either - empty or the key of the root is between @blo@ and @bhi@. ---------------------------------------------------------------------} -trim :: Ord k => MaybeS k -> MaybeS k -> Map k a -> Map k a -trim NothingS NothingS t = t -trim (JustS lk) NothingS t = greater lk t where greater lo (Bin _ k _ _ r) | k <= lo = greater lo r - greater _ t' = t' -trim NothingS (JustS hk) t = lesser hk t where lesser hi (Bin _ k _ l _) | k >= hi = lesser hi l - lesser _ t' = t' -trim (JustS lk) (JustS hk) t = middle lk hk t where middle lo hi (Bin _ k _ _ r) | k <= lo = middle lo hi r - middle lo hi (Bin _ k _ l _) | k >= hi = middle lo hi l - middle _ _ t' = t' -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE trim #-} -#endif - --- Helper function for 'mergeWithKey'. The @'trimLookupLo' lk hk t@ performs both --- @'trim' (JustS lk) hk t@ and @'lookup' lk t@. - --- See Note: Type of local 'go' function -trimLookupLo :: Ord k => k -> MaybeS k -> Map k a -> (Maybe a, Map k a) -trimLookupLo lk0 mhk0 t0 = toPair $ go lk0 mhk0 t0 - where - go lk NothingS t = greater lk t - where greater :: Ord k => k -> Map k a -> StrictPair (Maybe a) (Map k a) - greater lo t'@(Bin _ kx x l r) = case compare lo kx of - LT -> lookup lo l :*: t' - EQ -> (Just x :*: r) - GT -> greater lo r - greater _ Tip = (Nothing :*: Tip) - go lk (JustS hk) t = middle lk hk t - where middle :: Ord k => k -> k -> Map k a -> StrictPair (Maybe a) (Map k a) - middle lo hi t'@(Bin _ kx x l r) = case compare lo kx of - LT | kx < hi -> lookup lo l :*: t' - | otherwise -> middle lo hi l - EQ -> Just x :*: lesser hi r - GT -> middle lo hi r - middle _ _ Tip = (Nothing :*: Tip) - - lesser :: Ord k => k -> Map k a -> Map k a - lesser hi (Bin _ k _ l _) | k >= hi = lesser hi l - lesser _ t' = t' -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE trimLookupLo #-} -#endif - - -{-------------------------------------------------------------------- - [filterGt b t] filter all keys >[b] from tree [t] - [filterLt b t] filter all keys <[b] from tree [t] ---------------------------------------------------------------------} -filterGt :: Ord k => MaybeS k -> Map k v -> Map k v -filterGt NothingS t = t -filterGt (JustS b) t = filter' b t - where filter' _ Tip = Tip - filter' b' (Bin _ kx x l r) = - case compare b' kx of LT -> link kx x (filter' b' l) r - EQ -> r - GT -> filter' b' r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE filterGt #-} -#endif - -filterLt :: Ord k => MaybeS k -> Map k v -> Map k v -filterLt NothingS t = t -filterLt (JustS b) t = filter' b t - where filter' _ Tip = Tip - filter' b' (Bin _ kx x l r) = - case compare kx b' of LT -> link kx x l (filter' b' r) - EQ -> l - GT -> filter' b' l -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE filterLt #-} -#endif - -{-------------------------------------------------------------------- - Split ---------------------------------------------------------------------} --- | /O(log n)/. 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 :: Ord k => k -> Map k a -> (Map k a,Map k a) -split k0 t0 = k0 `seq` toPair $ go k0 t0 - where - go k t = - case t of - Tip -> (Tip :*: Tip) - Bin _ kx x l r -> case compare k kx of - LT -> let (lt :*: gt) = go k l in lt :*: link kx x gt r - GT -> let (lt :*: gt) = go k r in link kx x l lt :*: gt - EQ -> (l :*: r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE split #-} -#endif - --- | /O(log n)/. 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) - -splitLookup :: Ord k => k -> Map k a -> (Map k a,Maybe a,Map k a) -splitLookup k t = k `seq` - case t of - Tip -> (Tip,Nothing,Tip) - Bin _ kx x l r -> case compare k kx of - LT -> let (lt,z,gt) = splitLookup k l - gt' = link kx x gt r - in gt' `seq` (lt,z,gt') - GT -> let (lt,z,gt) = splitLookup k r - lt' = link kx x l lt - in lt' `seq` (lt',z,gt) - EQ -> (l,Just x,r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE splitLookup #-} -#endif - -{-------------------------------------------------------------------- - Utility functions that maintain the balance properties of the tree. - All constructors assume that all values in [l] < [k] and all values - in [r] > [k], and that [l] and [r] are valid trees. - - In order of sophistication: - [Bin sz k x l r] The type constructor. - [bin k x l r] Maintains the correct size, assumes that both [l] - and [r] are balanced with respect to each other. - [balance k x l r] Restores the balance and size. - Assumes that the original tree was balanced and - that [l] or [r] has changed by at most one element. - [link k x l r] Restores balance and size. - - Furthermore, we can construct a new tree from two trees. Both operations - assume that all values in [l] < all values in [r] and that [l] and [r] - are valid: - [glue l r] Glues [l] and [r] together. Assumes that [l] and - [r] are already balanced with respect to each other. - [merge l r] Merges two trees and restores balance. - - Note: in contrast to Adam's paper, we use (<=) comparisons instead - of (<) comparisons in [link], [merge] and [balance]. - Quickcheck (on [difference]) showed that this was necessary in order - to maintain the invariants. It is quite unsatisfactory that I haven't - been able to find out why this is actually the case! Fortunately, it - doesn't hurt to be a bit more conservative. ---------------------------------------------------------------------} - -{-------------------------------------------------------------------- - Link ---------------------------------------------------------------------} -link :: k -> a -> Map k a -> Map k a -> Map k a -link kx x Tip r = insertMin kx x r -link kx x l Tip = insertMax kx x l -link kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz) - | delta*sizeL < sizeR = balanceL kz z (link kx x l lz) rz - | delta*sizeR < sizeL = balanceR ky y ly (link kx x ry r) - | otherwise = bin kx x l r - - --- insertMin and insertMax don't perform potentially expensive comparisons. -insertMax,insertMin :: k -> a -> Map k a -> Map k a -insertMax kx x t - = case t of - Tip -> singleton kx x - Bin _ ky y l r - -> balanceR ky y l (insertMax kx x r) - -insertMin kx x t - = case t of - Tip -> singleton kx x - Bin _ ky y l r - -> balanceL ky y (insertMin kx x l) r - -{-------------------------------------------------------------------- - [merge l r]: merges two trees. ---------------------------------------------------------------------} -merge :: Map k a -> Map k a -> Map k a -merge Tip r = r -merge l Tip = l -merge l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) - | delta*sizeL < sizeR = balanceL ky y (merge l ly) ry - | delta*sizeR < sizeL = balanceR 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 :: Map k a -> Map k a -> Map k a -glue Tip r = r -glue l Tip = l -glue l r - | size l > size r = let ((km,m),l') = deleteFindMax l in balanceR km m l' r - | otherwise = let ((km,m),r') = deleteFindMin r in balanceL km m l r' - - --- | /O(log n)/. 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 - -deleteFindMin :: Map k a -> ((k,a),Map k a) -deleteFindMin t - = case t of - Bin _ k x Tip r -> ((k,x),r) - Bin _ k x l r -> let (km,l') = deleteFindMin l in (km,balanceR k x l' r) - Tip -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip) - --- | /O(log n)/. Delete and find the maximal element. --- --- > deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((10,"c"), fromList [(3,"b"), (5,"a")]) --- > deleteFindMax empty Error: can not return the maximal element of an empty map - -deleteFindMax :: Map k a -> ((k,a),Map k a) -deleteFindMax t - = case t of - Bin _ k x l Tip -> ((k,x),l) - Bin _ k x l r -> let (km,r') = deleteFindMax r in (km,balanceL k x l r') - Tip -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip) - - -{-------------------------------------------------------------------- - [balance l x r] balances two trees with value x. - The sizes of the trees should balance after decreasing the - size of one of them. (a rotation). - - [delta] is the maximal relative difference between the sizes of - two trees, it corresponds with the [w] in Adams' paper. - [ratio] is the ratio between an outer and inner sibling of the - heavier subtree in an unbalanced setting. It determines - whether a double or single rotation should be performed - to restore balance. It is corresponds with the inverse - of $\alpha$ in Adam's article. - - Note that according to the Adam's paper: - - [delta] should be larger than 4.646 with a [ratio] of 2. - - [delta] should be larger than 3.745 with a [ratio] of 1.534. - - But the Adam's paper is erroneous: - - It can be proved that for delta=2 and delta>=5 there does - not exist any ratio that would work. - - Delta=4.5 and ratio=2 does not work. - - That leaves two reasonable variants, delta=3 and delta=4, - both with ratio=2. - - - A lower [delta] leads to a more 'perfectly' balanced tree. - - A higher [delta] performs less rebalancing. - - In the benchmarks, delta=3 is faster on insert operations, - and delta=4 has slightly better deletes. As the insert speedup - is larger, we currently use delta=3. - ---------------------------------------------------------------------} -delta,ratio :: Int -delta = 3 -ratio = 2 - --- The balance function is equivalent to the following: --- --- balance :: k -> a -> Map k a -> Map k a -> Map k a --- balance k x l r --- | sizeL + sizeR <= 1 = Bin sizeX k x l r --- | sizeR > delta*sizeL = rotateL k x l r --- | sizeL > delta*sizeR = rotateR k x l r --- | otherwise = Bin sizeX k x l r --- where --- sizeL = size l --- sizeR = size r --- sizeX = sizeL + sizeR + 1 --- --- rotateL :: a -> b -> Map a b -> Map a b -> Map a b --- rotateL k x l r@(Bin _ _ _ ly ry) | size ly < ratio*size ry = singleL k x l r --- | otherwise = doubleL k x l r --- --- rotateR :: a -> b -> Map a b -> Map a b -> Map a b --- rotateR k x l@(Bin _ _ _ ly ry) r | size ry < ratio*size ly = singleR k x l r --- | otherwise = doubleR k x l r --- --- singleL, singleR :: a -> b -> Map a b -> Map a b -> Map a b --- singleL k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin k2 x2 (bin k1 x1 t1 t2) t3 --- singleR k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3) --- --- doubleL, doubleR :: a -> b -> Map a b -> Map a b -> Map a b --- 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) --- 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) --- --- It is only written in such a way that every node is pattern-matched only once. - -balance :: k -> a -> Map k a -> Map k a -> Map k a -balance k x l r = case l of - Tip -> case r of - Tip -> Bin 1 k x Tip Tip - (Bin _ _ _ Tip Tip) -> Bin 2 k x Tip r - (Bin _ rk rx Tip rr@(Bin _ _ _ _ _)) -> Bin 3 rk rx (Bin 1 k x Tip Tip) rr - (Bin _ rk rx (Bin _ rlk rlx _ _) Tip) -> Bin 3 rlk rlx (Bin 1 k x Tip Tip) (Bin 1 rk rx Tip Tip) - (Bin rs rk rx rl@(Bin rls rlk rlx rll rlr) rr@(Bin rrs _ _ _ _)) - | rls < ratio*rrs -> Bin (1+rs) rk rx (Bin (1+rls) k x Tip rl) rr - | otherwise -> Bin (1+rs) rlk rlx (Bin (1+size rll) k x Tip rll) (Bin (1+rrs+size rlr) rk rx rlr rr) - - (Bin ls lk lx ll lr) -> case r of - Tip -> case (ll, lr) of - (Tip, Tip) -> Bin 2 k x l Tip - (Tip, (Bin _ lrk lrx _ _)) -> Bin 3 lrk lrx (Bin 1 lk lx Tip Tip) (Bin 1 k x Tip Tip) - ((Bin _ _ _ _ _), Tip) -> Bin 3 lk lx ll (Bin 1 k x Tip Tip) - ((Bin lls _ _ _ _), (Bin lrs lrk lrx lrl lrr)) - | lrs < ratio*lls -> Bin (1+ls) lk lx ll (Bin (1+lrs) k x lr Tip) - | otherwise -> Bin (1+ls) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+size lrr) k x lrr Tip) - (Bin rs rk rx rl rr) - | rs > delta*ls -> case (rl, rr) of - (Bin rls rlk rlx rll rlr, Bin rrs _ _ _ _) - | rls < ratio*rrs -> Bin (1+ls+rs) rk rx (Bin (1+ls+rls) k x l rl) rr - | otherwise -> Bin (1+ls+rs) rlk rlx (Bin (1+ls+size rll) k x l rll) (Bin (1+rrs+size rlr) rk rx rlr rr) - (_, _) -> error "Failure in Data.Map.balance" - | ls > delta*rs -> case (ll, lr) of - (Bin lls _ _ _ _, Bin lrs lrk lrx lrl lrr) - | lrs < ratio*lls -> Bin (1+ls+rs) lk lx ll (Bin (1+rs+lrs) k x lr r) - | otherwise -> Bin (1+ls+rs) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+rs+size lrr) k x lrr r) - (_, _) -> error "Failure in Data.Map.balance" - | otherwise -> Bin (1+ls+rs) k x l r -{-# NOINLINE balance #-} - --- Functions balanceL and balanceR are specialised versions of balance. --- balanceL only checks whether the left subtree is too big, --- balanceR only checks whether the right subtree is too big. - --- balanceL is called when left subtree might have been inserted to or when --- right subtree might have been deleted from. -balanceL :: k -> a -> Map k a -> Map k a -> Map k a -balanceL k x l r = case r of - Tip -> case l of - Tip -> Bin 1 k x Tip Tip - (Bin _ _ _ Tip Tip) -> Bin 2 k x l Tip - (Bin _ lk lx Tip (Bin _ lrk lrx _ _)) -> Bin 3 lrk lrx (Bin 1 lk lx Tip Tip) (Bin 1 k x Tip Tip) - (Bin _ lk lx ll@(Bin _ _ _ _ _) Tip) -> Bin 3 lk lx ll (Bin 1 k x Tip Tip) - (Bin ls lk lx ll@(Bin lls _ _ _ _) lr@(Bin lrs lrk lrx lrl lrr)) - | lrs < ratio*lls -> Bin (1+ls) lk lx ll (Bin (1+lrs) k x lr Tip) - | otherwise -> Bin (1+ls) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+size lrr) k x lrr Tip) - - (Bin rs _ _ _ _) -> case l of - Tip -> Bin (1+rs) k x Tip r - - (Bin ls lk lx ll lr) - | ls > delta*rs -> case (ll, lr) of - (Bin lls _ _ _ _, Bin lrs lrk lrx lrl lrr) - | lrs < ratio*lls -> Bin (1+ls+rs) lk lx ll (Bin (1+rs+lrs) k x lr r) - | otherwise -> Bin (1+ls+rs) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+rs+size lrr) k x lrr r) - (_, _) -> error "Failure in Data.Map.balanceL" - | otherwise -> Bin (1+ls+rs) k x l r -{-# NOINLINE balanceL #-} - --- balanceR is called when right subtree might have been inserted to or when --- left subtree might have been deleted from. -balanceR :: k -> a -> Map k a -> Map k a -> Map k a -balanceR k x l r = case l of - Tip -> case r of - Tip -> Bin 1 k x Tip Tip - (Bin _ _ _ Tip Tip) -> Bin 2 k x Tip r - (Bin _ rk rx Tip rr@(Bin _ _ _ _ _)) -> Bin 3 rk rx (Bin 1 k x Tip Tip) rr - (Bin _ rk rx (Bin _ rlk rlx _ _) Tip) -> Bin 3 rlk rlx (Bin 1 k x Tip Tip) (Bin 1 rk rx Tip Tip) - (Bin rs rk rx rl@(Bin rls rlk rlx rll rlr) rr@(Bin rrs _ _ _ _)) - | rls < ratio*rrs -> Bin (1+rs) rk rx (Bin (1+rls) k x Tip rl) rr - | otherwise -> Bin (1+rs) rlk rlx (Bin (1+size rll) k x Tip rll) (Bin (1+rrs+size rlr) rk rx rlr rr) - - (Bin ls _ _ _ _) -> case r of - Tip -> Bin (1+ls) k x l Tip - - (Bin rs rk rx rl rr) - | rs > delta*ls -> case (rl, rr) of - (Bin rls rlk rlx rll rlr, Bin rrs _ _ _ _) - | rls < ratio*rrs -> Bin (1+ls+rs) rk rx (Bin (1+ls+rls) k x l rl) rr - | otherwise -> Bin (1+ls+rs) rlk rlx (Bin (1+ls+size rll) k x l rll) (Bin (1+rrs+size rlr) rk rx rlr rr) - (_, _) -> error "Failure in Data.Map.balanceR" - | otherwise -> Bin (1+ls+rs) k x l r -{-# NOINLINE balanceR #-} - - -{-------------------------------------------------------------------- - The bin constructor maintains the size of the tree ---------------------------------------------------------------------} -bin :: k -> a -> Map k a -> Map k a -> Map k a -bin k x l r - = Bin (size l + size r + 1) k x l r -{-# INLINE bin #-} - - -{-------------------------------------------------------------------- - Eq converts the tree to a list. In a lazy setting, this - actually seems one of the faster methods to compare two trees - and it is certainly the simplest :-) ---------------------------------------------------------------------} -instance (Eq k,Eq a) => Eq (Map k a) where - t1 == t2 = (size t1 == size t2) && (toAscList t1 == toAscList t2) - -{-------------------------------------------------------------------- - Ord ---------------------------------------------------------------------} - -instance (Ord k, Ord v) => Ord (Map k v) where - compare m1 m2 = compare (toAscList m1) (toAscList m2) - -{-------------------------------------------------------------------- - Functor ---------------------------------------------------------------------} -instance Functor (Map k) where - fmap f m = map f m - -instance Traversable (Map k) where - traverse f = traverseWithKey (\_ -> f) - {-# INLINE traverse #-} - -instance Foldable.Foldable (Map k) where - fold = go - where go Tip = mempty - go (Bin 1 _ v _ _) = v - go (Bin _ _ v l r) = go l `mappend` (v `mappend` go r) - {-# INLINABLE fold #-} - foldr = foldr - {-# INLINE foldr #-} - foldl = foldl - {-# INLINE foldl #-} - foldMap f t = go t - where go Tip = mempty - go (Bin 1 _ v _ _) = f v - go (Bin _ _ v l r) = go l `mappend` (f v `mappend` go r) - {-# INLINE foldMap #-} - -#if MIN_VERSION_base(4,6,0) - foldl' = foldl' - {-# INLINE foldl' #-} - foldr' = foldr' - {-# INLINE foldr' #-} -#endif -#if MIN_VERSION_base(4,8,0) - length = size - {-# INLINE length #-} - null = null - {-# INLINE null #-} - toList = elems -- NB: Foldable.toList /= Map.toList - {-# INLINE toList #-} - elem = go - where STRICT_1_OF_2(go) - go _ Tip = False - go x (Bin _ _ v l r) = x == v || go x l || go x r - {-# INLINABLE elem #-} - maximum = start - where start Tip = error "Map.Foldable.maximum: called with empty map" - start (Bin _ _ v l r) = go (go v l) r - - STRICT_1_OF_2(go) - go m Tip = m - go m (Bin _ _ v l r) = go (go (max m v) l) r - {-# INLINABLE maximum #-} - minimum = start - where start Tip = error "Map.Foldable.minumum: called with empty map" - start (Bin _ _ v l r) = go (go v l) r - - STRICT_1_OF_2(go) - go m Tip = m - go m (Bin _ _ v l r) = go (go (min m v) l) r - {-# INLINABLE minimum #-} - sum = foldl' (+) 0 - {-# INLINABLE sum #-} - product = foldl' (*) 1 - {-# INLINABLE product #-} -#endif - -instance (NFData k, NFData a) => NFData (Map k a) where - rnf Tip = () - rnf (Bin _ kx x l r) = rnf kx `seq` rnf x `seq` rnf l `seq` rnf r - -{-------------------------------------------------------------------- - Read ---------------------------------------------------------------------} -instance (Ord k, Read k, Read e) => Read (Map k e) where -#ifdef __GLASGOW_HASKELL__ - readPrec = parens $ prec 10 $ do - Ident "fromList" <- lexP - xs <- readPrec - return (fromList xs) - - readListPrec = readListPrecDefault -#else - readsPrec p = readParen (p > 10) $ \ r -> do - ("fromList",s) <- lex r - (xs,t) <- reads s - return (fromList xs,t) -#endif - -{-------------------------------------------------------------------- - Show ---------------------------------------------------------------------} -instance (Show k, Show a) => Show (Map k a) where - showsPrec d m = showParen (d > 10) $ - showString "fromList " . shows (toList m) - --- | /O(n)/. Show the tree that implements the map. The tree is shown --- in a compressed, hanging format. See 'showTreeWith'. -showTree :: (Show k,Show a) => Map k a -> String -showTree m - = showTreeWith showElem True False m - where - showElem k x = show k ++ ":=" ++ show x - - -{- | /O(n)/. The expression (@'showTreeWith' showelem hang wide map@) shows - the tree that implements the map. Elements are shown using the @showElem@ function. If @hang@ is - 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If - @wide@ is 'True', an extra wide version is shown. - -> Map> let t = fromDistinctAscList [(x,()) | x <- [1..5]] -> Map> putStrLn $ showTreeWith (\k x -> show (k,x)) True False t -> (4,()) -> +--(2,()) -> | +--(1,()) -> | +--(3,()) -> +--(5,()) -> -> Map> putStrLn $ showTreeWith (\k x -> show (k,x)) True True t -> (4,()) -> | -> +--(2,()) -> | | -> | +--(1,()) -> | | -> | +--(3,()) -> | -> +--(5,()) -> -> Map> putStrLn $ showTreeWith (\k x -> show (k,x)) False True t -> +--(5,()) -> | -> (4,()) -> | -> | +--(3,()) -> | | -> +--(2,()) -> | -> +--(1,()) - --} -showTreeWith :: (k -> a -> String) -> Bool -> Bool -> Map k a -> String -showTreeWith showelem hang wide t - | hang = (showsTreeHang showelem wide [] t) "" - | otherwise = (showsTree showelem wide [] [] t) "" - -showsTree :: (k -> a -> String) -> Bool -> [String] -> [String] -> Map k a -> ShowS -showsTree showelem wide lbars rbars t - = case t of - Tip -> showsBars lbars . showString "|\n" - Bin _ kx x Tip Tip - -> showsBars lbars . showString (showelem kx x) . showString "\n" - Bin _ kx x l r - -> showsTree showelem wide (withBar rbars) (withEmpty rbars) r . - showWide wide rbars . - showsBars lbars . showString (showelem kx x) . showString "\n" . - showWide wide lbars . - showsTree showelem wide (withEmpty lbars) (withBar lbars) l - -showsTreeHang :: (k -> a -> String) -> Bool -> [String] -> Map k a -> ShowS -showsTreeHang showelem wide bars t - = case t of - Tip -> showsBars bars . showString "|\n" - Bin _ kx x Tip Tip - -> showsBars bars . showString (showelem kx x) . showString "\n" - Bin _ kx x l r - -> showsBars bars . showString (showelem kx x) . showString "\n" . - showWide wide bars . - showsTreeHang showelem wide (withBar bars) l . - showWide wide bars . - showsTreeHang showelem wide (withEmpty bars) r - -showWide :: Bool -> [String] -> String -> String -showWide wide bars - | wide = showString (concat (reverse bars)) . showString "|\n" - | otherwise = id - -showsBars :: [String] -> ShowS -showsBars bars - = case bars of - [] -> id - _ -> showString (concat (reverse (tail bars))) . showString node - -node :: String -node = "+--" - -withBar, withEmpty :: [String] -> [String] -withBar bars = "| ":bars -withEmpty bars = " ":bars - -{-------------------------------------------------------------------- - Typeable ---------------------------------------------------------------------} - -INSTANCE_TYPEABLE2(Map,mapTc,"Map") +{-# LANGUAGE BangPatterns #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +#if __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE RoleAnnotations #-} +{-# LANGUAGE TypeFamilies #-} +#define USE_MAGIC_PROXY 1 +#endif + +#if USE_MAGIC_PROXY +{-# LANGUAGE MagicHash #-} +#endif + +#include "containers.h" + +#if !(WORD_SIZE_IN_BITS >= 61) +#define DEFINE_ALTERF_FALLBACK 1 +#endif + +{-# OPTIONS_HADDOCK hide #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Map.Base +-- Copyright : (c) Daan Leijen 2002 +-- (c) Andriy Palamarchuk 2008 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. +-- +-- = Description +-- +-- An efficient implementation of maps from keys to values (dictionaries). +-- +-- Since many function names (but not the type name) clash with +-- "Prelude" names, this module is usually imported @qualified@, e.g. +-- +-- > import Data.Map (Map) +-- > import qualified Data.Map as Map +-- +-- The implementation of 'Map' is based on /size balanced/ binary trees (or +-- trees of /bounded balance/) as described by: +-- +-- * Stephen Adams, \"/Efficient sets: a balancing act/\", +-- Journal of Functional Programming 3(4):553-562, October 1993, +-- . +-- * J. Nievergelt and E.M. Reingold, +-- \"/Binary search trees of bounded balance/\", +-- SIAM journal of computing 2(1), March 1973. +-- +-- Bounds for 'union', 'intersection', and 'difference' are as given +-- by +-- +-- * Guy Blelloch, Daniel Ferizovic, and Yihan Sun, +-- \"/Just Join for Parallel Ordered Sets/\", +-- . +-- +-- Note that the implementation is /left-biased/ -- the elements of a +-- first argument are always preferred to the second, for example in +-- 'union' or 'insert'. +-- +-- Operation comments contain the operation time complexity in +-- the Big-O notation . +----------------------------------------------------------------------------- + +-- [Note: Using INLINABLE] +-- ~~~~~~~~~~~~~~~~~~~~~~~ +-- It is crucial to the performance that the functions specialize on the Ord +-- type when possible. GHC 7.0 and higher does this by itself when it sees th +-- unfolding of a function -- that is why all public functions are marked +-- INLINABLE (that exposes the unfolding). + + +-- [Note: Using INLINE] +-- ~~~~~~~~~~~~~~~~~~~~ +-- For other compilers and GHC pre 7.0, we mark some of the functions INLINE. +-- We mark the functions that just navigate down the tree (lookup, insert, +-- delete and similar). That navigation code gets inlined and thus specialized +-- when possible. There is a price to pay -- code growth. The code INLINED is +-- therefore only the tree navigation, all the real work (rebalancing) is not +-- INLINED by using a NOINLINE. +-- +-- All methods marked INLINE have to be nonrecursive -- a 'go' function doing +-- the real work is provided. + + +-- [Note: Type of local 'go' function] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- If the local 'go' function uses an Ord class, it sometimes heap-allocates +-- the Ord dictionary when the 'go' function does not have explicit type. +-- In that case we give 'go' explicit type. But this slightly decrease +-- performance, as the resulting 'go' function can float out to top level. + + +-- [Note: Local 'go' functions and capturing] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- As opposed to Map, when 'go' function captures an argument, increased +-- heap-allocation can occur: sometimes in a polymorphic function, the 'go' +-- floats out of its enclosing function and then it heap-allocates the +-- dictionary and the argument. Maybe it floats out too late and strictness +-- analyzer cannot see that these could be passed on stack. +-- + +-- [Note: Order of constructors] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- The order of constructors of Map matters when considering performance. +-- Currently in GHC 7.0, when type has 2 constructors, a forward conditional +-- jump is made when successfully matching second constructor. Successful match +-- of first constructor results in the forward jump not taken. +-- On GHC 7.0, reordering constructors from Tip | Bin to Bin | Tip +-- improves the benchmark by up to 10% on x86. + +module Data.Map.Base ( + -- * Map type + Map(..) -- instance Eq,Show,Read + + -- * Operators + , (!), (\\) + + -- * Query + , null + , size + , member + , notMember + , lookup + , findWithDefault + , lookupLT + , lookupGT + , lookupLE + , lookupGE + + -- * Construction + , empty + , singleton + + -- ** Insertion + , insert + , insertWith + , insertWithKey + , insertLookupWithKey + + -- ** Delete\/Update + , delete + , adjust + , adjustWithKey + , update + , updateWithKey + , updateLookupWithKey + , alter + , alterF + + -- * Combine + + -- ** Union + , union + , unionWith + , unionWithKey + , unions + , unionsWith + + -- ** Difference + , difference + , differenceWith + , differenceWithKey + + -- ** Intersection + , intersection + , intersectionWith + , intersectionWithKey + + -- ** General combining function + , SimpleWhenMissing + , SimpleWhenMatched + , runWhenMatched + , runWhenMissing + , merge + -- *** @WhenMatched@ tactics + , zipWithMaybeMatched + , zipWithMatched + -- *** @WhenMissing@ tactics + , mapMaybeMissing + , dropMissing + , preserveMissing + , mapMissing + , filterMissing + + -- ** Applicative general combining function + , WhenMissing (..) + , WhenMatched (..) + , mergeA + + -- *** @WhenMatched@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , zipWithMaybeAMatched + , zipWithAMatched + + -- *** @WhenMissing@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , traverseMaybeMissing + , traverseMissing + , filterAMissing + + -- ** Deprecated general combining function + + , mergeWithKey + + -- * Traversal + -- ** Map + , map + , mapWithKey + , traverseWithKey + , traverseMaybeWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeys + , mapKeysWith + , mapKeysMonotonic + + -- * Folds + , foldr + , foldl + , foldrWithKey + , foldlWithKey + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + , fromSet + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + , fromDescList + , fromDescListWith + , fromDescListWithKey + , fromDistinctDescList + + -- * Filter + , filter + , filterWithKey + + , takeWhileAntitone + , dropWhileAntitone + , spanAntitone + + , restrictKeys + , withoutKeys + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , updateAt + , deleteAt + , take + , drop + , splitAt + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + , valid + + -- Used by the strict version + , AreWeStrict (..) + , atKeyImpl +#if __GLASGOW_HASKELL__ && MIN_VERSION_base(4,8,0) + , atKeyPlain +#endif + , bin + , balance + , balanced + , balanceL + , balanceR + , delta + , insertMax + , link + , link2 + , glue + , MaybeS(..) + , Identity(..) + + -- Used by Map.Lazy.Merge + , mapWhenMissing + , mapWhenMatched + , lmapWhenMissing + , contramapFirstWhenMatched + , contramapSecondWhenMatched + , mapGentlyWhenMissing + , mapGentlyWhenMatched + ) where + +#if MIN_VERSION_base(4,8,0) +import Data.Functor.Identity (Identity (..)) +#else +import Control.Applicative (Applicative(..), (<$>)) +import Data.Monoid (Monoid(..)) +import Data.Traversable (Traversable(traverse)) +#endif +#if MIN_VERSION_base(4,9,0) +import Data.Semigroup (Semigroup((<>), stimes), stimesIdempotentMonoid) +#endif +import Control.Applicative (Const (..)) +import Control.DeepSeq (NFData(rnf)) +import Data.Bits (shiftL, shiftR) +import qualified Data.Foldable as Foldable +import Data.Typeable +import Prelude hiding (lookup, map, filter, foldr, foldl, null, splitAt, take, drop) + +import qualified Data.Set.Base as Set +import Data.Set.Base (Set) +import Data.Utils.PtrEquality (ptrEq) +import Data.Utils.StrictFold +import Data.Utils.StrictPair +import Data.Utils.StrictMaybe +import Data.Utils.BitQueue +#if DEFINE_ALTERF_FALLBACK +import Data.Utils.BitUtil (wordSize) +#endif + +#if __GLASGOW_HASKELL__ +import GHC.Exts (build) +#if !MIN_VERSION_base(4,8,0) +import Data.Functor ((<$)) +#endif +#if USE_MAGIC_PROXY +import GHC.Exts (Proxy#, proxy# ) +#endif +#if __GLASGOW_HASKELL__ >= 708 +import qualified GHC.Exts as GHCExts +#endif +import Text.Read hiding (lift) +import Data.Data +import qualified Control.Category as Category +#endif +#if __GLASGOW_HASKELL__ >= 708 +import Data.Coerce +#endif + + +{-------------------------------------------------------------------- + Operators +--------------------------------------------------------------------} +infixl 9 !,\\ -- + +-- | /O(log n)/. Find the value at a key. +-- Calls 'error' when the element can not be found. +-- +-- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map +-- > fromList [(5,'a'), (3,'b')] ! 5 == 'a' + +(!) :: Ord k => Map k a -> k -> a +(!) m k = find k m +#if __GLASGOW_HASKELL__ +{-# INLINABLE (!) #-} +#endif + +-- | Same as 'difference'. +(\\) :: Ord k => Map k a -> Map k b -> Map k a +m1 \\ m2 = difference m1 m2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE (\\) #-} +#endif + +{-------------------------------------------------------------------- + Size balanced trees. +--------------------------------------------------------------------} +-- | A Map from keys @k@ to values @a@. + +-- See Note: Order of constructors +data Map k a = Bin {-# UNPACK #-} !Size !k a !(Map k a) !(Map k a) + | Tip + +type Size = Int + +#if __GLASGOW_HASKELL__ >= 708 +type role Map nominal representational +#endif + +instance (Ord k) => Monoid (Map k v) where + mempty = empty + mconcat = unions +#if !(MIN_VERSION_base(4,9,0)) + mappend = union +#else + mappend = (<>) + +instance (Ord k) => Semigroup (Map k v) where + (<>) = union + stimes = stimesIdempotentMonoid +#endif + +#if __GLASGOW_HASKELL__ + +{-------------------------------------------------------------------- + A Data instance +--------------------------------------------------------------------} + +-- This instance preserves data abstraction at the cost of inefficiency. +-- We provide limited reflection services for the sake of data abstraction. + +instance (Data k, Data a, Ord k) => Data (Map k a) where + gfoldl f z m = z fromList `f` toList m + toConstr _ = fromListConstr + gunfold k z c = case constrIndex c of + 1 -> k (z fromList) + _ -> error "gunfold" + dataTypeOf _ = mapDataType + dataCast2 f = gcast2 f + +fromListConstr :: Constr +fromListConstr = mkConstr mapDataType "fromList" [] Prefix + +mapDataType :: DataType +mapDataType = mkDataType "Data.Map.Base.Map" [fromListConstr] + +#endif + +{-------------------------------------------------------------------- + Query +--------------------------------------------------------------------} +-- | /O(1)/. Is the map empty? +-- +-- > Data.Map.null (empty) == True +-- > Data.Map.null (singleton 1 'a') == False + +null :: Map k a -> Bool +null Tip = True +null (Bin {}) = False +{-# INLINE null #-} + +-- | /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 :: Map k a -> Int +size Tip = 0 +size (Bin sz _ _ _ _) = sz +{-# INLINE size #-} + + +-- | /O(log n)/. Lookup the value at a key in the map. +-- +-- The function will return the corresponding value as @('Just' value)@, +-- or 'Nothing' if the key isn't in the map. +-- +-- An example of using @lookup@: +-- +-- > import Prelude hiding (lookup) +-- > import Data.Map +-- > +-- > employeeDept = fromList([("John","Sales"), ("Bob","IT")]) +-- > deptCountry = fromList([("IT","USA"), ("Sales","France")]) +-- > countryCurrency = fromList([("USA", "Dollar"), ("France", "Euro")]) +-- > +-- > employeeCurrency :: String -> Maybe String +-- > employeeCurrency name = do +-- > dept <- lookup name employeeDept +-- > country <- lookup dept deptCountry +-- > lookup country countryCurrency +-- > +-- > main = do +-- > putStrLn $ "John's currency: " ++ (show (employeeCurrency "John")) +-- > putStrLn $ "Pete's currency: " ++ (show (employeeCurrency "Pete")) +-- +-- The output of this program: +-- +-- > John's currency: Just "Euro" +-- > Pete's currency: Nothing +lookup :: Ord k => k -> Map k a -> Maybe a +lookup = go + where + go !_ Tip = Nothing + go k (Bin _ kx x l r) = case compare k kx of + LT -> go k l + GT -> go k r + EQ -> Just x +#if __GLASGOW_HASKELL__ +{-# INLINABLE lookup #-} +#else +{-# INLINE lookup #-} +#endif + +-- | /O(log n)/. 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 +member :: Ord k => k -> Map k a -> Bool +member = go + where + go !_ Tip = False + go k (Bin _ kx _ l r) = case compare k kx of + LT -> go k l + GT -> go k r + EQ -> True +#if __GLASGOW_HASKELL__ +{-# INLINABLE member #-} +#else +{-# INLINE member #-} +#endif + +-- | /O(log n)/. 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 + +notMember :: Ord k => k -> Map k a -> Bool +notMember k m = not $ member k m +#if __GLASGOW_HASKELL__ +{-# INLINABLE notMember #-} +#else +{-# INLINE notMember #-} +#endif + +-- | /O(log n)/. Find the value at a key. +-- Calls 'error' when the element can not be found. +find :: Ord k => k -> Map k a -> a +find = go + where + go !_ Tip = error "Map.!: given key is not an element in the map" + go k (Bin _ kx x l r) = case compare k kx of + LT -> go k l + GT -> go k r + EQ -> x +#if __GLASGOW_HASKELL__ +{-# INLINABLE find #-} +#else +{-# INLINE find #-} +#endif + +-- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns +-- the value at key @k@ or returns default value @def@ +-- when the key is not in the map. +-- +-- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' +-- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' +findWithDefault :: Ord k => a -> k -> Map k a -> a +findWithDefault = go + where + go def !_ Tip = def + go def k (Bin _ kx x l r) = case compare k kx of + LT -> go def k l + GT -> go def k r + EQ -> x +#if __GLASGOW_HASKELL__ +{-# INLINABLE findWithDefault #-} +#else +{-# INLINE findWithDefault #-} +#endif + +-- | /O(log n)/. Find largest key smaller than the given one and return the +-- corresponding (key, value) pair. +-- +-- > lookupLT 3 (fromList [(3,'a'), (5,'b')]) == Nothing +-- > lookupLT 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') +lookupLT :: Ord k => k -> Map k v -> Maybe (k, v) +lookupLT = goNothing + where + goNothing !_ Tip = Nothing + goNothing k (Bin _ kx x l r) | k <= kx = goNothing k l + | otherwise = goJust k kx x r + + goJust !_ kx' x' Tip = Just (kx', x') + goJust k kx' x' (Bin _ kx x l r) | k <= kx = goJust k kx' x' l + | otherwise = goJust k kx x r +#if __GLASGOW_HASKELL__ +{-# INLINABLE lookupLT #-} +#else +{-# INLINE lookupLT #-} +#endif + +-- | /O(log n)/. Find smallest key greater than the given one and return the +-- corresponding (key, value) pair. +-- +-- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') +-- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing +lookupGT :: Ord k => k -> Map k v -> Maybe (k, v) +lookupGT = goNothing + where + goNothing !_ Tip = Nothing + goNothing k (Bin _ kx x l r) | k < kx = goJust k kx x l + | otherwise = goNothing k r + + goJust !_ kx' x' Tip = Just (kx', x') + goJust k kx' x' (Bin _ kx x l r) | k < kx = goJust k kx x l + | otherwise = goJust k kx' x' r +#if __GLASGOW_HASKELL__ +{-# INLINABLE lookupGT #-} +#else +{-# INLINE lookupGT #-} +#endif + +-- | /O(log n)/. Find largest key smaller or equal to the given one and return +-- the corresponding (key, value) pair. +-- +-- > lookupLE 2 (fromList [(3,'a'), (5,'b')]) == Nothing +-- > lookupLE 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') +-- > lookupLE 5 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') +lookupLE :: Ord k => k -> Map k v -> Maybe (k, v) +lookupLE = goNothing + where + goNothing !_ Tip = Nothing + goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goNothing k l + EQ -> Just (kx, x) + GT -> goJust k kx x r + + goJust !_ kx' x' Tip = Just (kx', x') + goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx' x' l + EQ -> Just (kx, x) + GT -> goJust k kx x r +#if __GLASGOW_HASKELL__ +{-# INLINABLE lookupLE #-} +#else +{-# INLINE lookupLE #-} +#endif + +-- | /O(log n)/. Find smallest key greater or equal to the given one and return +-- the corresponding (key, value) pair. +-- +-- > lookupGE 3 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') +-- > lookupGE 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') +-- > lookupGE 6 (fromList [(3,'a'), (5,'b')]) == Nothing +lookupGE :: Ord k => k -> Map k v -> Maybe (k, v) +lookupGE = goNothing + where + goNothing !_ Tip = Nothing + goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l + EQ -> Just (kx, x) + GT -> goNothing k r + + goJust !_ kx' x' Tip = Just (kx', x') + goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l + EQ -> Just (kx, x) + GT -> goJust k kx' x' r +#if __GLASGOW_HASKELL__ +{-# INLINABLE lookupGE #-} +#else +{-# INLINE lookupGE #-} +#endif + +{-------------------------------------------------------------------- + Construction +--------------------------------------------------------------------} +-- | /O(1)/. The empty map. +-- +-- > empty == fromList [] +-- > size empty == 0 + +empty :: Map k a +empty = Tip +{-# INLINE empty #-} + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +-- > size (singleton 1 'a') == 1 + +singleton :: k -> a -> Map k a +singleton k x = Bin 1 k x Tip Tip +{-# INLINE singleton #-} + +{-------------------------------------------------------------------- + Insertion +--------------------------------------------------------------------} +-- | /O(log n)/. 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' + +-- See Note: Type of local 'go' function +insert :: Ord k => k -> a -> Map k a -> Map k a +insert = go + where + -- Unlike insertR, we only get sharing here + -- when the inserted value is at the same address + -- as the present value. We try anyway. If we decide + -- not to, then Data.Map.Strict should probably + -- get its own union implementation. + go :: Ord k => k -> a -> Map k a -> Map k a + go !kx x Tip = singleton kx x + go !kx x t@(Bin sz ky y l r) = + case compare kx ky of + LT | l' `ptrEq` l -> t + | otherwise -> balanceL ky y l' r + where !l' = go kx x l + GT | r' `ptrEq` r -> t + | otherwise -> balanceR ky y l r' + where !r' = go kx x r + EQ | kx `ptrEq` ky && x `ptrEq` y -> t + | otherwise -> Bin sz kx x l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE insert #-} +#else +{-# INLINE insert #-} +#endif + +-- Insert a new key and value in the map if it is not already present. +-- Used by `union`. + +-- See Note: Type of local 'go' function +insertR :: Ord k => k -> a -> Map k a -> Map k a +insertR = go + where + go :: Ord k => k -> a -> Map k a -> Map k a + go !kx x Tip = singleton kx x + go kx x t@(Bin _ ky y l r) = + case compare kx ky of + LT | l' `ptrEq` l -> t + | otherwise -> balanceL ky y l' r + where !l' = go kx x l + GT | r' `ptrEq` r -> t + | otherwise -> balanceR ky y l r' + where !r' = go kx x r + EQ -> t +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertR #-} +#else +{-# INLINE insertR #-} +#endif + +-- | /O(log n)/. 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" + +insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWith = go + where + -- We have no hope of making pointer equality tricks work + -- here, because lazy insertWith *always* changes the tree, + -- either adding a new entry or replacing an element with a + -- thunk. + go :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a + go _ !kx x Tip = singleton kx x + go f !kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go f kx x l) r + GT -> balanceR ky y l (go f kx x r) + EQ -> Bin sy kx (f x y) l r + +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertWith #-} +#else +{-# INLINE insertWith #-} +#endif + +-- | A helper function for 'unionWith'. When the key is already in +-- the map, the key is left alone, not replaced. The combining +-- function is flipped--it is applied to the old value and then the +-- new value. + +insertWithR :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWithR = go + where + go :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a + go _ !kx x Tip = singleton kx x + go f !kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go f kx x l) r + GT -> balanceR ky y l (go f kx x r) + EQ -> Bin sy ky (f y x) l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertWithR #-} +#else +{-# INLINE insertWithR #-} +#endif + +-- | /O(log n)/. 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" + +-- See Note: Type of local 'go' function +insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWithKey = go + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a + go _ !kx x Tip = singleton kx x + go f kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go f kx x l) r + GT -> balanceR ky y l (go f kx x r) + EQ -> Bin sy kx (f kx x y) l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertWithKey #-} +#else +{-# INLINE insertWithKey #-} +#endif + +-- | A helper function for 'unionWithKey'. When the key is already in +-- the map, the key is left alone, not replaced. The combining +-- function is flipped--it is applied to the old value and then the +-- new value. +insertWithKeyR :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWithKeyR = go + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a + go _ !kx x Tip = singleton kx x + go f kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go f kx x l) r + GT -> balanceR ky y l (go f kx x r) + EQ -> Bin sy ky (f ky y x) l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertWithKeyR #-} +#else +{-# INLINE insertWithKeyR #-} +#endif + +-- | /O(log n)/. 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") +-- +-- This is how to define @insertLookup@ using @insertLookupWithKey@: +-- +-- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t +-- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) +-- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) + +-- See Note: Type of local 'go' function +insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a + -> (Maybe a, Map k a) +insertLookupWithKey f0 k0 x0 = toPair . go f0 k0 x0 + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> StrictPair (Maybe a) (Map k a) + go _ !kx x Tip = (Nothing :*: singleton kx x) + go f kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> let !(found :*: l') = go f kx x l + !t' = balanceL ky y l' r + in (found :*: t') + GT -> let !(found :*: r') = go f kx x r + !t' = balanceR ky y l r' + in (found :*: t') + EQ -> (Just y :*: Bin sy kx (f kx x y) l r) +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertLookupWithKey #-} +#else +{-# INLINE insertLookupWithKey #-} +#endif + +{-------------------------------------------------------------------- + Deletion +--------------------------------------------------------------------} +-- | /O(log n)/. 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 + +-- See Note: Type of local 'go' function +delete :: Ord k => k -> Map k a -> Map k a +delete = go + where + go :: Ord k => k -> Map k a -> Map k a + go !_ Tip = Tip + go k t@(Bin _ kx x l r) = + case compare k kx of + LT | l' `ptrEq` l -> t + | otherwise -> balanceR kx x l' r + where !l' = go k l + GT | r' `ptrEq` r -> t + | otherwise -> balanceL kx x l r' + where !r' = go k r + EQ -> glue l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE delete #-} +#else +{-# INLINE delete #-} +#endif + +-- | /O(log n)/. 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 + +adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a +adjust f = adjustWithKey (\_ x -> f x) +#if __GLASGOW_HASKELL__ +{-# INLINABLE adjust #-} +#else +{-# INLINE adjust #-} +#endif + +-- | /O(log n)/. 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 + +adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a +adjustWithKey = go + where + go :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a + go _ !_ Tip = Tip + go f k (Bin sx kx x l r) = + case compare k kx of + LT -> Bin sx kx x (go f k l) r + GT -> Bin sx kx x l (go f k r) + EQ -> Bin sx kx (f kx x) l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE adjustWithKey #-} +#else +{-# INLINE adjustWithKey #-} +#endif + +-- | /O(log n)/. 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 :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a +update f = updateWithKey (\_ x -> f x) +#if __GLASGOW_HASKELL__ +{-# INLINABLE update #-} +#else +{-# INLINE update #-} +#endif + +-- | /O(log n)/. 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" + +-- See Note: Type of local 'go' function +updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a +updateWithKey = go + where + go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a + go _ !_ Tip = Tip + go f k(Bin sx kx x l r) = + case compare k kx of + LT -> balanceR kx x (go f k l) r + GT -> balanceL kx x l (go f k r) + EQ -> case f kx x of + Just x' -> Bin sx kx x' l r + Nothing -> glue l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE updateWithKey #-} +#else +{-# INLINE updateWithKey #-} +#endif + +-- | /O(log n)/. Lookup and update. See also 'updateWithKey'. +-- The function returns changed value, if it is updated. +-- Returns the original key value if the map entry is deleted. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "5:new a", fromList [(3, "b"), (5, "5:new a")]) +-- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) +-- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") + +-- See Note: Type of local 'go' function +updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) +updateLookupWithKey f0 k0 = toPair . go f0 k0 + where + go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> StrictPair (Maybe a) (Map k a) + go _ !_ Tip = (Nothing :*: Tip) + go f k (Bin sx kx x l r) = + case compare k kx of + LT -> let !(found :*: l') = go f k l + !t' = balanceR kx x l' r + in (found :*: t') + GT -> let !(found :*: r') = go f k r + !t' = balanceL kx x l r' + in (found :*: t') + EQ -> case f kx x of + Just x' -> (Just x' :*: Bin sx kx x' l r) + Nothing -> let !glued = glue l r + in (Just x :*: glued) +#if __GLASGOW_HASKELL__ +{-# INLINABLE updateLookupWithKey #-} +#else +{-# INLINE updateLookupWithKey #-} +#endif + +-- | /O(log n)/. 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 'Map'. +-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. +-- +-- > let f _ = Nothing +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > +-- > let f _ = Just "c" +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "c")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "c")] + +-- See Note: Type of local 'go' function +alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a +alter = go + where + go :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a + go f !k Tip = case f Nothing of + Nothing -> Tip + Just x -> singleton k x + + go f k (Bin sx kx x l r) = case compare k kx of + LT -> balance kx x (go f k l) r + GT -> balance kx x l (go f k r) + EQ -> case f (Just x) of + Just x' -> Bin sx kx x' l r + Nothing -> glue l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE alter #-} +#else +{-# INLINE alter #-} +#endif + +-- Used to choose the appropriate alterF implementation. +data AreWeStrict = Strict | Lazy + +-- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at +-- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete, +-- or update a value in a 'Map'. In short: @'lookup' k \<$\> 'alterF' f k m = f +-- ('lookup' k m)@. +-- +-- Example: +-- +-- @ +-- interactiveAlter :: Int -> Map Int String -> IO (Map Int String) +-- interactiveAlter k m = alterF f k m where +-- f Nothing -> do +-- putStrLn $ show k ++ +-- " was not found in the map. Would you like to add it?" +-- getUserResponse1 :: IO (Maybe String) +-- f (Just old) -> do +-- putStrLn "The key is currently bound to " ++ show old ++ +-- ". Would you like to change or delete it?" +-- getUserresponse2 :: IO (Maybe String) +-- @ +-- +-- 'alterF' is the most general operation for working with an individual +-- key that may or may not be in a given map. When used with trivial +-- functors like 'Identity' and 'Const', it is often slightly slower than +-- more specialized combinators like 'lookup' and 'insert'. However, when +-- the functor is non-trivial and key comparison is not particularly cheap, +-- it is the fastest way. +-- +-- Note on rewrite rules: +-- +-- This module includes GHC rewrite rules to optimize 'alterF' for +-- the 'Const' and 'Identity' functors. In general, these rules +-- improve performance. The sole exception is that when using +-- 'Identity', deleting a key that is already absent takes longer +-- than it would without the rules. If you expect this to occur +-- a very large fraction of the time, you might consider using a +-- private copy of the 'Identity' type. +-- +-- Note: 'alterF' is a flipped version of the 'at' combinator from +-- 'Control.Lens.At'. +-- +-- @since 0.5.8 +alterF :: (Functor f, Ord k) + => (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a) +alterF f k m = atKeyImpl Lazy k f m + +#ifndef __GLASGOW_HASKELL__ +{-# INLINE alterF #-} +#else +{-# INLINABLE [2] alterF #-} + +-- We can save a little time by recognizing the special case of +-- `Control.Applicative.Const` and just doing a lookup. +{-# RULES +"alterF/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterF f k = \m -> Const . getConst . f $ lookup k m + #-} + +#if MIN_VERSION_base(4,8,0) +-- base 4.8 and above include Data.Functor.Identity, so we can +-- save a pretty decent amount of time by handling it specially. +{-# RULES +"alterF/Identity" forall k f . alterF f k = atKeyIdentity k f + #-} +#endif +#endif + +atKeyImpl :: (Functor f, Ord k) => + AreWeStrict -> k -> (Maybe a -> f (Maybe a)) -> Map k a -> f (Map k a) +#if DEFINE_ALTERF_FALLBACK +atKeyImpl strict !k f m +-- It doesn't seem sensible to worry about overflowing the queue +-- if the word size is 61 or more. If I calculate it correctly, +-- that would take a map with nearly a quadrillion entries. + | wordSize < 61 && size m >= alterFCutoff = alterFFallback strict k f m +#endif +atKeyImpl strict !k f m = case lookupTrace k m of + TraceResult mv q -> (<$> f mv) $ \ fres -> + case fres of + Nothing -> case mv of + Nothing -> m + Just old -> deleteAlong old q m + Just new -> case strict of + Strict -> new `seq` case mv of + Nothing -> insertAlong q k new m + Just _ -> replaceAlong q new m + Lazy -> case mv of + Nothing -> insertAlong q k new m + Just _ -> replaceAlong q new m + +{-# INLINE atKeyImpl #-} + +#if DEFINE_ALTERF_FALLBACK +alterFCutoff :: Int +#if WORD_SIZE_IN_BITS == 32 +alterFCutoff = 55744454 +#else +alterFCutoff = case wordSize of + 30 -> 17637893 + 31 -> 31356255 + 32 -> 55744454 + x -> (4^(x*2-2)) `quot` (3^(x*2-2)) -- Unlikely +#endif +#endif + +data TraceResult a = TraceResult (Maybe a) {-# UNPACK #-} !BitQueue + +-- Look up a key and return a result indicating whether it was found +-- and what path was taken. +lookupTrace :: Ord k => k -> Map k a -> TraceResult a +lookupTrace = go emptyQB + where + go :: Ord k => BitQueueB -> k -> Map k a -> TraceResult a + go !q !_ Tip = TraceResult Nothing (buildQ q) + go q k (Bin _ kx x l r) = case compare k kx of + LT -> (go $! q `snocQB` False) k l + GT -> (go $! q `snocQB` True) k r + EQ -> TraceResult (Just x) (buildQ q) + +-- GHC 7.8 doesn't manage to unbox the queue properly +-- unless we explicitly inline this function. This stuff +-- is a bit touchy, unfortunately. +#if __GLASGOW_HASKELL__ >= 710 +{-# INLINABLE lookupTrace #-} +#else +{-# INLINE lookupTrace #-} +#endif + +-- Insert at a location (which will always be a leaf) +-- described by the path passed in. +insertAlong :: BitQueue -> k -> a -> Map k a -> Map k a +insertAlong !_ kx x Tip = singleton kx x +insertAlong q kx x (Bin sz ky y l r) = + case unconsQ q of + Just (False, tl) -> balanceL ky y (insertAlong tl kx x l) r + Just (True,tl) -> balanceR ky y l (insertAlong tl kx x r) + Nothing -> Bin sz kx x l r -- Shouldn't happen + +-- Delete from a location (which will always be a node) +-- described by the path passed in. +-- +-- This is fairly horrifying! We don't actually have any +-- use for the old value we're deleting. But if GHC sees +-- that, then it will allocate a thunk representing the +-- Map with the key deleted before we have any reason to +-- believe we'll actually want that. This transformation +-- enhances sharing, but we don't care enough about that. +-- So deleteAlong needs to take the old value, and we need +-- to convince GHC somehow that it actually uses it. We +-- can't NOINLINE deleteAlong, because that would prevent +-- the BitQueue from being unboxed. So instead we pass the +-- old value to a NOINLINE constant function and then +-- convince GHC that we use the result throughout the +-- computation. Doing the obvious thing and just passing +-- the value itself through the recursion costs 3-4% time, +-- so instead we convert the value to a magical zero-width +-- proxy that's ultimately erased. +deleteAlong :: any -> BitQueue -> Map k a -> Map k a +deleteAlong old !q0 !m = go (bogus old) q0 m where +#if USE_MAGIC_PROXY + go :: Proxy# () -> BitQueue -> Map k a -> Map k a +#else + go :: any -> BitQueue -> Map k a -> Map k a +#endif + go !_ !_ Tip = Tip + go foom q (Bin _ ky y l r) = + case unconsQ q of + Just (False, tl) -> balanceR ky y (go foom tl l) r + Just (True, tl) -> balanceL ky y l (go foom tl r) + Nothing -> glue l r + +#if USE_MAGIC_PROXY +{-# NOINLINE bogus #-} +bogus :: a -> Proxy# () +bogus _ = proxy# +#else +-- No point hiding in this case. +{-# INLINE bogus #-} +bogus :: a -> a +bogus a = a +#endif + +-- Replace the value found in the node described +-- by the given path with a new one. +replaceAlong :: BitQueue -> a -> Map k a -> Map k a +replaceAlong !_ _ Tip = Tip -- Should not happen +replaceAlong q x (Bin sz ky y l r) = + case unconsQ q of + Just (False, tl) -> Bin sz ky y (replaceAlong tl x l) r + Just (True,tl) -> Bin sz ky y l (replaceAlong tl x r) + Nothing -> Bin sz ky x l r + +#if __GLASGOW_HASKELL__ && MIN_VERSION_base(4,8,0) +atKeyIdentity :: Ord k => k -> (Maybe a -> Identity (Maybe a)) -> Map k a -> Identity (Map k a) +atKeyIdentity k f t = Identity $ atKeyPlain Lazy k (coerce f) t +{-# INLINABLE atKeyIdentity #-} + +atKeyPlain :: Ord k => AreWeStrict -> k -> (Maybe a -> Maybe a) -> Map k a -> Map k a +atKeyPlain strict k0 f0 t = case go k0 f0 t of + AltSmaller t' -> t' + AltBigger t' -> t' + AltAdj t' -> t' + AltSame -> t + where + go :: Ord k => k -> (Maybe a -> Maybe a) -> Map k a -> Altered k a + go !k f Tip = case f Nothing of + Nothing -> AltSame + Just x -> case strict of + Lazy -> AltBigger $ singleton k x + Strict -> x `seq` (AltBigger $ singleton k x) + + go k f (Bin sx kx x l r) = case compare k kx of + LT -> case go k f l of + AltSmaller l' -> AltSmaller $ balanceR kx x l' r + AltBigger l' -> AltBigger $ balanceL kx x l' r + AltAdj l' -> AltAdj $ Bin sx kx x l' r + AltSame -> AltSame + GT -> case go k f r of + AltSmaller r' -> AltSmaller $ balanceL kx x l r' + AltBigger r' -> AltBigger $ balanceR kx x l r' + AltAdj r' -> AltAdj $ Bin sx kx x l r' + AltSame -> AltSame + EQ -> case f (Just x) of + Just x' -> case strict of + Lazy -> AltAdj $ Bin sx kx x' l r + Strict -> x' `seq` (AltAdj $ Bin sx kx x' l r) + Nothing -> AltSmaller $ glue l r +{-# INLINE atKeyPlain #-} + +data Altered k a = AltSmaller !(Map k a) | AltBigger !(Map k a) | AltAdj !(Map k a) | AltSame +#endif + +#if DEFINE_ALTERF_FALLBACK +-- When the map is too large to use a bit queue, we fall back to +-- this much slower version which uses a more "natural" implementation +-- improved with Yoneda to avoid repeated fmaps. This works okayish for +-- some operations, but it's pretty lousy for lookups. +alterFFallback :: (Functor f, Ord k) + => AreWeStrict -> k -> (Maybe a -> f (Maybe a)) -> Map k a -> f (Map k a) +alterFFallback Lazy k f t = alterFYoneda k (\m q -> q <$> f m) t id +alterFFallback Strict k f t = alterFYoneda k (\m q -> q . forceMaybe <$> f m) t id + where + forceMaybe Nothing = Nothing + forceMaybe may@(Just !_) = may +{-# NOINLINE alterFFallback #-} + +alterFYoneda :: Ord k => + k -> (Maybe a -> (Maybe a -> b) -> f b) -> Map k a -> (Map k a -> b) -> f b +alterFYoneda = go + where + go :: Ord k => + k -> (Maybe a -> (Maybe a -> b) -> f b) -> Map k a -> (Map k a -> b) -> f b + go !k f Tip g = f Nothing $ \ mx -> case mx of + Nothing -> g Tip + Just x -> g (singleton k x) + go k f (Bin sx kx x l r) g = case compare k kx of + LT -> go k f l (\m -> g (balance kx x m r)) + GT -> go k f r (\m -> g (balance kx x l m)) + EQ -> f (Just x) $ \ mx' -> case mx' of + Just x' -> g (Bin sx kx x' l r) + Nothing -> g (glue l r) +{-# INLINE alterFYoneda #-} +#endif + +{-------------------------------------------------------------------- + Indexing +--------------------------------------------------------------------} +-- | /O(log n)/. Return the /index/ of a key, which is its zero-based index in +-- the sequence sorted by keys. 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 + +-- See Note: Type of local 'go' function +findIndex :: Ord k => k -> Map k a -> Int +findIndex = go 0 + where + go :: Ord k => Int -> k -> Map k a -> Int + go !_ !_ Tip = error "Map.findIndex: element is not in the map" + go idx k (Bin _ kx _ l r) = case compare k kx of + LT -> go idx k l + GT -> go (idx + size l + 1) k r + EQ -> idx + size l +#if __GLASGOW_HASKELL__ +{-# INLINABLE findIndex #-} +#endif + +-- | /O(log n)/. Lookup the /index/ of a key, which is its zero-based index in +-- the sequence sorted by keys. The index is a number from /0/ up to, but not +-- including, the 'size' of the map. +-- +-- > isJust (lookupIndex 2 (fromList [(5,"a"), (3,"b")])) == False +-- > fromJust (lookupIndex 3 (fromList [(5,"a"), (3,"b")])) == 0 +-- > fromJust (lookupIndex 5 (fromList [(5,"a"), (3,"b")])) == 1 +-- > isJust (lookupIndex 6 (fromList [(5,"a"), (3,"b")])) == False + +-- See Note: Type of local 'go' function +lookupIndex :: Ord k => k -> Map k a -> Maybe Int +lookupIndex = go 0 + where + go :: Ord k => Int -> k -> Map k a -> Maybe Int + go !_ !_ Tip = Nothing + go idx k (Bin _ kx _ l r) = case compare k kx of + LT -> go idx k l + GT -> go (idx + size l + 1) k r + EQ -> Just $! idx + size l +#if __GLASGOW_HASKELL__ +{-# INLINABLE lookupIndex #-} +#endif + +-- | /O(log n)/. Retrieve an element by its /index/, i.e. by its zero-based +-- index in the sequence sorted by keys. If the /index/ is out of range (less +-- than zero, greater or equal to 'size' of the map), 'error' is called. +-- +-- > 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 + +elemAt :: Int -> Map k a -> (k,a) +elemAt !_ Tip = error "Map.elemAt: index out of range" +elemAt i (Bin _ kx x l r) + = case compare i sizeL of + LT -> elemAt i l + GT -> elemAt (i-sizeL-1) r + EQ -> (kx,x) + where + sizeL = size l + +-- | Take a given number of entries in key order, beginning +-- with the smallest keys. +-- +-- @ +-- take n = 'fromDistinctAscList' . 'Prelude.take' n . 'toAscList' +-- @ + +take :: Int -> Map k a -> Map k a +take i m | i >= size m = m +take i0 m0 = go i0 m0 + where + go i !_ | i <= 0 = Tip + go !_ Tip = Tip + go i (Bin _ kx x l r) = + case compare i sizeL of + LT -> go i l + GT -> link kx x l (go (i - sizeL - 1) r) + EQ -> l + where sizeL = size l + +-- | Drop a given number of entries in key order, beginning +-- with the smallest keys. +-- +-- @ +-- drop n = 'fromDistinctAscList' . 'Prelude.drop' n . 'toAscList' +-- @ +drop :: Int -> Map k a -> Map k a +drop i m | i >= size m = Tip +drop i0 m0 = go i0 m0 + where + go i m | i <= 0 = m + go !_ Tip = Tip + go i (Bin _ kx x l r) = + case compare i sizeL of + LT -> link kx x (go i l) r + GT -> go (i - sizeL - 1) r + EQ -> insertMin kx x r + where sizeL = size l + +-- | /O(log n)/. Split a map at a particular index. +-- +-- @ +-- splitAt !n !xs = ('take' n xs, 'drop' n xs) +-- @ +splitAt :: Int -> Map k a -> (Map k a, Map k a) +splitAt i0 m0 + | i0 >= size m0 = (m0, Tip) + | otherwise = toPair $ go i0 m0 + where + go i m | i <= 0 = Tip :*: m + go !_ Tip = Tip :*: Tip + go i (Bin _ kx x l r) + = case compare i sizeL of + LT -> case go i l of + ll :*: lr -> ll :*: link kx x lr r + GT -> case go (i - sizeL - 1) r of + rl :*: rr -> link kx x l rl :*: rr + EQ -> l :*: insertMin kx x r + where sizeL = size l + +-- | /O(log n)/. Update the element at /index/, i.e. by its zero-based index in +-- the sequence sorted by keys. If the /index/ is out of range (less than zero, +-- greater or equal to 'size' of the map), 'error' is called. +-- +-- > 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 + +updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a +updateAt f !i t = + case t of + Tip -> error "Map.updateAt: index out of range" + Bin sx kx x l r -> case compare i sizeL of + LT -> balanceR kx x (updateAt f i l) r + GT -> balanceL kx x l (updateAt f (i-sizeL-1) r) + EQ -> case f kx x of + Just x' -> Bin sx kx x' l r + Nothing -> glue l r + where + sizeL = size l + +-- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based index in +-- the sequence sorted by keys. If the /index/ is out of range (less than zero, +-- greater or equal to 'size' of the map), 'error' is called. +-- +-- > 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 + +deleteAt :: Int -> Map k a -> Map k a +deleteAt !i t = + case t of + Tip -> error "Map.deleteAt: index out of range" + Bin _ kx x l r -> case compare i sizeL of + LT -> balanceR kx x (deleteAt i l) r + GT -> balanceL kx x l (deleteAt (i-sizeL-1) r) + EQ -> glue l r + where + sizeL = size l + + +{-------------------------------------------------------------------- + Minimal, Maximal +--------------------------------------------------------------------} +-- | /O(log n)/. 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 + +findMin :: Map k a -> (k,a) +findMin (Bin _ kx x Tip _) = (kx,x) +findMin (Bin _ _ _ l _) = findMin l +findMin Tip = error "Map.findMin: empty map has no minimal element" + +-- | /O(log n)/. 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 + +findMax :: Map k a -> (k,a) +findMax (Bin _ kx x _ Tip) = (kx,x) +findMax (Bin _ _ _ _ r) = findMax r +findMax Tip = error "Map.findMax: empty map has no maximal element" + +-- | /O(log n)/. 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 + +deleteMin :: Map k a -> Map k a +deleteMin (Bin _ _ _ Tip r) = r +deleteMin (Bin _ kx x l r) = balanceR kx x (deleteMin l) r +deleteMin Tip = Tip + +-- | /O(log n)/. 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 + +deleteMax :: Map k a -> Map k a +deleteMax (Bin _ _ _ l Tip) = l +deleteMax (Bin _ kx x l r) = balanceL kx x l (deleteMax r) +deleteMax Tip = Tip + +-- | /O(log n)/. 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" + +updateMin :: (a -> Maybe a) -> Map k a -> Map k a +updateMin f m + = updateMinWithKey (\_ x -> f x) m + +-- | /O(log n)/. 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" + +updateMax :: (a -> Maybe a) -> Map k a -> Map k a +updateMax f m + = updateMaxWithKey (\_ x -> f x) m + + +-- | /O(log n)/. 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" + +updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a +updateMinWithKey _ Tip = Tip +updateMinWithKey f (Bin sx kx x Tip r) = case f kx x of + Nothing -> r + Just x' -> Bin sx kx x' Tip r +updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r + +-- | /O(log n)/. 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" + +updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a +updateMaxWithKey _ Tip = Tip +updateMaxWithKey f (Bin sx kx x l Tip) = case f kx x of + Nothing -> l + Just x' -> Bin sx kx x' l Tip +updateMaxWithKey f (Bin _ kx x l r) = balanceL kx x l (updateMaxWithKey f r) + +-- | /O(log n)/. 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 :: Map k a -> Maybe ((k,a), Map k a) +minViewWithKey Tip = Nothing +minViewWithKey x = Just $! deleteFindMin x + +-- | /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 :: Map k a -> Maybe ((k,a), Map k a) +maxViewWithKey Tip = Nothing +maxViewWithKey x = Just $! deleteFindMax x + +-- | /O(log n)/. 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 + +minView :: Map k a -> Maybe (a, Map k a) +minView Tip = Nothing +minView x = Just $! (first snd $ deleteFindMin x) + +-- | /O(log n)/. Retrieves the value associated with maximal key of the +-- map, and the map stripped of that element, or 'Nothing' if passed an +-- empty map. +-- +-- > maxView (fromList [(5,"a"), (3,"b")]) == Just ("a", singleton 3 "b") +-- > maxView empty == Nothing + +maxView :: Map k a -> Maybe (a, Map k a) +maxView Tip = Nothing +maxView x = Just $! (first snd $ deleteFindMax x) + +-- Update the 1st component of a tuple (stricter version of +-- Control.Arrow.first) +first :: (a -> b) -> (a,c) -> (b,c) +first f (x,y) = (f x, y) + +{-------------------------------------------------------------------- + Union. +--------------------------------------------------------------------} +-- | The union of a list of maps: +-- (@'unions' == 'Prelude.foldl' 'union' 'empty'@). +-- +-- > unions [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] +-- > == fromList [(3, "b"), (5, "a"), (7, "C")] +-- > unions [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])] +-- > == fromList [(3, "B3"), (5, "A3"), (7, "C")] + +unions :: Ord k => [Map k a] -> Map k a +unions ts + = foldlStrict union empty ts +#if __GLASGOW_HASKELL__ +{-# INLINABLE unions #-} +#endif + +-- | The union of a list of maps, with a combining operation: +-- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@). +-- +-- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] +-- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] + +unionsWith :: Ord k => (a->a->a) -> [Map k a] -> Map k a +unionsWith f ts + = foldlStrict (unionWith f) empty ts +#if __GLASGOW_HASKELL__ +{-# INLINABLE unionsWith #-} +#endif + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- 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'@). +-- +-- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")] + +union :: Ord k => Map k a -> Map k a -> Map k a +union t1 Tip = t1 +union t1 (Bin _ k x Tip Tip) = insertR k x t1 +union (Bin _ k x Tip Tip) t2 = insert k x t2 +union Tip t2 = t2 +union t1@(Bin _ k1 x1 l1 r1) t2 = case split k1 t2 of + (l2, r2) | l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 -> t1 + | otherwise -> link k1 x1 l1l2 r1r2 + where !l1l2 = union l1 l2 + !r1r2 = union r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE union #-} +#endif + +{-------------------------------------------------------------------- + Union with a combining function +--------------------------------------------------------------------} +-- | /O(m*log(n\/m + 1)), m <= n/. Union with a combining function. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] + +unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a +-- QuickCheck says pointer equality never happens here. +unionWith _f t1 Tip = t1 +unionWith f t1 (Bin _ k x Tip Tip) = insertWithR f k x t1 +unionWith f (Bin _ k x Tip Tip) t2 = insertWith f k x t2 +unionWith _f Tip t2 = t2 +unionWith f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of + (l2, mb, r2) -> case mb of + Nothing -> link k1 x1 l1l2 r1r2 + Just x2 -> link k1 (f x1 x2) l1l2 r1r2 + where !l1l2 = unionWith f l1 l2 + !r1r2 = unionWith f r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE unionWith #-} +#endif + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- Union with a combining function. +-- +-- > 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")] + +unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a +unionWithKey _f t1 Tip = t1 +unionWithKey f t1 (Bin _ k x Tip Tip) = insertWithKeyR f k x t1 +unionWithKey f (Bin _ k x Tip Tip) t2 = insertWithKey f k x t2 +unionWithKey _f Tip t2 = t2 +unionWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of + (l2, mb, r2) -> case mb of + Nothing -> link k1 x1 l1l2 r1r2 + Just x2 -> link k1 (f k1 x1 x2) l1l2 r1r2 + where !l1l2 = unionWithKey f l1 l2 + !r1r2 = unionWithKey f r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE unionWithKey #-} +#endif + +{-------------------------------------------------------------------- + Difference +--------------------------------------------------------------------} + +-- | /O(m*log(n\/m + 1)), m <= n/. Difference of two maps. +-- Return elements of the first map not existing in the second map. +-- +-- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b" + +difference :: Ord k => Map k a -> Map k b -> Map k a +difference Tip _ = Tip +difference t1 Tip = t1 +difference t1 (Bin _ k _ l2 r2) = case split k t1 of + (l1, r1) + | size l1l2 + size r1r2 == size t1 -> t1 + | otherwise -> link2 l1l2 r1r2 + where + !l1l2 = difference l1 l2 + !r1r2 = difference r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE difference #-} +#endif + +-- | /O(m*log(n/m + 1)), m <= n/. Remove all keys in a 'Set' from a 'Map'. +-- +-- @ +-- m `withoutKeys` s = 'filterWithKey' (\k _ -> k `'Set.notMember'` s) m +-- @ +-- +-- @since 0.5.8 + +withoutKeys :: Ord k => Map k a -> Set k -> Map k a +withoutKeys Tip _ = Tip +withoutKeys m Set.Tip = m +withoutKeys m (Set.Bin _ k ls rs) = case splitMember k m of + (lm, b, rm) + | not b && lm' `ptrEq` lm && rm' `ptrEq` rm -> m + | otherwise -> link2 lm' rm' + where + !lm' = withoutKeys lm ls + !rm' = withoutKeys rm rs +#if __GLASGOW_HASKELL__ +{-# INLINABLE withoutKeys #-} +#endif + +-- | /O(n+m)/. 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@. +-- +-- > 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" +differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a +differenceWith f = merge preserveMissing dropMissing $ + zipWithMaybeMatched (\_ x y -> f x y) +#if __GLASGOW_HASKELL__ +{-# INLINABLE differenceWith #-} +#endif + +-- | /O(n+m)/. 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@. +-- +-- > 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" + +differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a +differenceWithKey f = + merge preserveMissing dropMissing (zipWithMaybeMatched f) +#if __GLASGOW_HASKELL__ +{-# INLINABLE differenceWithKey #-} +#endif + + +{-------------------------------------------------------------------- + Intersection +--------------------------------------------------------------------} +-- | /O(m*log(n\/m + 1)), m <= n/. 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" + +intersection :: Ord k => Map k a -> Map k b -> Map k a +intersection Tip _ = Tip +intersection _ Tip = Tip +intersection t1@(Bin _ k x l1 r1) t2 + | mb = if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 + then t1 + else link k x l1l2 r1r2 + | otherwise = link2 l1l2 r1r2 + where + !(l2, mb, r2) = splitMember k t2 + !l1l2 = intersection l1 l2 + !r1r2 = intersection r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE intersection #-} +#endif + +-- | /O(m*log(n/m + 1)), m <= n/. Restrict a 'Map' to only those keys +-- found in a 'Set'. +-- +-- @ +-- m `restrictKeys` s = 'filterWithKey' (\k _ -> k `'Set.member'` s) m +-- @ +-- +-- @since 0.5.8 +restrictKeys :: Ord k => Map k a -> Set k -> Map k a +restrictKeys Tip _ = Tip +restrictKeys _ Set.Tip = Tip +restrictKeys m@(Bin _ k x l1 r1) s + | b = if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 + then m + else link k x l1l2 r1r2 + | otherwise = link2 l1l2 r1r2 + where + !(l2, b, r2) = Set.splitMember k s + !l1l2 = restrictKeys l1 l2 + !r1r2 = restrictKeys r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE restrictKeys #-} +#endif + +-- | /O(m*log(n\/m + 1)), m <= n/. Intersection with a combining function. +-- +-- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" + +intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c +-- We have no hope of pointer equality tricks here because every single +-- element in the result will be a thunk. +intersectionWith _f Tip _ = Tip +intersectionWith _f _ Tip = Tip +intersectionWith f (Bin _ k x1 l1 r1) t2 = case mb of + Just x2 -> link k (f x1 x2) l1l2 r1r2 + Nothing -> link2 l1l2 r1r2 + where + !(l2, mb, r2) = splitLookup k t2 + !l1l2 = intersectionWith f l1 l2 + !r1r2 = intersectionWith f r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE intersectionWith #-} +#endif + +-- | /O(m*log(n\/m + 1)), m <= n/. Intersection with a combining function. +-- +-- > 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" + +intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c +intersectionWithKey _f Tip _ = Tip +intersectionWithKey _f _ Tip = Tip +intersectionWithKey f (Bin _ k x1 l1 r1) t2 = case mb of + Just x2 -> link k (f k x1 x2) l1l2 r1r2 + Nothing -> link2 l1l2 r1r2 + where + !(l2, mb, r2) = splitLookup k t2 + !l1l2 = intersectionWithKey f l1 l2 + !r1r2 = intersectionWithKey f r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE intersectionWithKey #-} +#endif + +#if !MIN_VERSION_base (4,8,0) +-- | The identity type. +newtype Identity a = Identity { runIdentity :: a } +#if __GLASGOW_HASKELL__ == 708 +instance Functor Identity where + fmap = coerce +instance Applicative Identity where + (<*>) = coerce + pure = Identity +#else +instance Functor Identity where + fmap f (Identity a) = Identity (f a) +instance Applicative Identity where + Identity f <*> Identity x = Identity (f x) + pure = Identity +#endif +#endif + +-- | A tactic for dealing with keys present in one map but not the other in +-- 'merge' or 'mergeA'. +-- +-- A tactic of type @ WhenMissing f k x z @ is an abstract representation +-- of a function of type @ k -> x -> f (Maybe z) @. + +data WhenMissing f k x y = WhenMissing + { missingSubtree :: Map k x -> f (Map k y) + , missingKey :: k -> x -> f (Maybe y)} + +instance (Applicative f, Monad f) => Functor (WhenMissing f k x) where + fmap = mapWhenMissing + {-# INLINE fmap #-} + +instance (Applicative f, Monad f) + => Category.Category (WhenMissing f k) where + id = preserveMissing + f . g = traverseMaybeMissing $ + \ k x -> missingKey g k x >>= \y -> + case y of + Nothing -> pure Nothing + Just q -> missingKey f k q + {-# INLINE id #-} + {-# INLINE (.) #-} + +-- | Equivalent to @ ReaderT k (ReaderT x (MaybeT f)) @. +instance (Applicative f, Monad f) => Applicative (WhenMissing f k x) where + pure x = mapMissing (\ _ _ -> x) + f <*> g = traverseMaybeMissing $ \k x -> do + res1 <- missingKey f k x + case res1 of + Nothing -> pure Nothing + Just r -> (pure $!) . fmap r =<< missingKey g k x + {-# INLINE pure #-} + {-# INLINE (<*>) #-} + +-- | Equivalent to @ ReaderT k (ReaderT x (MaybeT f)) @. +instance (Applicative f, Monad f) => Monad (WhenMissing f k x) where +#if !MIN_VERSION_base(4,8,0) + return = pure +#endif + m >>= f = traverseMaybeMissing $ \k x -> do + res1 <- missingKey m k x + case res1 of + Nothing -> pure Nothing + Just r -> missingKey (f r) k x + {-# INLINE (>>=) #-} + +-- | Map covariantly over a @'WhenMissing' f k x@. +mapWhenMissing :: (Applicative f, Monad f) + => (a -> b) + -> WhenMissing f k x a -> WhenMissing f k x b +mapWhenMissing f t = WhenMissing + { missingSubtree = \m -> missingSubtree t m >>= \m' -> pure $! fmap f m' + , missingKey = \k x -> missingKey t k x >>= \q -> (pure $! fmap f q) } +{-# INLINE mapWhenMissing #-} + +-- | Map covariantly over a @'WhenMissing' f k x@, using only a 'Functor f' +-- constraint. +mapGentlyWhenMissing :: Functor f + => (a -> b) + -> WhenMissing f k x a -> WhenMissing f k x b +mapGentlyWhenMissing f t = WhenMissing + { missingSubtree = \m -> fmap f <$> missingSubtree t m + , missingKey = \k x -> fmap f <$> missingKey t k x } +{-# INLINE mapGentlyWhenMissing #-} + +-- | Map covariantly over a @'WhenMatched' f k x@, using only a 'Functor f' +-- constraint. +mapGentlyWhenMatched :: Functor f + => (a -> b) + -> WhenMatched f k x y a -> WhenMatched f k x y b +mapGentlyWhenMatched f t = zipWithMaybeAMatched $ + \k x y -> fmap f <$> runWhenMatched t k x y +{-# INLINE mapGentlyWhenMatched #-} + +-- | Map contravariantly over a @'WhenMissing' f k _ x@. +lmapWhenMissing :: (b -> a) -> WhenMissing f k a x -> WhenMissing f k b x +lmapWhenMissing f t = WhenMissing + { missingSubtree = \m -> missingSubtree t (fmap f m) + , missingKey = \k x -> missingKey t k (f x) } +{-# INLINE lmapWhenMissing #-} + +-- | Map contravariantly over a @'WhenMatched' f k _ y z@. +contramapFirstWhenMatched :: (b -> a) + -> WhenMatched f k a y z + -> WhenMatched f k b y z +contramapFirstWhenMatched f t = WhenMatched $ + \k x y -> runWhenMatched t k (f x) y +{-# INLINE contramapFirstWhenMatched #-} + +-- | Map contravariantly over a @'WhenMatched' f k x _ z@. +contramapSecondWhenMatched :: (b -> a) + -> WhenMatched f k x a z + -> WhenMatched f k x b z +contramapSecondWhenMatched f t = WhenMatched $ + \k x y -> runWhenMatched t k x (f y) +{-# INLINE contramapSecondWhenMatched #-} + +-- | A tactic for dealing with keys present in one map but not the other in +-- 'merge'. +-- +-- A tactic of type @ SimpleWhenMissing k x z @ is an abstract representation +-- of a function of type @ k -> x -> Maybe z @. +type SimpleWhenMissing = WhenMissing Identity + +-- | A tactic for dealing with keys present in both +-- maps in 'merge' or 'mergeA'. +-- +-- A tactic of type @ WhenMatched f k x y z @ is an abstract representation +-- of a function of type @ k -> x -> y -> f (Maybe z) @. +newtype WhenMatched f k x y z = WhenMatched + { matchedKey :: k -> x -> y -> f (Maybe z) } + +-- | Along with zipWithMaybeAMatched, witnesses the isomorphism between +-- @WhenMatched f k x y z@ and @k -> x -> y -> f (Maybe z)@. +runWhenMatched :: WhenMatched f k x y z -> k -> x -> y -> f (Maybe z) +runWhenMatched = matchedKey +{-# INLINE runWhenMatched #-} + +-- | Along with traverseMaybeMissing, witnesses the isomorphism between +-- @WhenMissing f k x y@ and @k -> x -> f (Maybe y)@. +runWhenMissing :: WhenMissing f k x y -> k -> x -> f (Maybe y) +runWhenMissing = missingKey +{-# INLINE runWhenMissing #-} + +instance Functor f => Functor (WhenMatched f k x y) where + fmap = mapWhenMatched + {-# INLINE fmap #-} + +instance (Monad f, Applicative f) => Category.Category (WhenMatched f k x) where + id = zipWithMatched (\_ _ y -> y) + f . g = zipWithMaybeAMatched $ + \k x y -> do + res <- runWhenMatched g k x y + case res of + Nothing -> pure Nothing + Just r -> runWhenMatched f k x r + {-# INLINE id #-} + {-# INLINE (.) #-} + +-- | Equivalent to @ ReaderT k (ReaderT x (ReaderT y (MaybeT f))) @ +instance (Monad f, Applicative f) => Applicative (WhenMatched f k x y) where + pure x = zipWithMatched (\_ _ _ -> x) + fs <*> xs = zipWithMaybeAMatched $ \k x y -> do + res <- runWhenMatched fs k x y + case res of + Nothing -> pure Nothing + Just r -> (pure $!) . fmap r =<< runWhenMatched xs k x y + {-# INLINE pure #-} + {-# INLINE (<*>) #-} + +-- | Equivalent to @ ReaderT k (ReaderT x (ReaderT y (MaybeT f))) @ +instance (Monad f, Applicative f) => Monad (WhenMatched f k x y) where +#if !MIN_VERSION_base(4,8,0) + return = pure +#endif + m >>= f = zipWithMaybeAMatched $ \k x y -> do + res <- runWhenMatched m k x y + case res of + Nothing -> pure Nothing + Just r -> runWhenMatched (f r) k x y + {-# INLINE (>>=) #-} + +-- | Map covariantly over a @'WhenMatched' f k x y@. +mapWhenMatched :: Functor f + => (a -> b) + -> WhenMatched f k x y a + -> WhenMatched f k x y b +mapWhenMatched f (WhenMatched g) = WhenMatched $ \k x y -> fmap (fmap f) (g k x y) +{-# INLINE mapWhenMatched #-} + +-- | A tactic for dealing with keys present in both maps in 'merge'. +-- +-- A tactic of type @ SimpleWhenMatched k x y z @ is an abstract representation +-- of a function of type @ k -> x -> y -> Maybe z @. +type SimpleWhenMatched = WhenMatched Identity + +-- | When a key is found in both maps, apply a function to the +-- key and values and use the result in the merged map. +-- +-- @ +-- zipWithMatched :: (k -> x -> y -> z) +-- -> SimpleWhenMatched k x y z +-- @ +zipWithMatched :: Applicative f + => (k -> x -> y -> z) + -> WhenMatched f k x y z +zipWithMatched f = WhenMatched $ \ k x y -> pure . Just $ f k x y +{-# INLINE zipWithMatched #-} + +-- | When a key is found in both maps, apply a function to the +-- key and values to produce an action and use its result in the merged map. +zipWithAMatched :: Applicative f + => (k -> x -> y -> f z) + -> WhenMatched f k x y z +zipWithAMatched f = WhenMatched $ \ k x y -> Just <$> f k x y +{-# INLINE zipWithAMatched #-} + +-- | When a key is found in both maps, apply a function to the +-- key and values and maybe use the result in the merged map. +-- +-- @ +-- zipWithMaybeMatched :: (k -> x -> y -> Maybe z) +-- -> SimpleWhenMatched k x y z +-- @ +zipWithMaybeMatched :: Applicative f + => (k -> x -> y -> Maybe z) + -> WhenMatched f k x y z +zipWithMaybeMatched f = WhenMatched $ \ k x y -> pure $ f k x y +{-# INLINE zipWithMaybeMatched #-} + +-- | When a key is found in both maps, apply a function to the +-- key and values, perform the resulting action, and maybe use +-- the result in the merged map. +-- +-- This is the fundamental 'WhenMatched' tactic. +zipWithMaybeAMatched :: (k -> x -> y -> f (Maybe z)) + -> WhenMatched f k x y z +zipWithMaybeAMatched f = WhenMatched $ \ k x y -> f k x y +{-# INLINE zipWithMaybeAMatched #-} + +-- | Drop all the entries whose keys are missing from the other +-- map. +-- +-- @ +-- dropMissing :: SimpleWhenMissing k x y +-- @ +-- +-- prop> dropMissing = mapMaybeMissing (\_ _ -> Nothing) +-- +-- but @dropMissing@ is much faster. +dropMissing :: Applicative f => WhenMissing f k x y +dropMissing = WhenMissing + { missingSubtree = const (pure Tip) + , missingKey = \_ _ -> pure Nothing } +{-# INLINE dropMissing #-} + +-- | Preserve, unchanged, the entries whose keys are missing from +-- the other map. +-- +-- @ +-- preserveMissing :: SimpleWhenMissing k x x +-- @ +-- +-- prop> preserveMissing = Lazy.Merge.mapMaybeMissing (\_ x -> Just x) +-- +-- but @preserveMissing@ is much faster. +preserveMissing :: Applicative f => WhenMissing f k x x +preserveMissing = WhenMissing + { missingSubtree = pure + , missingKey = \_ v -> pure (Just v) } +{-# INLINE preserveMissing #-} + +-- | Map over the entries whose keys are missing from the other map. +-- +-- @ +-- mapMissing :: (k -> x -> y) -> SimpleWhenMissing k x y +-- @ +-- +-- prop> mapMissing f = mapMaybeMissing (\k x -> Just $ f k x) +-- +-- but @mapMissing@ is somewhat faster. +mapMissing :: Applicative f => (k -> x -> y) -> WhenMissing f k x y +mapMissing f = WhenMissing + { missingSubtree = \m -> pure $! mapWithKey f m + , missingKey = \ k x -> pure $ Just (f k x) } +{-# INLINE mapMissing #-} + +-- | Map over the entries whose keys are missing from the other map, +-- optionally removing some. This is the most powerful 'SimpleWhenMissing' +-- tactic, but others are usually more efficient. +-- +-- @ +-- mapMaybeMissing :: (k -> x -> Maybe y) -> SimpleWhenMissing k x y +-- @ +-- +-- prop> mapMaybeMissing f = traverseMaybeMissing (\k x -> pure (f k x)) +-- +-- but @mapMaybeMissing@ uses fewer unnecessary 'Applicative' operations. +mapMaybeMissing :: Applicative f => (k -> x -> Maybe y) -> WhenMissing f k x y +mapMaybeMissing f = WhenMissing + { missingSubtree = \m -> pure $! mapMaybeWithKey f m + , missingKey = \k x -> pure $! f k x } +{-# INLINE mapMaybeMissing #-} + +-- | Filter the entries whose keys are missing from the other map. +-- +-- @ +-- filterMissing :: (k -> x -> Bool) -> SimpleWhenMissing k x x +-- @ +-- +-- prop> filterMissing f = Lazy.Merge.mapMaybeMissing $ \k x -> guard (f k x) *> Just x +-- +-- but this should be a little faster. +filterMissing :: Applicative f + => (k -> x -> Bool) -> WhenMissing f k x x +filterMissing f = WhenMissing + { missingSubtree = \m -> pure $! filterWithKey f m + , missingKey = \k x -> pure $! if f k x then Just x else Nothing } +{-# INLINE filterMissing #-} + +-- | Filter the entries whose keys are missing from the other map +-- using some 'Applicative' action. +-- +-- @ +-- filterAMissing f = Lazy.Merge.traverseMaybeMissing $ +-- \k x -> (\b -> guard b *> Just x) <$> f k x +-- @ +-- +-- but this should be a little faster. +filterAMissing :: Applicative f + => (k -> x -> f Bool) -> WhenMissing f k x x +filterAMissing f = WhenMissing + { missingSubtree = \m -> filterWithKeyA f m + , missingKey = \k x -> bool Nothing (Just x) <$> f k x } +{-# INLINE filterAMissing #-} + +-- | This wasn't in Data.Bool until 4.7.0, so we define it here +bool :: a -> a -> Bool -> a +bool f _ False = f +bool _ t True = t + +-- | Traverse over the entries whose keys are missing from the other map. +traverseMissing :: Applicative f + => (k -> x -> f y) -> WhenMissing f k x y +traverseMissing f = WhenMissing + { missingSubtree = traverseWithKey f + , missingKey = \k x -> Just <$> f k x } +{-# INLINE traverseMissing #-} + +-- | Traverse over the entries whose keys are missing from the other map, +-- optionally producing values to put in the result. +-- This is the most powerful 'WhenMissing' tactic, but others are usually +-- more efficient. +traverseMaybeMissing :: Applicative f + => (k -> x -> f (Maybe y)) -> WhenMissing f k x y +traverseMaybeMissing f = WhenMissing + { missingSubtree = traverseMaybeWithKey f + , missingKey = f } +{-# INLINE traverseMaybeMissing #-} + +-- | Merge two maps. +-- +-- @merge@ takes two 'WhenMissing' tactics, a 'WhenMatched' +-- tactic and two maps. It uses the tactics to merge the maps. +-- Its behavior is best understood via its fundamental tactics, +-- 'mapMaybeMissing' and 'zipWithMaybeMatched'. +-- +-- Consider +-- +-- @ +-- merge (mapMaybeMissing g1) +-- (mapMaybeMissing g2) +-- (zipWithMaybeMatched f) +-- m1 m2 +-- @ +-- +-- Take, for example, +-- +-- @ +-- m1 = [(0, 'a'), (1, 'b'), (3,'c'), (4, 'd')] +-- m2 = [(1, "one"), (2, "two"), (4, "three")] +-- @ +-- +-- @merge@ will first ''align'' these maps by key: +-- +-- @ +-- m1 = [(0, 'a'), (1, 'b'), (3,'c'), (4, 'd')] +-- m2 = [(1, "one"), (2, "two"), (4, "three")] +-- @ +-- +-- It will then pass the individual entries and pairs of entries +-- to @g1@, @g2@, or @f@ as appropriate: +-- +-- @ +-- maybes = [g1 0 'a', f 1 'b' "one", g2 2 "two", g1 3 'c', f 4 'd' "three"] +-- @ +-- +-- This produces a 'Maybe' for each key: +-- +-- @ +-- keys = 0 1 2 3 4 +-- results = [Nothing, Just True, Just False, Nothing, Just True] +-- @ +-- +-- Finally, the @Just@ results are collected into a map: +-- +-- @ +-- return value = [(1, True), (2, False), (4, True)] +-- @ +-- +-- The other tactics below are optimizations or simplifications of +-- 'mapMaybeMissing' for special cases. Most importantly, +-- +-- * 'dropMissing' drops all the keys. +-- * 'preserveMissing' leaves all the entries alone. +-- +-- When 'merge' is given three arguments, it is inlined at the call +-- site. To prevent excessive inlining, you should typically use 'merge' +-- to define your custom combining functions. +-- +-- +-- Examples: +-- +-- prop> unionWithKey f = merge preserveMissing preserveMissing (zipWithMatched f) +-- prop> intersectionWithKey f = merge dropMissing dropMissing (zipWithMatched f) +-- prop> differenceWith f = merge diffPreserve diffDrop f +-- prop> symmetricDifference = merge diffPreserve diffPreserve (\ _ _ _ -> Nothing) +-- prop> mapEachPiece f g h = merge (diffMapWithKey f) (diffMapWithKey g) +-- +-- @since 0.5.8 +merge :: Ord k + => SimpleWhenMissing k a c -- ^ What to do with keys in @m1@ but not @m2@ + -> SimpleWhenMissing k b c -- ^ What to do with keys in @m2@ but not @m1@ + -> SimpleWhenMatched k a b c -- ^ What to do with keys in both @m1@ and @m2@ + -> Map k a -- ^ Map @m1@ + -> Map k b -- ^ Map @m2@ + -> Map k c +merge g1 g2 f m1 m2 = runIdentity $ + mergeA g1 g2 f m1 m2 +{-# INLINE merge #-} + +-- | An applicative version of 'merge'. +-- +-- @mergeA@ takes two 'WhenMissing' tactics, a 'WhenMatched' +-- tactic and two maps. It uses the tactics to merge the maps. +-- Its behavior is best understood via its fundamental tactics, +-- 'traverseMaybeMissing' and 'zipWithMaybeAMatched'. +-- +-- Consider +-- +-- @ +-- mergeA (traverseMaybeMissing g1) +-- (traverseMaybeMissing g2) +-- (zipWithMaybeAMatched f) +-- m1 m2 +-- @ +-- +-- Take, for example, +-- +-- @ +-- m1 = [(0, 'a'), (1, 'b'), (3,'c'), (4, 'd')] +-- m2 = [(1, "one"), (2, "two"), (4, "three")] +-- @ +-- +-- @mergeA@ will first ''align'' these maps by key: +-- +-- @ +-- m1 = [(0, 'a'), (1, 'b'), (3,'c'), (4, 'd')] +-- m2 = [(1, "one"), (2, "two"), (4, "three")] +-- @ +-- +-- It will then pass the individual entries and pairs of entries +-- to @g1@, @g2@, or @f@ as appropriate: +-- +-- @ +-- actions = [g1 0 'a', f 1 'b' "one", g2 2 "two", g1 3 'c', f 4 'd' "three"] +-- @ +-- +-- Next, it will perform the actions in the @actions@ list in order from +-- left to right. +-- +-- @ +-- keys = 0 1 2 3 4 +-- results = [Nothing, Just True, Just False, Nothing, Just True] +-- @ +-- +-- Finally, the @Just@ results are collected into a map: +-- +-- @ +-- return value = [(1, True), (2, False), (4, True)] +-- @ +-- +-- The other tactics below are optimizations or simplifications of +-- 'traverseMaybeMissing' for special cases. Most importantly, +-- +-- * 'dropMissing' drops all the keys. +-- * 'preserveMissing' leaves all the entries alone. +-- * 'mapMaybeMissing' does not use the 'Applicative' context. +-- +-- When 'mergeA' is given three arguments, it is inlined at the call +-- site. To prevent excessive inlining, you should generally only use +-- 'mergeA' to define custom combining functions. +-- +-- @since 0.5.8 +mergeA :: (Applicative f, Ord k) + => WhenMissing f k a c -- ^ What to do with keys in @m1@ but not @m2@ + -> WhenMissing f k b c -- ^ What to do with keys in @m2@ but not @m1@ + -> WhenMatched f k a b c -- ^ What to do with keys in both @m1@ and @m2@ + -> Map k a -- ^ Map @m1@ + -> Map k b -- ^ Map @m2@ + -> f (Map k c) +mergeA g1 WhenMissing{missingSubtree = g2} (WhenMatched f) = go + where + go t1 Tip = missingSubtree g1 t1 + go Tip t2 = g2 t2 + go (Bin _ kx x1 l1 r1) t2 = case splitLookup kx t2 of + (l2, mx2, r2) -> case mx2 of + Nothing -> (\l' mx' r' -> maybe link2 (link kx) mx' l' r') + <$> l1l2 <*> missingKey g1 kx x1 <*> r1r2 + Just x2 -> (\l' mx' r' -> maybe link2 (link kx) mx' l' r') + <$> l1l2 <*> f kx x1 x2 <*> r1r2 + where + !l1l2 = go l1 l2 + !r1r2 = go r1 r2 +{-# INLINE mergeA #-} + + +{-------------------------------------------------------------------- + MergeWithKey +--------------------------------------------------------------------} + +-- | /O(n+m)/. An unsafe general combining function. +-- +-- WARNING: This function can produce corrupt maps and its results +-- may depend on the internal structures of its inputs. Users should +-- prefer 'merge' or 'mergeA'. +-- +-- When 'mergeWithKey' is given three arguments, it is inlined to the call +-- site. You should therefore use 'mergeWithKey' only to define custom +-- combining functions. For example, you could define 'unionWithKey', +-- 'differenceWithKey' and 'intersectionWithKey' as +-- +-- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2 +-- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2 +-- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2 +-- +-- When calling @'mergeWithKey' combine only1 only2@, a function combining two +-- 'Map's is created, such that +-- +-- * if a key is present in both maps, it is passed with both corresponding +-- values to the @combine@ function. Depending on the result, the key is either +-- present in the result with specified value, or is left out; +-- +-- * a nonempty subtree present only in the first map is passed to @only1@ and +-- the output is added to the result; +-- +-- * a nonempty subtree present only in the second map is passed to @only2@ and +-- the output is added to the result. +-- +-- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/. +-- The values can be modified arbitrarily. Most common variants of @only1@ and +-- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@, +-- @'filterWithKey' f@, or @'mapMaybeWithKey' f@ could be used for any @f@. + +mergeWithKey :: Ord k + => (k -> a -> b -> Maybe c) + -> (Map k a -> Map k c) + -> (Map k b -> Map k c) + -> Map k a -> Map k b -> Map k c +mergeWithKey f g1 g2 = go + where + go Tip t2 = g2 t2 + go t1 Tip = g1 t1 + go (Bin _ kx x l1 r1) t2 = + case found of + Nothing -> case g1 (singleton kx x) of + Tip -> link2 l' r' + (Bin _ _ x' Tip Tip) -> link kx x' l' r' + _ -> error "mergeWithKey: Given function only1 does not fulfill required conditions (see documentation)" + Just x2 -> case f kx x x2 of + Nothing -> link2 l' r' + Just x' -> link kx x' l' r' + where + (l2, found, r2) = splitLookup kx t2 + l' = go l1 l2 + r' = go r1 r2 +{-# INLINE mergeWithKey #-} + +{-------------------------------------------------------------------- + Submap +--------------------------------------------------------------------} +-- | /O(m*log(n/m + 1)), m <= n/. +-- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). +-- +isSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool +isSubmapOf m1 m2 = isSubmapOfBy (==) m1 m2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE isSubmapOf #-} +#endif + +{- | /O(m*log(n/m + 1)), m <= n/. + 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)]) + + +-} +isSubmapOfBy :: Ord k => (a->b->Bool) -> Map k a -> Map k b -> Bool +isSubmapOfBy f t1 t2 + = (size t1 <= size t2) && (submap' f t1 t2) +#if __GLASGOW_HASKELL__ +{-# INLINABLE isSubmapOfBy #-} +#endif + +submap' :: Ord a => (b -> c -> Bool) -> Map a b -> Map a c -> Bool +submap' _ Tip _ = True +submap' _ _ Tip = False +submap' f (Bin _ kx x l r) t + = case found of + Nothing -> False + Just y -> f x y && submap' f l lt && submap' f r gt + where + (lt,found,gt) = splitLookup kx t +#if __GLASGOW_HASKELL__ +{-# INLINABLE submap' #-} +#endif + +-- | /O(m*log(n/m + 1)), m <= n/. Is this a proper submap? (ie. a submap but not equal). +-- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@). +isProperSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool +isProperSubmapOf m1 m2 + = isProperSubmapOfBy (==) m1 m2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE isProperSubmapOf #-} +#endif + +{- | /O(m*log(n/m + 1)), m <= n/. Is this a proper submap? (ie. a submap but not equal). + The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when + @m1@ and @m2@ are not equal, + all keys in @m1@ are in @m2@, and when @f@ returns 'True' when + applied to their respective values. For example, the following + expressions are all 'True': + + > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + + But the following are all 'False': + + > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) + > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)]) + > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + + +-} +isProperSubmapOfBy :: Ord k => (a -> b -> Bool) -> Map k a -> Map k b -> Bool +isProperSubmapOfBy f t1 t2 + = (size t1 < size t2) && (submap' f t1 t2) +#if __GLASGOW_HASKELL__ +{-# INLINABLE isProperSubmapOfBy #-} +#endif + +{-------------------------------------------------------------------- + Filter and partition +--------------------------------------------------------------------} +-- | /O(n)/. 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 + +filter :: (a -> Bool) -> Map k a -> Map k a +filter p m + = filterWithKey (\_ x -> p x) m + +-- | /O(n)/. Filter all keys\/values that satisfy the predicate. +-- +-- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +filterWithKey :: (k -> a -> Bool) -> Map k a -> Map k a +filterWithKey _ Tip = Tip +filterWithKey p t@(Bin _ kx x l r) + | p kx x = if pl `ptrEq` l && pr `ptrEq` r + then t + else link kx x pl pr + | otherwise = link2 pl pr + where !pl = filterWithKey p l + !pr = filterWithKey p r + +-- | /O(n)/. Filter keys and values using an 'Applicative' +-- predicate. +filterWithKeyA :: Applicative f => (k -> a -> f Bool) -> Map k a -> f (Map k a) +filterWithKeyA _ Tip = pure Tip +filterWithKeyA p t@(Bin _ kx x l r) = + combine <$> p kx x <*> filterWithKeyA p l <*> filterWithKeyA p r + where + combine True pl pr + | pl `ptrEq` l && pr `ptrEq` r = t + | otherwise = link kx x pl pr + combine False pl pr = link2 pl pr + +-- | /O(log n)/. Take while a predicate on the keys holds. +-- The user is responsible for ensuring that for all keys @j@ and @k@ in the map, +-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'. +-- +-- @ +-- takeWhileAntitone p = 'fromDistinctAscList' . 'Data.List.takeWhile' (p . fst) . 'toList' +-- takeWhileAntitone p = 'filterWithKey' (\k _ -> p k) +-- @ + +takeWhileAntitone :: (k -> Bool) -> Map k a -> Map k a +takeWhileAntitone _ Tip = Tip +takeWhileAntitone p (Bin _ kx x l r) + | p kx = link kx x l (takeWhileAntitone p r) + | otherwise = takeWhileAntitone p l + +-- | /O(log n)/. Drop while a predicate on the keys holds. +-- The user is responsible for ensuring that for all keys @j@ and @k@ in the map, +-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'. +-- +-- @ +-- dropWhileAntitone p = 'fromDistinctAscList' . 'Data.List.dropWhile' (p . fst) . 'toList' +-- dropWhileAntitone p = 'filterWithKey' (\k -> not (p k)) +-- @ + +dropWhileAntitone :: (k -> Bool) -> Map k a -> Map k a +dropWhileAntitone _ Tip = Tip +dropWhileAntitone p (Bin _ kx x l r) + | p kx = dropWhileAntitone p r + | otherwise = link kx x (dropWhileAntitone p l) r + +-- | /O(log n)/. Divide a map at the point where a predicate on the keys stops holding. +-- The user is responsible for ensuring that for all keys @j@ and @k@ in the map, +-- @j \< k ==\> p j \>= p k@. +-- +-- @ +-- spanAntitone p xs = ('takeWhileAntitone' p xs, 'dropWhileAntitone' p xs) +-- spanAntitone p xs = partition p xs +-- @ +-- +-- Note: if @p@ is not actually antitone, then @spanAntitone@ will split the map +-- at some /unspecified/ point where the predicate switches from holding to not +-- holding (where the predicate is seen to hold before the first key and to fail +-- after the last key). + +spanAntitone :: (k -> Bool) -> Map k a -> (Map k a, Map k a) +spanAntitone p0 m = toPair (go p0 m) + where + go _ Tip = Tip :*: Tip + go p (Bin _ kx x l r) + | p kx = let u :*: v = go p r in link kx x l u :*: v + | otherwise = let u :*: v = go p l in u :*: link kx x v r + +-- | /O(n)/. 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")]) + +partition :: (a -> Bool) -> Map k a -> (Map k a,Map k a) +partition p m + = partitionWithKey (\_ x -> p x) m + +-- | /O(n)/. 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'. +-- +-- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b") +-- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty) +-- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")]) + +partitionWithKey :: (k -> a -> Bool) -> Map k a -> (Map k a,Map k a) +partitionWithKey p0 t0 = toPair $ go p0 t0 + where + go _ Tip = (Tip :*: Tip) + go p t@(Bin _ kx x l r) + | p kx x = (if l1 `ptrEq` l && r1 `ptrEq` r + then t + else link kx x l1 r1) :*: link2 l2 r2 + | otherwise = link2 l1 r1 :*: + (if l2 `ptrEq` l && r2 `ptrEq` r + then t + else link kx x l2 r2) + where + (l1 :*: l2) = go p l + (r1 :*: r2) = go p r + +-- | /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" + +mapMaybe :: (a -> Maybe b) -> Map k a -> Map k b +mapMaybe f = mapMaybeWithKey (\_ x -> f x) + +-- | /O(n)/. 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" + +mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b +mapMaybeWithKey _ Tip = Tip +mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of + Just y -> link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r) + Nothing -> link2 (mapMaybeWithKey f l) (mapMaybeWithKey f r) + +-- | /O(n)/. Traverse keys\/values and collect the 'Just' results. + +traverseMaybeWithKey :: Applicative f + => (k -> a -> f (Maybe b)) -> Map k a -> f (Map k b) +traverseMaybeWithKey = go + where + go _ Tip = pure Tip + go f (Bin _ kx x Tip Tip) = maybe Tip (\x' -> Bin 1 kx x' Tip Tip) <$> f kx x + go f (Bin _ kx x l r) = combine <$> go f l <*> f kx x <*> go f r + where + combine !l' mx !r' = case mx of + Nothing -> link2 l' r' + Just x' -> link kx x' l' r' + +-- | /O(n)/. 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")]) + +mapEither :: (a -> Either b c) -> Map k a -> (Map k b, Map k c) +mapEither f m + = mapEitherWithKey (\_ x -> f x) m + +-- | /O(n)/. 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")]) + +mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) +mapEitherWithKey f0 t0 = toPair $ go f0 t0 + where + go _ Tip = (Tip :*: Tip) + go f (Bin _ kx x l r) = case f kx x of + Left y -> link kx y l1 r1 :*: link2 l2 r2 + Right z -> link2 l1 r1 :*: link kx z l2 r2 + where + (l1 :*: l2) = go f l + (r1 :*: r2) = go f r + +{-------------------------------------------------------------------- + Mapping +--------------------------------------------------------------------} +-- | /O(n)/. Map a function over all values in the map. +-- +-- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] + +map :: (a -> b) -> Map k a -> Map k b +map f = go where + go Tip = Tip + go (Bin sx kx x l r) = Bin sx kx (f x) (go l) (go r) +-- We use a `go` function to allow `map` to inline. This makes +-- a big difference if someone uses `map (const x) m` instead +-- of `x <$ m`; it doesn't seem to do any harm. + +#ifdef __GLASGOW_HASKELL__ +{-# NOINLINE [1] map #-} +{-# RULES +"map/map" forall f g xs . map f (map g xs) = map (f . g) xs + #-} +#endif +#if __GLASGOW_HASKELL__ >= 709 +-- Safe coercions were introduced in 7.8, but did not work well with RULES yet. +{-# RULES +"map/coerce" map coerce = coerce + #-} +#endif + +-- | /O(n)/. 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")] + +mapWithKey :: (k -> a -> b) -> Map k a -> Map k b +mapWithKey _ Tip = Tip +mapWithKey f (Bin sx kx x l r) = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r) + +#ifdef __GLASGOW_HASKELL__ +{-# NOINLINE [1] mapWithKey #-} +{-# RULES +"mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) = + mapWithKey (\k a -> f k (g k a)) xs +"mapWithKey/map" forall f g xs . mapWithKey f (map g xs) = + mapWithKey (\k a -> f k (g a)) xs +"map/mapWithKey" forall f g xs . map f (mapWithKey g xs) = + mapWithKey (\k a -> f (g k a)) xs + #-} +#endif + +-- | /O(n)/. +-- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@ +-- That is, behaves exactly like a regular 'traverse' except that the traversing +-- function also has access to the key associated with a value. +-- +-- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')]) +-- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')]) == Nothing +traverseWithKey :: Applicative t => (k -> a -> t b) -> Map k a -> t (Map k b) +traverseWithKey f = go + where + go Tip = pure Tip + go (Bin 1 k v _ _) = (\v' -> Bin 1 k v' Tip Tip) <$> f k v + go (Bin s k v l r) = flip (Bin s k) <$> go l <*> f k v <*> go r +{-# INLINE traverseWithKey #-} + +-- | /O(n)/. The function 'mapAccum' threads an accumulating +-- argument through the map in ascending order of keys. +-- +-- > let f a b = (a ++ b, b ++ "X") +-- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) + +mapAccum :: (a -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccum f a m + = mapAccumWithKey (\a' _ x' -> f a' x') a m + +-- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating +-- argument through the map in ascending order of keys. +-- +-- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") +-- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) + +mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumWithKey f a t + = mapAccumL f a t + +-- | /O(n)/. The function 'mapAccumL' threads an accumulating +-- argument through the map in ascending order of keys. +mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumL _ a Tip = (a,Tip) +mapAccumL f a (Bin sx kx x l r) = + let (a1,l') = mapAccumL f a l + (a2,x') = f a1 kx x + (a3,r') = mapAccumL f a2 r + in (a3,Bin sx kx x' l' r') + +-- | /O(n)/. The function 'mapAccumR' threads an accumulating +-- argument through the map in descending order of keys. +mapAccumRWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumRWithKey _ a Tip = (a,Tip) +mapAccumRWithKey f a (Bin sx kx x l r) = + let (a1,r') = mapAccumRWithKey f a r + (a2,x') = f a1 kx x + (a3,l') = mapAccumRWithKey f a2 l + in (a3,Bin sx kx x' l' r') + +-- | /O(n*log n)/. +-- @'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 greatest of the +-- original 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" + +mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a +mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) [] +#if __GLASGOW_HASKELL__ +{-# INLINABLE mapKeys #-} +#endif + +-- | /O(n*log n)/. +-- @'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" + +mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a +mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] +#if __GLASGOW_HASKELL__ +{-# INLINABLE mapKeysWith #-} +#endif + + +-- | /O(n)/. +-- @'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")] +-- > valid (mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")])) == True +-- > valid (mapKeysMonotonic (\ _ -> 1) (fromList [(5,"a"), (3,"b")])) == False + +mapKeysMonotonic :: (k1->k2) -> Map k1 a -> Map k2 a +mapKeysMonotonic _ Tip = Tip +mapKeysMonotonic f (Bin sz k x l r) = + Bin sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r) + +{-------------------------------------------------------------------- + Folds +--------------------------------------------------------------------} + +-- | /O(n)/. Fold the values in the map using the given right-associative +-- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@. +-- +-- For example, +-- +-- > elems map = foldr (:) [] map +-- +-- > let f a len = len + (length a) +-- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4 +foldr :: (a -> b -> b) -> b -> Map k a -> b +foldr f z = go z + where + go z' Tip = z' + go z' (Bin _ _ x l r) = go (f x (go z' r)) l +{-# INLINE foldr #-} + +-- | /O(n)/. A strict version of 'foldr'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldr' :: (a -> b -> b) -> b -> Map k a -> b +foldr' f z = go z + where + go !z' Tip = z' + go z' (Bin _ _ x l r) = go (f x (go z' r)) l +{-# INLINE foldr' #-} + +-- | /O(n)/. Fold the values in the map using the given left-associative +-- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@. +-- +-- For example, +-- +-- > elems = reverse . foldl (flip (:)) [] +-- +-- > let f len a = len + (length a) +-- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4 +foldl :: (a -> b -> a) -> a -> Map k b -> a +foldl f z = go z + where + go z' Tip = z' + go z' (Bin _ _ x l r) = go (f (go z' l) x) r +{-# INLINE foldl #-} + +-- | /O(n)/. A strict version of 'foldl'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldl' :: (a -> b -> a) -> a -> Map k b -> a +foldl' f z = go z + where + go !z' Tip = z' + go z' (Bin _ _ x l r) = go (f (go z' l) x) r +{-# INLINE foldl' #-} + +-- | /O(n)/. Fold the keys and values in the map using the given right-associative +-- binary operator, such that +-- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@. +-- +-- For example, +-- +-- > keys map = foldrWithKey (\k x ks -> k:ks) [] map +-- +-- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")" +-- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)" +foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b +foldrWithKey f z = go z + where + go z' Tip = z' + go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l +{-# INLINE foldrWithKey #-} + +-- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldrWithKey' :: (k -> a -> b -> b) -> b -> Map k a -> b +foldrWithKey' f z = go z + where + go !z' Tip = z' + go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l +{-# INLINE foldrWithKey' #-} + +-- | /O(n)/. Fold the keys and values in the map using the given left-associative +-- binary operator, such that +-- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@. +-- +-- For example, +-- +-- > keys = reverse . foldlWithKey (\ks k x -> k:ks) [] +-- +-- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")" +-- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)" +foldlWithKey :: (a -> k -> b -> a) -> a -> Map k b -> a +foldlWithKey f z = go z + where + go z' Tip = z' + go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r +{-# INLINE foldlWithKey #-} + +-- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldlWithKey' :: (a -> k -> b -> a) -> a -> Map k b -> a +foldlWithKey' f z = go z + where + go !z' Tip = z' + go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r +{-# INLINE foldlWithKey' #-} + +-- | /O(n)/. Fold the keys and values in the map using the given monoid, such that +-- +-- @'foldMapWithKey' f = 'Prelude.fold' . 'mapWithKey' f@ +-- +-- This can be an asymptotically faster than 'foldrWithKey' or 'foldlWithKey' for some monoids. +foldMapWithKey :: Monoid m => (k -> a -> m) -> Map k a -> m +foldMapWithKey f = go + where + go Tip = mempty + go (Bin 1 k v _ _) = f k v + go (Bin _ k v l r) = go l `mappend` (f k v `mappend` go r) +{-# INLINE foldMapWithKey #-} + +{-------------------------------------------------------------------- + List variations +--------------------------------------------------------------------} +-- | /O(n)/. +-- Return all elements of the map in the ascending order of their keys. +-- Subject to list fusion. +-- +-- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"] +-- > elems empty == [] + +elems :: Map k a -> [a] +elems = foldr (:) [] + +-- | /O(n)/. Return all keys of the map in ascending order. Subject to list +-- fusion. +-- +-- > keys (fromList [(5,"a"), (3,"b")]) == [3,5] +-- > keys empty == [] + +keys :: Map k a -> [k] +keys = foldrWithKey (\k _ ks -> k : ks) [] + +-- | /O(n)/. An alias for 'toAscList'. Return all key\/value pairs in the map +-- in ascending key order. Subject to list fusion. +-- +-- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] +-- > assocs empty == [] + +assocs :: Map k a -> [(k,a)] +assocs m + = toAscList m + +-- | /O(n)/. The set of all keys of the map. +-- +-- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.Set.fromList [3,5] +-- > keysSet empty == Data.Set.empty + +keysSet :: Map k a -> Set.Set k +keysSet Tip = Set.Tip +keysSet (Bin sz kx _ l r) = Set.Bin sz kx (keysSet l) (keysSet r) + +-- | /O(n)/. Build a map from a set of keys and a function which for each key +-- computes its value. +-- +-- > fromSet (\k -> replicate k 'a') (Data.Set.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")] +-- > fromSet undefined Data.Set.empty == empty + +fromSet :: (k -> a) -> Set.Set k -> Map k a +fromSet _ Set.Tip = Tip +fromSet f (Set.Bin sz x l r) = Bin sz x (f x) (fromSet f l) (fromSet f r) + +{-------------------------------------------------------------------- + Lists + use [foldlStrict] to reduce demand on the control-stack +--------------------------------------------------------------------} +#if __GLASGOW_HASKELL__ >= 708 +instance (Ord k) => GHCExts.IsList (Map k v) where + type Item (Map k v) = (k,v) + fromList = fromList + toList = toList +#endif + +-- | /O(n*log n)/. 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. +-- +-- If the keys of the list are ordered, linear-time implementation is used, +-- with the performance equal to 'fromDistinctAscList'. +-- +-- > 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")] + +-- For some reason, when 'singleton' is used in fromList or in +-- create, it is not inlined, so we inline it manually. +fromList :: Ord k => [(k,a)] -> Map k a +fromList [] = Tip +fromList [(kx, x)] = Bin 1 kx x Tip Tip +fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip Tip) xs0 + | otherwise = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 + where + not_ordered _ [] = False + not_ordered kx ((ky,_) : _) = kx >= ky + {-# INLINE not_ordered #-} + + fromList' t0 xs = foldlStrict ins t0 xs + where ins t (k,x) = insert k x t + + go !_ t [] = t + go _ t [(kx, x)] = insertMax kx x t + go s l xs@((kx, x) : xss) | not_ordered kx xss = fromList' l xs + | otherwise = case create s xss of + (r, ys, []) -> go (s `shiftL` 1) (link kx x l r) ys + (r, _, ys) -> fromList' (link kx x l r) ys + + -- The create is returning a triple (tree, xs, ys). Both xs and ys + -- represent not yet processed elements and only one of them can be nonempty. + -- If ys is nonempty, the keys in ys are not ordered with respect to tree + -- and must be inserted using fromList'. Otherwise the keys have been + -- ordered so far. + create !_ [] = (Tip, [], []) + create s xs@(xp : xss) + | s == 1 = case xp of (kx, x) | not_ordered kx xss -> (Bin 1 kx x Tip Tip, [], xss) + | otherwise -> (Bin 1 kx x Tip Tip, xss, []) + | otherwise = case create (s `shiftR` 1) xs of + res@(_, [], _) -> res + (l, [(ky, y)], zs) -> (insertMax ky y l, [], zs) + (l, ys@((ky, y):yss), _) | not_ordered ky yss -> (l, [], ys) + | otherwise -> case create (s `shiftR` 1) yss of + (r, zs, ws) -> (link ky y l r, zs, ws) +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromList #-} +#endif + +-- | /O(n*log n)/. 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 + +fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a +fromListWith f xs + = fromListWithKey (\_ x y -> f x y) xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromListWith #-} +#endif + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. +-- +-- > let f k a1 a2 = (show k) ++ a1 ++ a2 +-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "3ab"), (5, "5a5ba")] +-- > fromListWithKey f [] == empty + +fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromListWithKey f xs + = foldlStrict ins empty xs + where + ins t (k,x) = insertWithKey f k x t +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromListWithKey #-} +#endif + +-- | /O(n)/. Convert the map to a list of key\/value pairs. Subject to list fusion. +-- +-- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] +-- > toList empty == [] + +toList :: Map k a -> [(k,a)] +toList = toAscList + +-- | /O(n)/. Convert the map to a list of key\/value pairs where the keys are +-- in ascending order. Subject to list fusion. +-- +-- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] + +toAscList :: Map k a -> [(k,a)] +toAscList = foldrWithKey (\k x xs -> (k,x):xs) [] + +-- | /O(n)/. Convert the map to a list of key\/value pairs where the keys +-- are in descending order. Subject to list fusion. +-- +-- > toDescList (fromList [(5,"a"), (3,"b")]) == [(5,"a"), (3,"b")] + +toDescList :: Map k a -> [(k,a)] +toDescList = foldlWithKey (\xs k x -> (k,x):xs) [] + +-- List fusion for the list generating functions. +#if __GLASGOW_HASKELL__ +-- The foldrFB and foldlFB are fold{r,l}WithKey equivalents, used for list fusion. +-- They are important to convert unfused methods back, see mapFB in prelude. +foldrFB :: (k -> a -> b -> b) -> b -> Map k a -> b +foldrFB = foldrWithKey +{-# INLINE[0] foldrFB #-} +foldlFB :: (a -> k -> b -> a) -> a -> Map k b -> a +foldlFB = foldlWithKey +{-# INLINE[0] foldlFB #-} + +-- Inline assocs and toList, so that we need to fuse only toAscList. +{-# INLINE assocs #-} +{-# INLINE toList #-} + +-- The fusion is enabled up to phase 2 included. If it does not succeed, +-- convert in phase 1 the expanded elems,keys,to{Asc,Desc}List calls back to +-- elems,keys,to{Asc,Desc}List. In phase 0, we inline fold{lr}FB (which were +-- used in a list fusion, otherwise it would go away in phase 1), and let compiler +-- do whatever it wants with elems,keys,to{Asc,Desc}List -- it was forbidden to +-- inline it before phase 0, otherwise the fusion rules would not fire at all. +{-# NOINLINE[0] elems #-} +{-# NOINLINE[0] keys #-} +{-# NOINLINE[0] toAscList #-} +{-# NOINLINE[0] toDescList #-} +{-# RULES "Map.elems" [~1] forall m . elems m = build (\c n -> foldrFB (\_ x xs -> c x xs) n m) #-} +{-# RULES "Map.elemsBack" [1] foldrFB (\_ x xs -> x : xs) [] = elems #-} +{-# RULES "Map.keys" [~1] forall m . keys m = build (\c n -> foldrFB (\k _ xs -> c k xs) n m) #-} +{-# RULES "Map.keysBack" [1] foldrFB (\k _ xs -> k : xs) [] = keys #-} +{-# RULES "Map.toAscList" [~1] forall m . toAscList m = build (\c n -> foldrFB (\k x xs -> c (k,x) xs) n m) #-} +{-# RULES "Map.toAscListBack" [1] foldrFB (\k x xs -> (k, x) : xs) [] = toAscList #-} +{-# RULES "Map.toDescList" [~1] forall m . toDescList m = build (\c n -> foldlFB (\xs k x -> c (k,x) xs) n m) #-} +{-# RULES "Map.toDescListBack" [1] foldlFB (\xs k x -> (k, x) : xs) [] = toDescList #-} +#endif + +{-------------------------------------------------------------------- + Building trees from ascending/descending lists can be done in linear time. + + Note that if [xs] is ascending that: + fromAscList xs == fromList xs + fromAscListWith f xs == fromListWith f xs +--------------------------------------------------------------------} +-- | /O(n)/. 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")] +-- > valid (fromAscList [(3,"b"), (5,"a"), (5,"b")]) == True +-- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False + +fromAscList :: Eq k => [(k,a)] -> Map k a +fromAscList xs + = fromAscListWithKey (\_ x _ -> x) xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromAscList #-} +#endif + +-- | /O(n)/. Build a map from a descending list in linear time. +-- /The precondition (input list is descending) is not checked./ +-- +-- > fromDescList [(5,"a"), (3,"b")] == fromList [(3, "b"), (5, "a")] +-- > fromDescList [(5,"a"), (5,"b"), (3,"b")] == fromList [(3, "b"), (5, "b")] +-- > valid (fromDescList [(5,"a"), (5,"b"), (3,"b")]) == True +-- > valid (fromDescList [(5,"a"), (3,"b"), (5,"b")]) == False + +fromDescList :: Eq k => [(k,a)] -> Map k a +fromDescList xs + = fromDescListWithKey (\_ x _ -> x) xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromDescList #-} +#endif + +-- | /O(n)/. 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")] +-- > valid (fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")]) == True +-- > valid (fromAscListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False + +fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a +fromAscListWith f xs + = fromAscListWithKey (\_ x y -> f x y) xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromAscListWith #-} +#endif + +-- | /O(n)/. Build a map from a descending list in linear time with a combining function for equal keys. +-- /The precondition (input list is descending) is not checked./ +-- +-- > fromDescListWith (++) [(5,"a"), (5,"b"), (3,"b")] == fromList [(3, "b"), (5, "ba")] +-- > valid (fromDescListWith (++) [(5,"a"), (5,"b"), (3,"b")]) == True +-- > valid (fromDescListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False + +fromDescListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a +fromDescListWith f xs + = fromDescListWithKey (\_ x y -> f x y) xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromDescListWith #-} +#endif + +-- | /O(n)/. 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./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] == fromList [(3, "b"), (5, "5:b5:ba")] +-- > valid (fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")]) == True +-- > valid (fromAscListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False + +fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromAscListWithKey f xs + = fromDistinctAscList (combineEq f xs) + where + -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] + combineEq _ xs' + = case xs' of + [] -> [] + [x] -> [x] + (x:xx) -> combineEq' x xx + + combineEq' z [] = [z] + combineEq' z@(kz,zz) (x@(kx,xx):xs') + | kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs' + | otherwise = z:combineEq' x xs' +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromAscListWithKey #-} +#endif + +-- | /O(n)/. Build a map from a descending list in linear time with a +-- combining function for equal keys. +-- /The precondition (input list is descending) is not checked./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromDescListWithKey f [(5,"a"), (5,"b"), (5,"b"), (3,"b")] == fromList [(3, "b"), (5, "5:b5:ba")] +-- > valid (fromDescListWithKey f [(5,"a"), (5,"b"), (5,"b"), (3,"b")]) == True +-- > valid (fromDescListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False +fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromDescListWithKey f xs + = fromDistinctDescList (combineEq f xs) + where + -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] + combineEq _ xs' + = case xs' of + [] -> [] + [x] -> [x] + (x:xx) -> combineEq' x xx + + combineEq' z [] = [z] + combineEq' z@(kz,zz) (x@(kx,xx):xs') + | kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs' + | otherwise = z:combineEq' x xs' +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromDescListWithKey #-} +#endif + + +-- | /O(n)/. 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")] +-- > valid (fromDistinctAscList [(3,"b"), (5,"a")]) == True +-- > valid (fromDistinctAscList [(3,"b"), (5,"a"), (5,"b")]) == False + +-- For some reason, when 'singleton' is used in fromDistinctAscList or in +-- create, it is not inlined, so we inline it manually. +fromDistinctAscList :: [(k,a)] -> Map k a +fromDistinctAscList [] = Tip +fromDistinctAscList ((kx0, x0) : xs0) = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 + where + go !_ t [] = t + go s l ((kx, x) : xs) = case create s xs of + (r, ys) -> go (s `shiftL` 1) (link kx x l r) ys + + create !_ [] = (Tip, []) + create s xs@(x' : xs') + | s == 1 = case x' of (kx, x) -> (Bin 1 kx x Tip Tip, xs') + | otherwise = case create (s `shiftR` 1) xs of + res@(_, []) -> res + (l, (ky, y):ys) -> case create (s `shiftR` 1) ys of + (r, zs) -> (link ky y l r, zs) + +-- | /O(n)/. Build a map from a descending list of distinct elements in linear time. +-- /The precondition is not checked./ +-- +-- > fromDistinctDescList [(5,"a"), (3,"b")] == fromList [(3, "b"), (5, "a")] +-- > valid (fromDistinctDescList [(5,"a"), (3,"b")]) == True +-- > valid (fromDistinctDescList [(5,"a"), (5,"b"), (3,"b")]) == False + +-- For some reason, when 'singleton' is used in fromDistinctDescList or in +-- create, it is not inlined, so we inline it manually. +fromDistinctDescList :: [(k,a)] -> Map k a +fromDistinctDescList [] = Tip +fromDistinctDescList ((kx0, x0) : xs0) = go (1 :: Int) (Bin 1 kx0 x0 Tip Tip) xs0 + where + go !_ t [] = t + go s r ((kx, x) : xs) = case create s xs of + (l, ys) -> go (s `shiftL` 1) (link kx x l r) ys + + create !_ [] = (Tip, []) + create s xs@(x' : xs') + | s == 1 = case x' of (kx, x) -> (Bin 1 kx x Tip Tip, xs') + | otherwise = case create (s `shiftR` 1) xs of + res@(_, []) -> res + (r, (ky, y):ys) -> case create (s `shiftR` 1) ys of + (l, zs) -> (link ky y l r, zs) + +{- +-- Functions very similar to these were used to implement +-- hedge union, intersection, and difference algorithms that we no +-- longer use. These functions, however, seem likely to be useful +-- in their own right, so I'm leaving them here in case we end up +-- exporting them. + +{-------------------------------------------------------------------- + [filterGt b t] filter all keys >[b] from tree [t] + [filterLt b t] filter all keys <[b] from tree [t] +--------------------------------------------------------------------} +filterGt :: Ord k => k -> Map k v -> Map k v +filterGt !_ Tip = Tip +filterGt !b (Bin _ kx x l r) = + case compare b kx of LT -> link kx x (filterGt b l) r + EQ -> r + GT -> filterGt b r +#if __GLASGOW_HASKELL__ +{-# INLINABLE filterGt #-} +#endif + +filterLt :: Ord k => k -> Map k v -> Map k v +filterLt !_ Tip = Tip +filterLt !b (Bin _ kx x l r) = + case compare kx b of LT -> link kx x l (filterLt b r) + EQ -> l + GT -> filterLt b l +#if __GLASGOW_HASKELL__ +{-# INLINABLE filterLt #-} +#endif +-} + +{-------------------------------------------------------------------- + Split +--------------------------------------------------------------------} +-- | /O(log n)/. 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 :: Ord k => k -> Map k a -> (Map k a,Map k a) +split !k0 t0 = toPair $ go k0 t0 + where + go k t = + case t of + Tip -> Tip :*: Tip + Bin _ kx x l r -> case compare k kx of + LT -> let (lt :*: gt) = go k l in lt :*: link kx x gt r + GT -> let (lt :*: gt) = go k r in link kx x l lt :*: gt + EQ -> (l :*: r) +#if __GLASGOW_HASKELL__ +{-# INLINABLE split #-} +#endif + +-- | /O(log n)/. 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) +splitLookup :: Ord k => k -> Map k a -> (Map k a,Maybe a,Map k a) +splitLookup k0 m = case go k0 m of + StrictTriple l mv r -> (l, mv, r) + where + go :: Ord k => k -> Map k a -> StrictTriple (Map k a) (Maybe a) (Map k a) + go !k t = + case t of + Tip -> StrictTriple Tip Nothing Tip + Bin _ kx x l r -> case compare k kx of + LT -> let StrictTriple lt z gt = go k l + !gt' = link kx x gt r + in StrictTriple lt z gt' + GT -> let StrictTriple lt z gt = go k r + !lt' = link kx x l lt + in StrictTriple lt' z gt + EQ -> StrictTriple l (Just x) r +#if __GLASGOW_HASKELL__ +{-# INLINABLE splitLookup #-} +#endif + +-- | A variant of 'splitLookup' that indicates only whether the +-- key was present, rather than producing its value. This is used to +-- implement 'intersection' to avoid allocating unnecessary 'Just' +-- constructors. +splitMember :: Ord k => k -> Map k a -> (Map k a,Bool,Map k a) +splitMember k0 m = case go k0 m of + StrictTriple l mv r -> (l, mv, r) + where + go :: Ord k => k -> Map k a -> StrictTriple (Map k a) Bool (Map k a) + go !k t = + case t of + Tip -> StrictTriple Tip False Tip + Bin _ kx x l r -> case compare k kx of + LT -> let StrictTriple lt z gt = go k l + !gt' = link kx x gt r + in StrictTriple lt z gt' + GT -> let StrictTriple lt z gt = go k r + !lt' = link kx x l lt + in StrictTriple lt' z gt + EQ -> StrictTriple l True r +#if __GLASGOW_HASKELL__ +{-# INLINABLE splitMember #-} +#endif + +data StrictTriple a b c = StrictTriple !a !b !c + +{-------------------------------------------------------------------- + Utility functions that maintain the balance properties of the tree. + All constructors assume that all values in [l] < [k] and all values + in [r] > [k], and that [l] and [r] are valid trees. + + In order of sophistication: + [Bin sz k x l r] The type constructor. + [bin k x l r] Maintains the correct size, assumes that both [l] + and [r] are balanced with respect to each other. + [balance k x l r] Restores the balance and size. + Assumes that the original tree was balanced and + that [l] or [r] has changed by at most one element. + [link k x l r] Restores balance and size. + + Furthermore, we can construct a new tree from two trees. Both operations + assume that all values in [l] < all values in [r] and that [l] and [r] + are valid: + [glue l r] Glues [l] and [r] together. Assumes that [l] and + [r] are already balanced with respect to each other. + [link2 l r] Merges two trees and restores balance. +--------------------------------------------------------------------} + +{-------------------------------------------------------------------- + Link +--------------------------------------------------------------------} +link :: k -> a -> Map k a -> Map k a -> Map k a +link kx x Tip r = insertMin kx x r +link kx x l Tip = insertMax kx x l +link kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz) + | delta*sizeL < sizeR = balanceL kz z (link kx x l lz) rz + | delta*sizeR < sizeL = balanceR ky y ly (link kx x ry r) + | otherwise = bin kx x l r + + +-- insertMin and insertMax don't perform potentially expensive comparisons. +insertMax,insertMin :: k -> a -> Map k a -> Map k a +insertMax kx x t + = case t of + Tip -> singleton kx x + Bin _ ky y l r + -> balanceR ky y l (insertMax kx x r) + +insertMin kx x t + = case t of + Tip -> singleton kx x + Bin _ ky y l r + -> balanceL ky y (insertMin kx x l) r + +{-------------------------------------------------------------------- + [link2 l r]: merges two trees. +--------------------------------------------------------------------} +link2 :: Map k a -> Map k a -> Map k a +link2 Tip r = r +link2 l Tip = l +link2 l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) + | delta*sizeL < sizeR = balanceL ky y (link2 l ly) ry + | delta*sizeR < sizeL = balanceR kx x lx (link2 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 :: Map k a -> Map k a -> Map k a +glue Tip r = r +glue l Tip = l +glue l r + | size l > size r = let ((km,m),l') = deleteFindMax l in balanceR km m l' r + | otherwise = let ((km,m),r') = deleteFindMin r in balanceL km m l r' + + +-- | /O(log n)/. 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 + +deleteFindMin :: Map k a -> ((k,a),Map k a) +deleteFindMin t + = case t of + Bin _ k x Tip r -> ((k,x),r) + Bin _ k x l r -> let !(km,l') = deleteFindMin l + !t' = balanceR k x l' r + in (km, t') + Tip -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip) + +-- | /O(log n)/. Delete and find the maximal element. +-- +-- > deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((10,"c"), fromList [(3,"b"), (5,"a")]) +-- > deleteFindMax empty Error: can not return the maximal element of an empty map + +deleteFindMax :: Map k a -> ((k,a),Map k a) +deleteFindMax t + = case t of + Bin _ k x l Tip -> ((k,x),l) + Bin _ k x l r -> let !(km,r') = deleteFindMax r + !t' = balanceL k x l r' + in (km, t') + Tip -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip) + + +{-------------------------------------------------------------------- + [balance l x r] balances two trees with value x. + The sizes of the trees should balance after decreasing the + size of one of them. (a rotation). + + [delta] is the maximal relative difference between the sizes of + two trees, it corresponds with the [w] in Adams' paper. + [ratio] is the ratio between an outer and inner sibling of the + heavier subtree in an unbalanced setting. It determines + whether a double or single rotation should be performed + to restore balance. It is corresponds with the inverse + of $\alpha$ in Adam's article. + + Note that according to the Adam's paper: + - [delta] should be larger than 4.646 with a [ratio] of 2. + - [delta] should be larger than 3.745 with a [ratio] of 1.534. + + But the Adam's paper is erroneous: + - It can be proved that for delta=2 and delta>=5 there does + not exist any ratio that would work. + - Delta=4.5 and ratio=2 does not work. + + That leaves two reasonable variants, delta=3 and delta=4, + both with ratio=2. + + - A lower [delta] leads to a more 'perfectly' balanced tree. + - A higher [delta] performs less rebalancing. + + In the benchmarks, delta=3 is faster on insert operations, + and delta=4 has slightly better deletes. As the insert speedup + is larger, we currently use delta=3. + +--------------------------------------------------------------------} +delta,ratio :: Int +delta = 3 +ratio = 2 + +-- The balance function is equivalent to the following: +-- +-- balance :: k -> a -> Map k a -> Map k a -> Map k a +-- balance k x l r +-- | sizeL + sizeR <= 1 = Bin sizeX k x l r +-- | sizeR > delta*sizeL = rotateL k x l r +-- | sizeL > delta*sizeR = rotateR k x l r +-- | otherwise = Bin sizeX k x l r +-- where +-- sizeL = size l +-- sizeR = size r +-- sizeX = sizeL + sizeR + 1 +-- +-- rotateL :: a -> b -> Map a b -> Map a b -> Map a b +-- rotateL k x l r@(Bin _ _ _ ly ry) | size ly < ratio*size ry = singleL k x l r +-- | otherwise = doubleL k x l r +-- +-- rotateR :: a -> b -> Map a b -> Map a b -> Map a b +-- rotateR k x l@(Bin _ _ _ ly ry) r | size ry < ratio*size ly = singleR k x l r +-- | otherwise = doubleR k x l r +-- +-- singleL, singleR :: a -> b -> Map a b -> Map a b -> Map a b +-- singleL k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin k2 x2 (bin k1 x1 t1 t2) t3 +-- singleR k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3) +-- +-- doubleL, doubleR :: a -> b -> Map a b -> Map a b -> Map a b +-- 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) +-- 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) +-- +-- It is only written in such a way that every node is pattern-matched only once. + +balance :: k -> a -> Map k a -> Map k a -> Map k a +balance k x l r = case l of + Tip -> case r of + Tip -> Bin 1 k x Tip Tip + (Bin _ _ _ Tip Tip) -> Bin 2 k x Tip r + (Bin _ rk rx Tip rr@(Bin _ _ _ _ _)) -> Bin 3 rk rx (Bin 1 k x Tip Tip) rr + (Bin _ rk rx (Bin _ rlk rlx _ _) Tip) -> Bin 3 rlk rlx (Bin 1 k x Tip Tip) (Bin 1 rk rx Tip Tip) + (Bin rs rk rx rl@(Bin rls rlk rlx rll rlr) rr@(Bin rrs _ _ _ _)) + | rls < ratio*rrs -> Bin (1+rs) rk rx (Bin (1+rls) k x Tip rl) rr + | otherwise -> Bin (1+rs) rlk rlx (Bin (1+size rll) k x Tip rll) (Bin (1+rrs+size rlr) rk rx rlr rr) + + (Bin ls lk lx ll lr) -> case r of + Tip -> case (ll, lr) of + (Tip, Tip) -> Bin 2 k x l Tip + (Tip, (Bin _ lrk lrx _ _)) -> Bin 3 lrk lrx (Bin 1 lk lx Tip Tip) (Bin 1 k x Tip Tip) + ((Bin _ _ _ _ _), Tip) -> Bin 3 lk lx ll (Bin 1 k x Tip Tip) + ((Bin lls _ _ _ _), (Bin lrs lrk lrx lrl lrr)) + | lrs < ratio*lls -> Bin (1+ls) lk lx ll (Bin (1+lrs) k x lr Tip) + | otherwise -> Bin (1+ls) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+size lrr) k x lrr Tip) + (Bin rs rk rx rl rr) + | rs > delta*ls -> case (rl, rr) of + (Bin rls rlk rlx rll rlr, Bin rrs _ _ _ _) + | rls < ratio*rrs -> Bin (1+ls+rs) rk rx (Bin (1+ls+rls) k x l rl) rr + | otherwise -> Bin (1+ls+rs) rlk rlx (Bin (1+ls+size rll) k x l rll) (Bin (1+rrs+size rlr) rk rx rlr rr) + (_, _) -> error "Failure in Data.Map.balance" + | ls > delta*rs -> case (ll, lr) of + (Bin lls _ _ _ _, Bin lrs lrk lrx lrl lrr) + | lrs < ratio*lls -> Bin (1+ls+rs) lk lx ll (Bin (1+rs+lrs) k x lr r) + | otherwise -> Bin (1+ls+rs) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+rs+size lrr) k x lrr r) + (_, _) -> error "Failure in Data.Map.balance" + | otherwise -> Bin (1+ls+rs) k x l r +{-# NOINLINE balance #-} + +-- Functions balanceL and balanceR are specialised versions of balance. +-- balanceL only checks whether the left subtree is too big, +-- balanceR only checks whether the right subtree is too big. + +-- balanceL is called when left subtree might have been inserted to or when +-- right subtree might have been deleted from. +balanceL :: k -> a -> Map k a -> Map k a -> Map k a +balanceL k x l r = case r of + Tip -> case l of + Tip -> Bin 1 k x Tip Tip + (Bin _ _ _ Tip Tip) -> Bin 2 k x l Tip + (Bin _ lk lx Tip (Bin _ lrk lrx _ _)) -> Bin 3 lrk lrx (Bin 1 lk lx Tip Tip) (Bin 1 k x Tip Tip) + (Bin _ lk lx ll@(Bin _ _ _ _ _) Tip) -> Bin 3 lk lx ll (Bin 1 k x Tip Tip) + (Bin ls lk lx ll@(Bin lls _ _ _ _) lr@(Bin lrs lrk lrx lrl lrr)) + | lrs < ratio*lls -> Bin (1+ls) lk lx ll (Bin (1+lrs) k x lr Tip) + | otherwise -> Bin (1+ls) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+size lrr) k x lrr Tip) + + (Bin rs _ _ _ _) -> case l of + Tip -> Bin (1+rs) k x Tip r + + (Bin ls lk lx ll lr) + | ls > delta*rs -> case (ll, lr) of + (Bin lls _ _ _ _, Bin lrs lrk lrx lrl lrr) + | lrs < ratio*lls -> Bin (1+ls+rs) lk lx ll (Bin (1+rs+lrs) k x lr r) + | otherwise -> Bin (1+ls+rs) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+rs+size lrr) k x lrr r) + (_, _) -> error "Failure in Data.Map.balanceL" + | otherwise -> Bin (1+ls+rs) k x l r +{-# NOINLINE balanceL #-} + +-- balanceR is called when right subtree might have been inserted to or when +-- left subtree might have been deleted from. +balanceR :: k -> a -> Map k a -> Map k a -> Map k a +balanceR k x l r = case l of + Tip -> case r of + Tip -> Bin 1 k x Tip Tip + (Bin _ _ _ Tip Tip) -> Bin 2 k x Tip r + (Bin _ rk rx Tip rr@(Bin _ _ _ _ _)) -> Bin 3 rk rx (Bin 1 k x Tip Tip) rr + (Bin _ rk rx (Bin _ rlk rlx _ _) Tip) -> Bin 3 rlk rlx (Bin 1 k x Tip Tip) (Bin 1 rk rx Tip Tip) + (Bin rs rk rx rl@(Bin rls rlk rlx rll rlr) rr@(Bin rrs _ _ _ _)) + | rls < ratio*rrs -> Bin (1+rs) rk rx (Bin (1+rls) k x Tip rl) rr + | otherwise -> Bin (1+rs) rlk rlx (Bin (1+size rll) k x Tip rll) (Bin (1+rrs+size rlr) rk rx rlr rr) + + (Bin ls _ _ _ _) -> case r of + Tip -> Bin (1+ls) k x l Tip + + (Bin rs rk rx rl rr) + | rs > delta*ls -> case (rl, rr) of + (Bin rls rlk rlx rll rlr, Bin rrs _ _ _ _) + | rls < ratio*rrs -> Bin (1+ls+rs) rk rx (Bin (1+ls+rls) k x l rl) rr + | otherwise -> Bin (1+ls+rs) rlk rlx (Bin (1+ls+size rll) k x l rll) (Bin (1+rrs+size rlr) rk rx rlr rr) + (_, _) -> error "Failure in Data.Map.balanceR" + | otherwise -> Bin (1+ls+rs) k x l r +{-# NOINLINE balanceR #-} + + +{-------------------------------------------------------------------- + The bin constructor maintains the size of the tree +--------------------------------------------------------------------} +bin :: k -> a -> Map k a -> Map k a -> Map k a +bin k x l r + = Bin (size l + size r + 1) k x l r +{-# INLINE bin #-} + + +{-------------------------------------------------------------------- + Eq converts the tree to a list. In a lazy setting, this + actually seems one of the faster methods to compare two trees + and it is certainly the simplest :-) +--------------------------------------------------------------------} +instance (Eq k,Eq a) => Eq (Map k a) where + t1 == t2 = (size t1 == size t2) && (toAscList t1 == toAscList t2) + +{-------------------------------------------------------------------- + Ord +--------------------------------------------------------------------} + +instance (Ord k, Ord v) => Ord (Map k v) where + compare m1 m2 = compare (toAscList m1) (toAscList m2) + +{-------------------------------------------------------------------- + Functor +--------------------------------------------------------------------} +instance Functor (Map k) where + fmap f m = map f m +#ifdef __GLASGOW_HASKELL__ + _ <$ Tip = Tip + a <$ (Bin sx kx _ l r) = Bin sx kx a (a <$ l) (a <$ r) +#endif + +instance Traversable (Map k) where + traverse f = traverseWithKey (\_ -> f) + {-# INLINE traverse #-} + +instance Foldable.Foldable (Map k) where + fold = go + where go Tip = mempty + go (Bin 1 _ v _ _) = v + go (Bin _ _ v l r) = go l `mappend` (v `mappend` go r) + {-# INLINABLE fold #-} + foldr = foldr + {-# INLINE foldr #-} + foldl = foldl + {-# INLINE foldl #-} + foldMap f t = go t + where go Tip = mempty + go (Bin 1 _ v _ _) = f v + go (Bin _ _ v l r) = go l `mappend` (f v `mappend` go r) + {-# INLINE foldMap #-} + +#if MIN_VERSION_base(4,6,0) + foldl' = foldl' + {-# INLINE foldl' #-} + foldr' = foldr' + {-# INLINE foldr' #-} +#endif +#if MIN_VERSION_base(4,8,0) + length = size + {-# INLINE length #-} + null = null + {-# INLINE null #-} + toList = elems -- NB: Foldable.toList /= Map.toList + {-# INLINE toList #-} + elem = go + where go !_ Tip = False + go x (Bin _ _ v l r) = x == v || go x l || go x r + {-# INLINABLE elem #-} + maximum = start + where start Tip = error "Map.Foldable.maximum: called with empty map" + start (Bin _ _ v l r) = go (go v l) r + + go !m Tip = m + go m (Bin _ _ v l r) = go (go (max m v) l) r + {-# INLINABLE maximum #-} + minimum = start + where start Tip = error "Map.Foldable.minumum: called with empty map" + start (Bin _ _ v l r) = go (go v l) r + + go !m Tip = m + go m (Bin _ _ v l r) = go (go (min m v) l) r + {-# INLINABLE minimum #-} + sum = foldl' (+) 0 + {-# INLINABLE sum #-} + product = foldl' (*) 1 + {-# INLINABLE product #-} +#endif + +instance (NFData k, NFData a) => NFData (Map k a) where + rnf Tip = () + rnf (Bin _ kx x l r) = rnf kx `seq` rnf x `seq` rnf l `seq` rnf r + +{-------------------------------------------------------------------- + Read +--------------------------------------------------------------------} +instance (Ord k, Read k, Read e) => Read (Map k e) where +#ifdef __GLASGOW_HASKELL__ + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- readPrec + return (fromList xs) + + readListPrec = readListPrecDefault +#else + readsPrec p = readParen (p > 10) $ \ r -> do + ("fromList",s) <- lex r + (xs,t) <- reads s + return (fromList xs,t) +#endif + +{-------------------------------------------------------------------- + Show +--------------------------------------------------------------------} +instance (Show k, Show a) => Show (Map k a) where + showsPrec d m = showParen (d > 10) $ + showString "fromList " . shows (toList m) + +-- | /O(n)/. Show the tree that implements the map. The tree is shown +-- in a compressed, hanging format. See 'showTreeWith'. +{-# DEPRECATED showTree "This function is being removed from the public API." #-} +showTree :: (Show k,Show a) => Map k a -> String +showTree m + = showTreeWith showElem True False m + where + showElem k x = show k ++ ":=" ++ show x + + +{- | /O(n)/. The expression (@'showTreeWith' showelem hang wide map@) shows + the tree that implements the map. Elements are shown using the @showElem@ function. If @hang@ is + 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If + @wide@ is 'True', an extra wide version is shown. + +> Map> let t = fromDistinctAscList [(x,()) | x <- [1..5]] +> Map> putStrLn $ showTreeWith (\k x -> show (k,x)) True False t +> (4,()) +> +--(2,()) +> | +--(1,()) +> | +--(3,()) +> +--(5,()) +> +> Map> putStrLn $ showTreeWith (\k x -> show (k,x)) True True t +> (4,()) +> | +> +--(2,()) +> | | +> | +--(1,()) +> | | +> | +--(3,()) +> | +> +--(5,()) +> +> Map> putStrLn $ showTreeWith (\k x -> show (k,x)) False True t +> +--(5,()) +> | +> (4,()) +> | +> | +--(3,()) +> | | +> +--(2,()) +> | +> +--(1,()) + +-} +{-# DEPRECATED showTreeWith "This function is being removed from the public API." #-} +showTreeWith :: (k -> a -> String) -> Bool -> Bool -> Map k a -> String +showTreeWith showelem hang wide t + | hang = (showsTreeHang showelem wide [] t) "" + | otherwise = (showsTree showelem wide [] [] t) "" + +showsTree :: (k -> a -> String) -> Bool -> [String] -> [String] -> Map k a -> ShowS +showsTree showelem wide lbars rbars t + = case t of + Tip -> showsBars lbars . showString "|\n" + Bin _ kx x Tip Tip + -> showsBars lbars . showString (showelem kx x) . showString "\n" + Bin _ kx x l r + -> showsTree showelem wide (withBar rbars) (withEmpty rbars) r . + showWide wide rbars . + showsBars lbars . showString (showelem kx x) . showString "\n" . + showWide wide lbars . + showsTree showelem wide (withEmpty lbars) (withBar lbars) l + +showsTreeHang :: (k -> a -> String) -> Bool -> [String] -> Map k a -> ShowS +showsTreeHang showelem wide bars t + = case t of + Tip -> showsBars bars . showString "|\n" + Bin _ kx x Tip Tip + -> showsBars bars . showString (showelem kx x) . showString "\n" + Bin _ kx x l r + -> showsBars bars . showString (showelem kx x) . showString "\n" . + showWide wide bars . + showsTreeHang showelem wide (withBar bars) l . + showWide wide bars . + showsTreeHang showelem wide (withEmpty bars) r + +showWide :: Bool -> [String] -> String -> String +showWide wide bars + | wide = showString (concat (reverse bars)) . showString "|\n" + | otherwise = id + +showsBars :: [String] -> ShowS +showsBars bars + = case bars of + [] -> id + _ -> showString (concat (reverse (tail bars))) . showString node + +node :: String +node = "+--" + +withBar, withEmpty :: [String] -> [String] +withBar bars = "| ":bars +withEmpty bars = " ":bars + +{-------------------------------------------------------------------- + Typeable +--------------------------------------------------------------------} + +INSTANCE_TYPEABLE2(Map) {-------------------------------------------------------------------- Assertions diff --git a/Data/Map/Lazy.hs b/Data/Map/Lazy.hs --- a/Data/Map/Lazy.hs +++ b/Data/Map/Lazy.hs @@ -1,5 +1,5 @@ {-# LANGUAGE CPP #-} -#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +#if __GLASGOW_HASKELL__ >= 703 {-# LANGUAGE Safe #-} #endif @@ -35,11 +35,17 @@ -- * Stephen Adams, \"/Efficient sets: a balancing act/\", -- Journal of Functional Programming 3(4):553-562, October 1993, -- . --- -- * J. Nievergelt and E.M. Reingold, -- \"/Binary search trees of bounded balance/\", -- SIAM journal of computing 2(1), March 1973. -- +-- Bounds for 'union', 'intersection', and 'difference' are as given +-- by +-- +-- * Guy Blelloch, Daniel Ferizovic, and Yihan Sun, +-- \"/Just Join for Parallel Ordered Sets/\", +-- . +-- -- Note that the implementation is /left-biased/ -- the elements of a -- first argument are always preferred to the second, for example in -- 'union' or 'insert'. @@ -57,11 +63,7 @@ -- $strictness -- * Map type -#if !defined(TESTING) Map -- instance Eq,Show,Read -#else - Map(..) -- instance Eq,Show,Read -#endif -- * Operators , (!), (\\) @@ -96,6 +98,7 @@ , updateWithKey , updateLookupWithKey , alter + , alterF -- * Combine @@ -116,7 +119,11 @@ , intersectionWith , intersectionWithKey - -- ** Universal combining function + -- ** General combining functions + -- | See "Data.Map.Lazy.Merge" + + -- ** Unsafe general combining function + , mergeWithKey -- * Traversal @@ -124,6 +131,7 @@ , M.map , mapWithKey , traverseWithKey + , traverseMaybeWithKey , mapAccum , mapAccumWithKey , mapAccumRWithKey @@ -164,12 +172,21 @@ , fromAscListWith , fromAscListWithKey , fromDistinctAscList + , fromDescList + , fromDescListWith + , fromDescListWithKey + , fromDistinctDescList -- * Filter , M.filter , filterWithKey + , restrictKeys + , withoutKeys , partition , partitionWithKey + , takeWhileAntitone + , dropWhileAntitone + , spanAntitone , mapMaybe , mapMaybeWithKey @@ -190,6 +207,9 @@ , elemAt , updateAt , deleteAt + , take + , drop + , splitAt -- * Min\/Max , findMin @@ -211,18 +231,10 @@ , showTree , showTreeWith , valid - -#if defined(TESTING) - -- * Internals - , bin - , balanced - , link - , merge -#endif - ) where import Data.Map.Base as M +import Prelude () -- $strictness -- diff --git a/Data/Map/Lazy/Merge.hs b/Data/Map/Lazy/Merge.hs new file mode 100644 --- /dev/null +++ b/Data/Map/Lazy/Merge.hs @@ -0,0 +1,103 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Safe #-} +#endif +#if __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE RoleAnnotations #-} +{-# LANGUAGE TypeFamilies #-} +#define USE_MAGIC_PROXY 1 +#endif + +#if USE_MAGIC_PROXY +{-# LANGUAGE MagicHash #-} +#endif + +#include "containers.h" + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Map.Lazy.Merge +-- Copyright : (c) David Feuer 2016 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- This module defines an API for writing functions that merge two +-- maps. The key functions are 'merge' and 'mergeA'. +-- Each of these can be used with several different "merge tactics". +-- +-- The 'merge' and 'mergeA' functions are shared by +-- the lazy and strict modules. Only the choice of merge tactics +-- determines strictness. If you use 'Data.Map.Strict.Merge.mapMissing' +-- from "Data.Map.Strict.Merge" then the results will be forced before +-- they are inserted. If you use 'Data.Map.Lazy.Merge.mapMissing' from +-- this module then they will not. +-- +-- == Efficiency note +-- +-- The 'Category', 'Applicative', and 'Monad' instances for 'WhenMissing' +-- tactics are included because they are valid. However, they are +-- inefficient in many cases and should usually be avoided. The instances +-- for 'WhenMatched' tactics should not pose any major efficiency problems. + +module Data.Map.Lazy.Merge ( + -- ** Simple merge tactic types + SimpleWhenMissing + , SimpleWhenMatched + + -- ** General combining function + , merge + + -- *** @WhenMatched@ tactics + , zipWithMaybeMatched + , zipWithMatched + + -- *** @WhenMissing@ tactics + , mapMaybeMissing + , dropMissing + , preserveMissing + , mapMissing + , filterMissing + + -- ** Applicative merge tactic types + , WhenMissing + , WhenMatched + + -- ** Applicative general combining function + , mergeA + + -- *** @WhenMatched@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , zipWithMaybeAMatched + , zipWithAMatched + + -- *** @WhenMissing@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , traverseMaybeMissing + , traverseMissing + , filterAMissing + + -- *** Covariant maps for tactics + , mapWhenMissing + , mapWhenMatched + + -- *** Contravariant maps for tactics + , lmapWhenMissing + , contramapFirstWhenMatched + , contramapSecondWhenMatched + + -- *** Miscellaneous tactic functions + , runWhenMatched + , runWhenMissing + ) where + +import Data.Map.Base diff --git a/Data/Map/Strict.hs b/Data/Map/Strict.hs --- a/Data/Map/Strict.hs +++ b/Data/Map/Strict.hs @@ -1,1209 +1,264 @@ {-# LANGUAGE CPP #-} -#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 -{-# LANGUAGE Trustworthy #-} -#endif - -#include "containers.h" - ------------------------------------------------------------------------------ --- | --- Module : Data.Map.Strict --- Copyright : (c) Daan Leijen 2002 --- (c) Andriy Palamarchuk 2008 --- License : BSD-style --- Maintainer : libraries@haskell.org --- Stability : provisional --- Portability : portable --- --- An efficient implementation of ordered maps from keys to values --- (dictionaries). --- --- API of this module is strict in both the keys and the values. --- If you need value-lazy maps, use "Data.Map.Lazy" instead. --- The 'Map' type is shared between the lazy and strict modules, --- meaning that the same 'Map' value can be passed to functions in --- both modules (although that is rarely needed). --- --- These modules are intended to be imported qualified, to avoid name --- clashes with Prelude functions, e.g. --- --- > import qualified Data.Map.Strict as Map --- --- The implementation of 'Map' is based on /size balanced/ binary trees (or --- trees of /bounded balance/) as described by: --- --- * Stephen Adams, \"/Efficient sets: a balancing act/\", --- Journal of Functional Programming 3(4):553-562, October 1993, --- . --- --- * J. Nievergelt and E.M. Reingold, --- \"/Binary search trees of bounded balance/\", --- SIAM journal of computing 2(1), March 1973. --- --- Note that the implementation is /left-biased/ -- the elements of a --- first argument are always preferred to the second, for example in --- 'union' or 'insert'. --- --- /Warning/: The size of the map must not exceed @maxBound::Int@. Violation of --- this condition is not detected and if the size limit is exceeded, its --- behaviour is undefined. --- --- Operation comments contain the operation time complexity in --- the Big-O notation (). --- --- Be aware that the 'Functor', 'Traversable' and 'Data' instances --- are the same as for the "Data.Map.Lazy" module, so if they are used --- on strict maps, the resulting maps will be lazy. ------------------------------------------------------------------------------ - --- See the notes at the beginning of Data.Map.Base. - -module Data.Map.Strict - ( - -- * Strictness properties - -- $strictness - - -- * Map type -#if !defined(TESTING) - Map -- instance Eq,Show,Read -#else - Map(..) -- instance Eq,Show,Read -#endif - - -- * Operators - , (!), (\\) - - -- * Query - , null - , size - , member - , notMember - , lookup - , findWithDefault - , lookupLT - , lookupGT - , lookupLE - , lookupGE - - -- * Construction - , empty - , singleton - - -- ** Insertion - , insert - , insertWith - , insertWithKey - , insertLookupWithKey - - -- ** Delete\/Update - , delete - , adjust - , adjustWithKey - , update - , updateWithKey - , updateLookupWithKey - , alter - - -- * Combine - - -- ** Union - , union - , unionWith - , unionWithKey - , unions - , unionsWith - - -- ** Difference - , difference - , differenceWith - , differenceWithKey - - -- ** Intersection - , intersection - , intersectionWith - , intersectionWithKey - - -- ** Universal combining function - , mergeWithKey - - -- * Traversal - -- ** Map - , map - , mapWithKey - , traverseWithKey - , mapAccum - , mapAccumWithKey - , mapAccumRWithKey - , mapKeys - , mapKeysWith - , mapKeysMonotonic - - -- * Folds - , foldr - , foldl - , foldrWithKey - , foldlWithKey - , foldMapWithKey - - -- ** Strict folds - , foldr' - , foldl' - , foldrWithKey' - , foldlWithKey' - - -- * Conversion - , elems - , keys - , assocs - , keysSet - , fromSet - - -- ** Lists - , toList - , fromList - , fromListWith - , fromListWithKey - - -- ** Ordered lists - , toAscList - , toDescList - , fromAscList - , fromAscListWith - , fromAscListWithKey - , fromDistinctAscList - - -- * Filter - , filter - , filterWithKey - , partition - , partitionWithKey - - , mapMaybe - , mapMaybeWithKey - , mapEither - , mapEitherWithKey - - , split - , splitLookup - , splitRoot - - -- * Submap - , isSubmapOf, isSubmapOfBy - , isProperSubmapOf, isProperSubmapOfBy - - -- * Indexed - , lookupIndex - , findIndex - , elemAt - , updateAt - , deleteAt - - -- * Min\/Max - , findMin - , findMax - , deleteMin - , deleteMax - , deleteFindMin - , deleteFindMax - , updateMin - , updateMax - , updateMinWithKey - , updateMaxWithKey - , minView - , maxView - , minViewWithKey - , maxViewWithKey - - -- * Debugging - , showTree - , showTreeWith - , valid - -#if defined(TESTING) - -- * Internals - , bin - , balanced - , link - , merge -#endif - ) where - -import Prelude hiding (lookup,map,filter,foldr,foldl,null) - -import Data.Map.Base hiding - ( findWithDefault - , singleton - , insert - , insertWith - , insertWithKey - , insertLookupWithKey - , adjust - , adjustWithKey - , update - , updateWithKey - , updateLookupWithKey - , alter - , unionWith - , unionWithKey - , unionsWith - , differenceWith - , differenceWithKey - , intersectionWith - , intersectionWithKey - , mergeWithKey - , map - , mapWithKey - , mapAccum - , mapAccumWithKey - , mapAccumRWithKey - , mapKeysWith - , fromSet - , fromList - , fromListWith - , fromListWithKey - , fromAscList - , fromAscListWith - , fromAscListWithKey - , fromDistinctAscList - , mapMaybe - , mapMaybeWithKey - , mapEither - , mapEitherWithKey - , updateAt - , updateMin - , updateMax - , updateMinWithKey - , updateMaxWithKey - ) -import qualified Data.Set.Base as Set -import Data.Utils.StrictFold -import Data.Utils.StrictPair - -import Data.Bits (shiftL, shiftR) -#if __GLASGOW_HASKELL__ >= 709 -import Data.Coerce -#endif - - --- $strictness --- --- This module satisfies the following strictness properties: --- --- 1. Key arguments are evaluated to WHNF; --- --- 2. Keys and values are evaluated to WHNF before they are stored in --- the map. --- --- Here's an example illustrating the first property: --- --- > delete undefined m == undefined --- --- Here are some examples that illustrate the second property: --- --- > map (\ v -> undefined) m == undefined -- m is not empty --- > mapKeys (\ k -> undefined) m == undefined -- m is not empty - -{-------------------------------------------------------------------- - Query ---------------------------------------------------------------------} - --- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns --- the value at key @k@ or returns default value @def@ --- when the key is not in the map. --- --- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' --- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' - --- See Map.Base.Note: Local 'go' functions and capturing -findWithDefault :: Ord k => a -> k -> Map k a -> a -findWithDefault def k = k `seq` go - where - go Tip = def - go (Bin _ kx x l r) = case compare k kx of - LT -> go l - GT -> go r - EQ -> x -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE findWithDefault #-} -#else -{-# INLINE findWithDefault #-} -#endif - -{-------------------------------------------------------------------- - Construction ---------------------------------------------------------------------} - --- | /O(1)/. A map with a single element. --- --- > singleton 1 'a' == fromList [(1, 'a')] --- > size (singleton 1 'a') == 1 - -singleton :: k -> a -> Map k a -singleton k x = x `seq` Bin 1 k x Tip Tip -{-# INLINE singleton #-} - -{-------------------------------------------------------------------- - Insertion ---------------------------------------------------------------------} --- | /O(log n)/. 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' - --- See Map.Base.Note: Type of local 'go' function -insert :: Ord k => k -> a -> Map k a -> Map k a -insert = go - where - go :: Ord k => k -> a -> Map k a -> Map k a - STRICT_1_OF_3(go) - STRICT_2_OF_3(go) - go kx x Tip = singleton kx x - go kx x (Bin sz ky y l r) = - case compare kx ky of - LT -> balanceL ky y (go kx x l) r - GT -> balanceR ky y l (go kx x r) - EQ -> Bin sz kx x l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insert #-} -#else -{-# INLINE insert #-} -#endif - --- | /O(log n)/. 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" - -insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWith f = insertWithKey (\_ x' y' -> f x' y') -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertWith #-} -#else -{-# INLINE insertWith #-} -#endif - --- | /O(log n)/. 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" - --- See Map.Base.Note: Type of local 'go' function -insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWithKey = go - where - go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a - STRICT_2_OF_4(go) - go _ kx x Tip = singleton kx x - go f kx x (Bin sy ky y l r) = - case compare kx ky of - LT -> balanceL ky y (go f kx x l) r - GT -> balanceR ky y l (go f kx x r) - EQ -> let x' = f kx x y - in x' `seq` Bin sy kx x' l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertWithKey #-} -#else -{-# INLINE insertWithKey #-} -#endif - --- | /O(log n)/. 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") --- --- This is how to define @insertLookup@ using @insertLookupWithKey@: --- --- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t --- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) --- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) - --- See Map.Base.Note: Type of local 'go' function -insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a - -> (Maybe a, Map k a) -insertLookupWithKey f0 kx0 x0 t0 = toPair $ go f0 kx0 x0 t0 - where - go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> StrictPair (Maybe a) (Map k a) - STRICT_2_OF_4(go) - go _ kx x Tip = Nothing :*: singleton kx x - go f kx x (Bin sy ky y l r) = - case compare kx ky of - LT -> let (found :*: l') = go f kx x l - in found :*: balanceL ky y l' r - GT -> let (found :*: r') = go f kx x r - in found :*: balanceR ky y l r' - EQ -> let x' = f kx x y - in x' `seq` (Just y :*: Bin sy kx x' l r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertLookupWithKey #-} -#else -{-# INLINE insertLookupWithKey #-} -#endif - -{-------------------------------------------------------------------- - Deletion ---------------------------------------------------------------------} - --- | /O(log n)/. 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 - -adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a -adjust f = adjustWithKey (\_ x -> f x) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE adjust #-} -#else -{-# INLINE adjust #-} -#endif - --- | /O(log n)/. 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 - -adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a -adjustWithKey f = updateWithKey (\k' x' -> Just (f k' x')) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE adjustWithKey #-} -#else -{-# INLINE adjustWithKey #-} -#endif - --- | /O(log n)/. 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 :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a -update f = updateWithKey (\_ x -> f x) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE update #-} -#else -{-# INLINE update #-} -#endif - --- | /O(log n)/. 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" - --- See Map.Base.Note: Type of local 'go' function -updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a -updateWithKey = go - where - go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a - STRICT_2_OF_3(go) - go _ _ Tip = Tip - go f k(Bin sx kx x l r) = - case compare k kx of - LT -> balanceR kx x (go f k l) r - GT -> balanceL kx x l (go f k r) - EQ -> case f kx x of - Just x' -> x' `seq` Bin sx kx x' l r - Nothing -> glue l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE updateWithKey #-} -#else -{-# INLINE updateWithKey #-} -#endif - --- | /O(log n)/. Lookup and update. See also 'updateWithKey'. --- The function returns changed value, if it is updated. --- Returns the original key value if the map entry is deleted. --- --- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing --- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "5:new a", fromList [(3, "b"), (5, "5:new a")]) --- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) --- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") - --- See Map.Base.Note: Type of local 'go' function -updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) -updateLookupWithKey f0 k0 t0 = toPair $ go f0 k0 t0 - where - go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> StrictPair (Maybe a) (Map k a) - STRICT_2_OF_3(go) - go _ _ Tip = (Nothing :*: Tip) - go f k (Bin sx kx x l r) = - case compare k kx of - LT -> let (found :*: l') = go f k l - in found :*: balanceR kx x l' r - GT -> let (found :*: r') = go f k r - in found :*: balanceL kx x l r' - EQ -> case f kx x of - Just x' -> x' `seq` (Just x' :*: Bin sx kx x' l r) - Nothing -> (Just x :*: glue l r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE updateLookupWithKey #-} -#else -{-# INLINE updateLookupWithKey #-} -#endif - --- | /O(log n)/. 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 'Map'. --- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. --- --- > let f _ = Nothing --- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] --- > alter f 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" --- > --- > let f _ = Just "c" --- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "c")] --- > alter f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "c")] - --- See Map.Base.Note: Type of local 'go' function -alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a -alter = go - where - go :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a - STRICT_2_OF_3(go) - go f k Tip = case f Nothing of - Nothing -> Tip - Just x -> singleton k x - - go f k (Bin sx kx x l r) = case compare k kx of - LT -> balance kx x (go f k l) r - GT -> balance kx x l (go f k r) - EQ -> case f (Just x) of - Just x' -> x' `seq` Bin sx kx x' l r - Nothing -> glue l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE alter #-} -#else -{-# INLINE alter #-} -#endif - -{-------------------------------------------------------------------- - Indexing ---------------------------------------------------------------------} - --- | /O(log n)/. 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 - -updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a -updateAt f i t = i `seq` - case t of - Tip -> error "Map.updateAt: index out of range" - Bin sx kx x l r -> case compare i sizeL of - LT -> balanceR kx x (updateAt f i l) r - GT -> balanceL kx x l (updateAt f (i-sizeL-1) r) - EQ -> case f kx x of - Just x' -> x' `seq` Bin sx kx x' l r - Nothing -> glue l r - where - sizeL = size l - -{-------------------------------------------------------------------- - Minimal, Maximal ---------------------------------------------------------------------} - --- | /O(log n)/. 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" - -updateMin :: (a -> Maybe a) -> Map k a -> Map k a -updateMin f m - = updateMinWithKey (\_ x -> f x) m - --- | /O(log n)/. 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" - -updateMax :: (a -> Maybe a) -> Map k a -> Map k a -updateMax f m - = updateMaxWithKey (\_ x -> f x) m - - --- | /O(log n)/. 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" - -updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a -updateMinWithKey _ Tip = Tip -updateMinWithKey f (Bin sx kx x Tip r) = case f kx x of - Nothing -> r - Just x' -> x' `seq` Bin sx kx x' Tip r -updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r - --- | /O(log n)/. 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" - -updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a -updateMaxWithKey _ Tip = Tip -updateMaxWithKey f (Bin sx kx x l Tip) = case f kx x of - Nothing -> l - Just x' -> x' `seq` Bin sx kx x' l Tip -updateMaxWithKey f (Bin _ kx x l r) = balanceL kx x l (updateMaxWithKey f r) - -{-------------------------------------------------------------------- - Union. ---------------------------------------------------------------------} - --- | The union of a list of maps, with a combining operation: --- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@). --- --- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] --- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] - -unionsWith :: Ord k => (a->a->a) -> [Map k a] -> Map k a -unionsWith f ts - = foldlStrict (unionWith f) empty ts -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE unionsWith #-} -#endif - -{-------------------------------------------------------------------- - Union with a combining function ---------------------------------------------------------------------} --- | /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")] - -unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a -unionWith f m1 m2 - = unionWithKey (\_ x y -> f x y) m1 m2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE unionWith #-} -#endif - --- | /O(n+m)/. --- Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm. --- --- > 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")] - -unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a -unionWithKey f t1 t2 = mergeWithKey (\k x1 x2 -> Just $ f k x1 x2) id id t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE unionWithKey #-} -#endif - -{-------------------------------------------------------------------- - Difference ---------------------------------------------------------------------} - --- | /O(n+m)/. 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" - -differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a -differenceWith f m1 m2 - = differenceWithKey (\_ x y -> f x y) m1 m2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE differenceWith #-} -#endif - --- | /O(n+m)/. 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" - -differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a -differenceWithKey f t1 t2 = mergeWithKey f id (const Tip) t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE differenceWithKey #-} -#endif - - -{-------------------------------------------------------------------- - Intersection ---------------------------------------------------------------------} - --- | /O(n+m)/. Intersection with a combining function. The implementation uses --- an efficient /hedge/ algorithm comparable with /hedge-union/. --- --- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" - -intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c -intersectionWith f m1 m2 - = intersectionWithKey (\_ x y -> f x y) m1 m2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE intersectionWith #-} -#endif - --- | /O(n+m)/. Intersection with a combining function. The implementation uses --- an efficient /hedge/ algorithm comparable with /hedge-union/. --- --- > 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" - - -intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c -intersectionWithKey f t1 t2 = mergeWithKey (\k x1 x2 -> Just $ f k x1 x2) (const Tip) (const Tip) t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE intersectionWithKey #-} -#endif - - -{-------------------------------------------------------------------- - MergeWithKey ---------------------------------------------------------------------} - --- | /O(n+m)/. A high-performance universal combining function. This function --- is used to define 'unionWith', 'unionWithKey', 'differenceWith', --- 'differenceWithKey', 'intersectionWith', 'intersectionWithKey' and can be --- used to define other custom combine functions. --- --- Please make sure you know what is going on when using 'mergeWithKey', --- otherwise you can be surprised by unexpected code growth or even --- corruption of the data structure. --- --- When 'mergeWithKey' is given three arguments, it is inlined to the call --- site. You should therefore use 'mergeWithKey' only to define your custom --- combining functions. For example, you could define 'unionWithKey', --- 'differenceWithKey' and 'intersectionWithKey' as --- --- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2 --- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2 --- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2 --- --- When calling @'mergeWithKey' combine only1 only2@, a function combining two --- 'IntMap's is created, such that --- --- * if a key is present in both maps, it is passed with both corresponding --- values to the @combine@ function. Depending on the result, the key is either --- present in the result with specified value, or is left out; --- --- * a nonempty subtree present only in the first map is passed to @only1@ and --- the output is added to the result; --- --- * a nonempty subtree present only in the second map is passed to @only2@ and --- the output is added to the result. --- --- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/. --- The values can be modified arbitrarily. Most common variants of @only1@ and --- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or --- @'filterWithKey' f@ could be used for any @f@. - -mergeWithKey :: Ord k => (k -> a -> b -> Maybe c) -> (Map k a -> Map k c) -> (Map k b -> Map k c) - -> Map k a -> Map k b -> Map k c -mergeWithKey f g1 g2 = go - where - go Tip t2 = g2 t2 - go t1 Tip = g1 t1 - go t1 t2 = hedgeMerge NothingS NothingS t1 t2 - - hedgeMerge _ _ t1 Tip = g1 t1 - hedgeMerge blo bhi Tip (Bin _ kx x l r) = g2 $ link kx x (filterGt blo l) (filterLt bhi r) - hedgeMerge blo bhi (Bin _ kx x l r) t2 = let l' = hedgeMerge blo bmi l (trim blo bmi t2) - (found, trim_t2) = trimLookupLo kx bhi t2 - r' = hedgeMerge bmi bhi r trim_t2 - in case found of - Nothing -> case g1 (singleton kx x) of - Tip -> merge l' r' - (Bin _ _ x' Tip Tip) -> link kx x' l' r' - _ -> error "mergeWithKey: Given function only1 does not fulfil required conditions (see documentation)" - Just x2 -> case f kx x x2 of - Nothing -> merge l' r' - Just x' -> x' `seq` link kx x' l' r' - where bmi = JustS kx -{-# INLINE mergeWithKey #-} - -{-------------------------------------------------------------------- - Filter and partition ---------------------------------------------------------------------} - --- | /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" - -mapMaybe :: (a -> Maybe b) -> Map k a -> Map k b -mapMaybe f = mapMaybeWithKey (\_ x -> f x) - --- | /O(n)/. 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" - -mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b -mapMaybeWithKey _ Tip = Tip -mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of - Just y -> y `seq` link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r) - Nothing -> merge (mapMaybeWithKey f l) (mapMaybeWithKey f r) - --- | /O(n)/. 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")]) - -mapEither :: (a -> Either b c) -> Map k a -> (Map k b, Map k c) -mapEither f m - = mapEitherWithKey (\_ x -> f x) m - --- | /O(n)/. 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")]) - -mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) -mapEitherWithKey f0 t0 = toPair $ go f0 t0 - where - go _ Tip = (Tip :*: Tip) - go f (Bin _ kx x l r) = case f kx x of - Left y -> y `seq` (link kx y l1 r1 :*: merge l2 r2) - Right z -> z `seq` (merge l1 r1 :*: link kx z l2 r2) - where - (l1 :*: l2) = go f l - (r1 :*: r2) = go f r - -{-------------------------------------------------------------------- - Mapping ---------------------------------------------------------------------} --- | /O(n)/. Map a function over all values in the map. --- --- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] - -map :: (a -> b) -> Map k a -> Map k b -map _ Tip = Tip -map f (Bin sx kx x l r) = let x' = f x in x' `seq` Bin sx kx x' (map f l) (map f r) -#ifdef __GLASGOW_HASKELL__ -{-# NOINLINE [1] map #-} -{-# RULES -"map/map" forall f g xs . map f (map g xs) = map (f . g) xs - #-} -#endif -#if __GLASGOW_HASKELL__ >= 709 --- Safe coercions were introduced in 7.8, but did not work well with RULES yet. -{-# RULES -"mapSeq/coerce" map coerce = coerce - #-} -#endif - --- | /O(n)/. 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")] - -mapWithKey :: (k -> a -> b) -> Map k a -> Map k b -mapWithKey _ Tip = Tip -mapWithKey f (Bin sx kx x l r) = - let x' = f kx x - in x' `seq` Bin sx kx x' (mapWithKey f l) (mapWithKey f r) - -#ifdef __GLASGOW_HASKELL__ -{-# NOINLINE [1] mapWithKey #-} -{-# RULES -"mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) = - mapWithKey (\k a -> f k (g k a)) xs -"mapWithKey/map" forall f g xs . mapWithKey f (map g xs) = - mapWithKey (\k a -> f k (g a)) xs -"map/mapWithKey" forall f g xs . map f (mapWithKey g xs) = - mapWithKey (\k a -> f (g k a)) xs - #-} -#endif - --- | /O(n)/. The function 'mapAccum' threads an accumulating --- argument through the map in ascending order of keys. --- --- > let f a b = (a ++ b, b ++ "X") --- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) - -mapAccum :: (a -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) -mapAccum f a m - = mapAccumWithKey (\a' _ x' -> f a' x') a m - --- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating --- argument through the map in ascending order of keys. --- --- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") --- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) - -mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) -mapAccumWithKey f a t - = mapAccumL f a t - --- | /O(n)/. The function 'mapAccumL' threads an accumulating --- argument through the map in ascending order of keys. -mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) -mapAccumL _ a Tip = (a,Tip) -mapAccumL f a (Bin sx kx x l r) = - let (a1,l') = mapAccumL f a l - (a2,x') = f a1 kx x - (a3,r') = mapAccumL f a2 r - in x' `seq` (a3,Bin sx kx x' l' r') - --- | /O(n)/. The function 'mapAccumR' threads an accumulating --- argument through the map in descending order of keys. -mapAccumRWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) -mapAccumRWithKey _ a Tip = (a,Tip) -mapAccumRWithKey f a (Bin sx kx x l r) = - let (a1,r') = mapAccumRWithKey f a r - (a2,x') = f a1 kx x - (a3,l') = mapAccumRWithKey f a2 l - in x' `seq` (a3,Bin sx kx x' l' r') - --- | /O(n*log n)/. --- @'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" - -mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a -mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapKeysWith #-} -#endif - -{-------------------------------------------------------------------- - Conversions ---------------------------------------------------------------------} - --- | /O(n)/. Build a map from a set of keys and a function which for each key --- computes its value. --- --- > fromSet (\k -> replicate k 'a') (Data.Set.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")] --- > fromSet undefined Data.Set.empty == empty - -fromSet :: (k -> a) -> Set.Set k -> Map k a -fromSet _ Set.Tip = Tip -fromSet f (Set.Bin sz x l r) = case f x of v -> v `seq` Bin sz x v (fromSet f l) (fromSet f r) - -{-------------------------------------------------------------------- - Lists - use [foldlStrict] to reduce demand on the control-stack ---------------------------------------------------------------------} --- | /O(n*log n)/. 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. --- --- If the keys of the list are ordered, linear-time implementation is used, --- with the performance equal to 'fromDistinctAscList'. --- --- > 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")] - --- For some reason, when 'singleton' is used in fromList or in --- create, it is not inlined, so we inline it manually. -fromList :: Ord k => [(k,a)] -> Map k a -fromList [] = Tip -fromList [(kx, x)] = x `seq` Bin 1 kx x Tip Tip -fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = x0 `seq` fromList' (Bin 1 kx0 x0 Tip Tip) xs0 - | otherwise = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 - where - not_ordered _ [] = False - not_ordered kx ((ky,_) : _) = kx >= ky - {-# INLINE not_ordered #-} - - fromList' t0 xs = foldlStrict ins t0 xs - where ins t (k,x) = insert k x t - - STRICT_1_OF_3(go) - go _ t [] = t - go _ t [(kx, x)] = x `seq` insertMax kx x t - go s l xs@((kx, x) : xss) | not_ordered kx xss = fromList' l xs - | otherwise = case create s xss of - (r, ys, []) -> x `seq` go (s `shiftL` 1) (link kx x l r) ys - (r, _, ys) -> x `seq` fromList' (link kx x l r) ys - - -- The create is returning a triple (tree, xs, ys). Both xs and ys - -- represent not yet processed elements and only one of them can be nonempty. - -- If ys is nonempty, the keys in ys are not ordered with respect to tree - -- and must be inserted using fromList'. Otherwise the keys have been - -- ordered so far. - STRICT_1_OF_2(create) - create _ [] = (Tip, [], []) - create s xs@(xp : xss) - | s == 1 = case xp of (kx, x) | not_ordered kx xss -> x `seq` (Bin 1 kx x Tip Tip, [], xss) - | otherwise -> x `seq` (Bin 1 kx x Tip Tip, xss, []) - | otherwise = case create (s `shiftR` 1) xs of - res@(_, [], _) -> res - (l, [(ky, y)], zs) -> y `seq` (insertMax ky y l, [], zs) - (l, ys@((ky, y):yss), _) | not_ordered ky yss -> (l, [], ys) - | otherwise -> case create (s `shiftR` 1) yss of - (r, zs, ws) -> y `seq` (link ky y l r, zs, ws) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromList #-} -#endif - --- | /O(n*log n)/. 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 - -fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a -fromListWith f xs - = fromListWithKey (\_ x y -> f x y) xs -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromListWith #-} -#endif - --- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. --- --- > let f k a1 a2 = (show k) ++ a1 ++ a2 --- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "3ab"), (5, "5a5ba")] --- > fromListWithKey f [] == empty - -fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a -fromListWithKey f xs - = foldlStrict ins empty xs - where - ins t (k,x) = insertWithKey f k x t -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromListWithKey #-} -#endif - -{-------------------------------------------------------------------- - Building trees from ascending/descending lists can be done in linear time. - - Note that if [xs] is ascending that: - fromAscList xs == fromList xs - fromAscListWith f xs == fromListWith f xs ---------------------------------------------------------------------} --- | /O(n)/. 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")] --- > valid (fromAscList [(3,"b"), (5,"a"), (5,"b")]) == True --- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False - -fromAscList :: Eq k => [(k,a)] -> Map k a -fromAscList xs - = fromAscListWithKey (\_ x _ -> x) xs -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromAscList #-} -#endif - --- | /O(n)/. 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")] --- > valid (fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")]) == True --- > valid (fromAscListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False - -fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a -fromAscListWith f xs - = fromAscListWithKey (\_ x y -> f x y) xs -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromAscListWith #-} -#endif - --- | /O(n)/. 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./ --- --- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 --- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] == fromList [(3, "b"), (5, "5:b5:ba")] --- > valid (fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")]) == True --- > valid (fromAscListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False - -fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a -fromAscListWithKey f xs - = fromDistinctAscList (combineEq f xs) - where - -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] - combineEq _ xs' - = case xs' of - [] -> [] - [x] -> [x] - (x:xx) -> combineEq' x xx - - combineEq' z [] = [z] - combineEq' z@(kz,zz) (x@(kx,xx):xs') - | kx==kz = let yy = f kx xx zz in yy `seq` combineEq' (kx,yy) xs' - | otherwise = z:combineEq' x xs' -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromAscListWithKey #-} -#endif - --- | /O(n)/. 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")] --- > valid (fromDistinctAscList [(3,"b"), (5,"a")]) == True --- > valid (fromDistinctAscList [(3,"b"), (5,"a"), (5,"b")]) == False - --- For some reason, when 'singleton' is used in fromDistinctAscList or in --- create, it is not inlined, so we inline it manually. -fromDistinctAscList :: [(k,a)] -> Map k a -fromDistinctAscList [] = Tip -fromDistinctAscList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 - where - STRICT_1_OF_3(go) - go _ t [] = t - go s l ((kx, x) : xs) = case create s xs of - (r, ys) -> x `seq` go (s `shiftL` 1) (link kx x l r) ys - - STRICT_1_OF_2(create) - create _ [] = (Tip, []) - create s xs@(x' : xs') - | s == 1 = case x' of (kx, x) -> x `seq` (Bin 1 kx x Tip Tip, xs') - | otherwise = case create (s `shiftR` 1) xs of - res@(_, []) -> res - (l, (ky, y):ys) -> case create (s `shiftR` 1) ys of - (r, zs) -> y `seq` (link ky y l r, zs) +{-# LANGUAGE BangPatterns #-} +#if __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Safe #-} +#endif + +#include "containers.h" + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Map.Strict +-- Copyright : (c) Daan Leijen 2002 +-- (c) Andriy Palamarchuk 2008 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- An efficient implementation of ordered maps from keys to values +-- (dictionaries). +-- +-- API of this module is strict in both the keys and the values. +-- If you need value-lazy maps, use "Data.Map.Lazy" instead. +-- The 'Map' type is shared between the lazy and strict modules, +-- meaning that the same 'Map' value can be passed to functions in +-- both modules (although that is rarely needed). +-- +-- These modules are intended to be imported qualified, to avoid name +-- clashes with Prelude functions, e.g. +-- +-- > import qualified Data.Map.Strict as Map +-- +-- The implementation of 'Map' is based on /size balanced/ binary trees (or +-- trees of /bounded balance/) as described by: +-- +-- * Stephen Adams, \"/Efficient sets: a balancing act/\", +-- Journal of Functional Programming 3(4):553-562, October 1993, +-- . +-- * J. Nievergelt and E.M. Reingold, +-- \"/Binary search trees of bounded balance/\", +-- SIAM journal of computing 2(1), March 1973. +-- +-- Bounds for 'union', 'intersection', and 'difference' are as given +-- by +-- +-- * Guy Blelloch, Daniel Ferizovic, and Yihan Sun, +-- \"/Just Join for Parallel Ordered Sets/\", +-- . +-- +-- Note that the implementation is /left-biased/ -- the elements of a +-- first argument are always preferred to the second, for example in +-- 'union' or 'insert'. +-- +-- /Warning/: The size of the map must not exceed @maxBound::Int@. Violation of +-- this condition is not detected and if the size limit is exceeded, its +-- behaviour is undefined. +-- +-- Operation comments contain the operation time complexity in +-- the Big-O notation (). +-- +-- Be aware that the 'Functor', 'Traversable' and 'Data' instances +-- are the same as for the "Data.Map.Lazy" module, so if they are used +-- on strict maps, the resulting maps will be lazy. +----------------------------------------------------------------------------- + +-- See the notes at the beginning of Data.Map.Base. + +module Data.Map.Strict + ( + -- * Strictness properties + -- $strictness + + -- * Map type + Map -- instance Eq,Show,Read + + -- * Operators + , (!), (\\) + + -- * Query + , null + , size + , member + , notMember + , lookup + , findWithDefault + , lookupLT + , lookupGT + , lookupLE + , lookupGE + + -- * Construction + , empty + , singleton + + -- ** Insertion + , insert + , insertWith + , insertWithKey + , insertLookupWithKey + + -- ** Delete\/Update + , delete + , adjust + , adjustWithKey + , update + , updateWithKey + , updateLookupWithKey + , alter + , alterF + + -- * Combine + + -- ** Union + , union + , unionWith + , unionWithKey + , unions + , unionsWith + + -- ** Difference + , difference + , differenceWith + , differenceWithKey + + -- ** Intersection + , intersection + , intersectionWith + , intersectionWithKey + + -- ** General combining functions + -- | See "Data.Map.Strict.Merge" + + -- ** Deprecated general combining function + + , mergeWithKey + + -- * Traversal + -- ** Map + , map + , mapWithKey + , traverseWithKey + , traverseMaybeWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeys + , mapKeysWith + , mapKeysMonotonic + + -- * Folds + , foldr + , foldl + , foldrWithKey + , foldlWithKey + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + , fromSet + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + , fromDescList + , fromDescListWith + , fromDescListWithKey + , fromDistinctDescList + + -- * Filter + , filter + , filterWithKey + , restrictKeys + , withoutKeys + , partition + , partitionWithKey + + , takeWhileAntitone + , dropWhileAntitone + , spanAntitone + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , updateAt + , deleteAt + , take + , drop + , splitAt + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + , valid + ) where + +import Data.Map.Strict.Internal +import Prelude () + +-- $strictness +-- +-- This module satisfies the following strictness properties: +-- +-- 1. Key arguments are evaluated to WHNF; +-- +-- 2. Keys and values are evaluated to WHNF before they are stored in +-- the map. +-- +-- Here's an example illustrating the first property: +-- +-- > delete undefined m == undefined +-- +-- Here are some examples that illustrate the second property: +-- +-- > map (\ v -> undefined) m == undefined -- m is not empty +-- > mapKeys (\ k -> undefined) m == undefined -- m is not empty diff --git a/Data/Map/Strict/Internal.hs b/Data/Map/Strict/Internal.hs new file mode 100644 --- /dev/null +++ b/Data/Map/Strict/Internal.hs @@ -0,0 +1,1716 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} +#if __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif + +{-# OPTIONS_HADDOCK hide #-} + +#include "containers.h" + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Map.Strict.Internal +-- Copyright : (c) Daan Leijen 2002 +-- (c) Andriy Palamarchuk 2008 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. +-- +-- = Description +-- +-- An efficient implementation of ordered maps from keys to values +-- (dictionaries). +-- +-- API of this module is strict in both the keys and the values. +-- If you need value-lazy maps, use "Data.Map.Lazy" instead. +-- The 'Map' type is shared between the lazy and strict modules, +-- meaning that the same 'Map' value can be passed to functions in +-- both modules (although that is rarely needed). +-- +-- These modules are intended to be imported qualified, to avoid name +-- clashes with Prelude functions, e.g. +-- +-- > import qualified Data.Map.Strict as Map +-- +-- The implementation of 'Map' is based on /size balanced/ binary trees (or +-- trees of /bounded balance/) as described by: +-- +-- * Stephen Adams, \"/Efficient sets: a balancing act/\", +-- Journal of Functional Programming 3(4):553-562, October 1993, +-- . +-- * J. Nievergelt and E.M. Reingold, +-- \"/Binary search trees of bounded balance/\", +-- SIAM journal of computing 2(1), March 1973. +-- +-- Bounds for 'union', 'intersection', and 'difference' are as given +-- by +-- +-- * Guy Blelloch, Daniel Ferizovic, and Yihan Sun, +-- \"/Just Join for Parallel Ordered Sets/\", +-- . +-- +-- Note that the implementation is /left-biased/ -- the elements of a +-- first argument are always preferred to the second, for example in +-- 'union' or 'insert'. +-- +-- /Warning/: The size of the map must not exceed @maxBound::Int@. Violation of +-- this condition is not detected and if the size limit is exceeded, its +-- behaviour is undefined. +-- +-- Operation comments contain the operation time complexity in +-- the Big-O notation (). +-- +-- Be aware that the 'Functor', 'Traversable' and 'Data' instances +-- are the same as for the "Data.Map.Lazy" module, so if they are used +-- on strict maps, the resulting maps will be lazy. +----------------------------------------------------------------------------- + +-- See the notes at the beginning of Data.Map.Base. + +module Data.Map.Strict.Internal + ( + -- * Strictness properties + -- $strictness + + -- * Map type + Map(..) -- instance Eq,Show,Read + + -- * Operators + , (!), (\\) + + -- * Query + , null + , size + , member + , notMember + , lookup + , findWithDefault + , lookupLT + , lookupGT + , lookupLE + , lookupGE + + -- * Construction + , empty + , singleton + + -- ** Insertion + , insert + , insertWith + , insertWithKey + , insertLookupWithKey + + -- ** Delete\/Update + , delete + , adjust + , adjustWithKey + , update + , updateWithKey + , updateLookupWithKey + , alter + , alterF + + -- * Combine + + -- ** Union + , union + , unionWith + , unionWithKey + , unions + , unionsWith + + -- ** Difference + , difference + , differenceWith + , differenceWithKey + + -- ** Intersection + , intersection + , intersectionWith + , intersectionWithKey + + -- ** General combining function + , SimpleWhenMissing + , SimpleWhenMatched + , merge + , runWhenMatched + , runWhenMissing + + -- *** @WhenMatched@ tactics + , zipWithMaybeMatched + , zipWithMatched + + -- *** @WhenMissing@ tactics + , mapMaybeMissing + , dropMissing + , preserveMissing + , mapMissing + , filterMissing + + -- ** Applicative general combining function + , WhenMissing (..) + , WhenMatched (..) + , mergeA + + -- *** @WhenMatched@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , zipWithMaybeAMatched + , zipWithAMatched + + -- *** @WhenMissing@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , traverseMaybeMissing + , traverseMissing + , filterAMissing + + -- *** Covariant maps for tactics + , mapWhenMissing + , mapWhenMatched + + -- ** Deprecated general combining function + + , mergeWithKey + + -- * Traversal + -- ** Map + , map + , mapWithKey + , traverseWithKey + , traverseMaybeWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeys + , mapKeysWith + , mapKeysMonotonic + + -- * Folds + , foldr + , foldl + , foldrWithKey + , foldlWithKey + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + , fromSet + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + , fromDescList + , fromDescListWith + , fromDescListWithKey + , fromDistinctDescList + + -- * Filter + , filter + , filterWithKey + , restrictKeys + , withoutKeys + , partition + , partitionWithKey + , takeWhileAntitone + , dropWhileAntitone + , spanAntitone + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , updateAt + , deleteAt + , take + , drop + , splitAt + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + , valid + + , bin + , balanced + , link + , link2 + ) where + +import Prelude hiding (lookup,map,filter,foldr,foldl,null,take,drop,splitAt) + +import Data.Map.Base + ( Map (..) + , AreWeStrict (..) + , WhenMissing (..) + , WhenMatched (..) + , runWhenMatched + , runWhenMissing + , SimpleWhenMissing + , SimpleWhenMatched + , preserveMissing + , dropMissing + , filterMissing + , filterAMissing + , merge + , mergeA + , (!) + , (\\) + , assocs + , atKeyImpl +#if MIN_VERSION_base(4,8,0) + , atKeyPlain +#endif + , balance + , balanceL + , balanceR + , elemAt + , elems + , empty + , delete + , deleteAt + , deleteFindMax + , deleteFindMin + , deleteMin + , deleteMax + , difference + , drop + , dropWhileAntitone + , filter + , filterWithKey + , findIndex + , findMax + , findMin + , foldl + , foldl' + , foldlWithKey + , foldlWithKey' + , foldMapWithKey + , foldr + , foldr' + , foldrWithKey + , foldrWithKey' + , glue + , insertMax + , intersection + , isProperSubmapOf + , isProperSubmapOfBy + , isSubmapOf + , isSubmapOfBy + , keys + , keysSet + , link + , lookup + , lookupGE + , lookupGT + , lookupIndex + , lookupLE + , lookupLT + , mapKeys + , mapKeysMonotonic + , maxView + , maxViewWithKey + , member + , link2 + , minView + , minViewWithKey + , notMember + , null + , partition + , partitionWithKey + , restrictKeys + , showTree + , showTreeWith + , size + , spanAntitone + , split + , splitAt + , splitLookup + , splitRoot + , take + , takeWhileAntitone + , toList + , toAscList + , toDescList + , union + , unions + , valid + , withoutKeys ) + +import Data.Map.Base (bin, balanced) + +import Control.Applicative (Const (..)) +#if !MIN_VERSION_base(4,8,0) +import Control.Applicative (Applicative (..), (<$>)) +#endif +import qualified Data.Set.Base as Set +import Data.Utils.StrictFold +import Data.Utils.StrictPair + +import Data.Bits (shiftL, shiftR) +#if __GLASGOW_HASKELL__ >= 709 +import Data.Coerce +#endif + +#if __GLASGOW_HASKELL__ && MIN_VERSION_base(4,8,0) +import Data.Functor.Identity (Identity (..)) +#endif + + +-- $strictness +-- +-- This module satisfies the following strictness properties: +-- +-- 1. Key arguments are evaluated to WHNF; +-- +-- 2. Keys and values are evaluated to WHNF before they are stored in +-- the map. +-- +-- Here's an example illustrating the first property: +-- +-- > delete undefined m == undefined +-- +-- Here are some examples that illustrate the second property: +-- +-- > map (\ v -> undefined) m == undefined -- m is not empty +-- > mapKeys (\ k -> undefined) m == undefined -- m is not empty + +-- [Note: Pointer equality for sharing] +-- +-- We use pointer equality to enhance sharing between the arguments +-- of some functions and their results. Notably, we use it +-- for insert, delete, union, intersection, and difference. We do +-- *not* use it for functions, like insertWith, unionWithKey, +-- intersectionWith, etc., that allow the user to modify the elements. +-- While we *could* do so, we would only get sharing under fairly +-- narrow conditions and at a relatively high cost. It does not seem +-- worth the price. + +{-------------------------------------------------------------------- + Query +--------------------------------------------------------------------} + +-- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns +-- the value at key @k@ or returns default value @def@ +-- when the key is not in the map. +-- +-- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' +-- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' + +-- See Map.Base.Note: Local 'go' functions and capturing +findWithDefault :: Ord k => a -> k -> Map k a -> a +findWithDefault def k = k `seq` go + where + go Tip = def + go (Bin _ kx x l r) = case compare k kx of + LT -> go l + GT -> go r + EQ -> x +#if __GLASGOW_HASKELL__ +{-# INLINABLE findWithDefault #-} +#else +{-# INLINE findWithDefault #-} +#endif + +{-------------------------------------------------------------------- + Construction +--------------------------------------------------------------------} + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +-- > size (singleton 1 'a') == 1 + +singleton :: k -> a -> Map k a +singleton k x = x `seq` Bin 1 k x Tip Tip +{-# INLINE singleton #-} + +{-------------------------------------------------------------------- + Insertion +--------------------------------------------------------------------} +-- | /O(log n)/. 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' + +-- See Map.Base.Note: Type of local 'go' function +insert :: Ord k => k -> a -> Map k a -> Map k a +insert = go + where + go :: Ord k => k -> a -> Map k a -> Map k a + go !kx !x Tip = singleton kx x + go kx x (Bin sz ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go kx x l) r + GT -> balanceR ky y l (go kx x r) + EQ -> Bin sz kx x l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE insert #-} +#else +{-# INLINE insert #-} +#endif + +-- | /O(log n)/. 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" + +insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWith = go + where + go :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a + go _ !kx x Tip = singleton kx x + go f !kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go f kx x l) r + GT -> balanceR ky y l (go f kx x r) + EQ -> let !y' = f x y in Bin sy kx y' l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertWith #-} +#else +{-# INLINE insertWith #-} +#endif + +insertWithR :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWithR = go + where + go :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a + go _ !kx x Tip = singleton kx x + go f !kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go f kx x l) r + GT -> balanceR ky y l (go f kx x r) + EQ -> let !y' = f y x in Bin sy ky y' l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertWithR #-} +#else +{-# INLINE insertWithR #-} +#endif + +-- | /O(log n)/. 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" + +-- See Map.Base.Note: Type of local 'go' function +insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWithKey = go + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a + -- Forcing `kx` may look redundant, but it's possible `compare` will + -- be lazy. + go _ !kx x Tip = singleton kx x + go f kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go f kx x l) r + GT -> balanceR ky y l (go f kx x r) + EQ -> let !x' = f kx x y + in Bin sy kx x' l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertWithKey #-} +#else +{-# INLINE insertWithKey #-} +#endif + +insertWithKeyR :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWithKeyR = go + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a + -- Forcing `kx` may look redundant, but it's possible `compare` will + -- be lazy. + go _ !kx x Tip = singleton kx x + go f kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go f kx x l) r + GT -> balanceR ky y l (go f kx x r) + EQ -> let !y' = f ky y x + in Bin sy ky y' l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertWithKeyR #-} +#else +{-# INLINE insertWithKeyR #-} +#endif + +-- | /O(log n)/. 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") +-- +-- This is how to define @insertLookup@ using @insertLookupWithKey@: +-- +-- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t +-- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) +-- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) + +-- See Map.Base.Note: Type of local 'go' function +insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a + -> (Maybe a, Map k a) +insertLookupWithKey f0 kx0 x0 t0 = toPair $ go f0 kx0 x0 t0 + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> StrictPair (Maybe a) (Map k a) + go _ !kx x Tip = Nothing :*: singleton kx x + go f kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> let (found :*: l') = go f kx x l + in found :*: balanceL ky y l' r + GT -> let (found :*: r') = go f kx x r + in found :*: balanceR ky y l r' + EQ -> let x' = f kx x y + in x' `seq` (Just y :*: Bin sy kx x' l r) +#if __GLASGOW_HASKELL__ +{-# INLINABLE insertLookupWithKey #-} +#else +{-# INLINE insertLookupWithKey #-} +#endif + +{-------------------------------------------------------------------- + Deletion +--------------------------------------------------------------------} + +-- | /O(log n)/. 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 + +adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a +adjust f = adjustWithKey (\_ x -> f x) +#if __GLASGOW_HASKELL__ +{-# INLINABLE adjust #-} +#else +{-# INLINE adjust #-} +#endif + +-- | /O(log n)/. 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 + +adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a +adjustWithKey = go + where + go :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a + go _ !_ Tip = Tip + go f k (Bin sx kx x l r) = + case compare k kx of + LT -> Bin sx kx x (go f k l) r + GT -> Bin sx kx x l (go f k r) + EQ -> Bin sx kx x' l r + where !x' = f kx x +#if __GLASGOW_HASKELL__ +{-# INLINABLE adjustWithKey #-} +#else +{-# INLINE adjustWithKey #-} +#endif + +-- | /O(log n)/. 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 :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a +update f = updateWithKey (\_ x -> f x) +#if __GLASGOW_HASKELL__ +{-# INLINABLE update #-} +#else +{-# INLINE update #-} +#endif + +-- | /O(log n)/. 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" + +-- See Map.Base.Note: Type of local 'go' function +updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a +updateWithKey = go + where + go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a + go _ !_ Tip = Tip + go f k(Bin sx kx x l r) = + case compare k kx of + LT -> balanceR kx x (go f k l) r + GT -> balanceL kx x l (go f k r) + EQ -> case f kx x of + Just x' -> x' `seq` Bin sx kx x' l r + Nothing -> glue l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE updateWithKey #-} +#else +{-# INLINE updateWithKey #-} +#endif + +-- | /O(log n)/. Lookup and update. See also 'updateWithKey'. +-- The function returns changed value, if it is updated. +-- Returns the original key value if the map entry is deleted. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "5:new a", fromList [(3, "b"), (5, "5:new a")]) +-- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) +-- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") + +-- See Map.Base.Note: Type of local 'go' function +updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) +updateLookupWithKey f0 k0 t0 = toPair $ go f0 k0 t0 + where + go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> StrictPair (Maybe a) (Map k a) + go _ !_ Tip = (Nothing :*: Tip) + go f k (Bin sx kx x l r) = + case compare k kx of + LT -> let (found :*: l') = go f k l + in found :*: balanceR kx x l' r + GT -> let (found :*: r') = go f k r + in found :*: balanceL kx x l r' + EQ -> case f kx x of + Just x' -> x' `seq` (Just x' :*: Bin sx kx x' l r) + Nothing -> (Just x :*: glue l r) +#if __GLASGOW_HASKELL__ +{-# INLINABLE updateLookupWithKey #-} +#else +{-# INLINE updateLookupWithKey #-} +#endif + +-- | /O(log n)/. 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 'Map'. +-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. +-- +-- > let f _ = Nothing +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > +-- > let f _ = Just "c" +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "c")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "c")] + +-- See Map.Base.Note: Type of local 'go' function +alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a +alter = go + where + go :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a + go f !k Tip = case f Nothing of + Nothing -> Tip + Just x -> singleton k x + + go f k (Bin sx kx x l r) = case compare k kx of + LT -> balance kx x (go f k l) r + GT -> balance kx x l (go f k r) + EQ -> case f (Just x) of + Just x' -> x' `seq` Bin sx kx x' l r + Nothing -> glue l r +#if __GLASGOW_HASKELL__ +{-# INLINABLE alter #-} +#else +{-# INLINE alter #-} +#endif + +-- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at @k@, or absence thereof. +-- 'alterF' can be used to inspect, insert, delete, or update a value in a 'Map'. +-- In short: @'lookup' k \<$\> 'alterF' f k m = f ('lookup' k m)@. +-- +-- Example: +-- +-- @ +-- interactiveAlter :: Int -> Map Int String -> IO (Map Int String) +-- interactiveAlter k m = alterF f k m where +-- f Nothing -> do +-- putStrLn $ show k ++ +-- " was not found in the map. Would you like to add it?" +-- getUserResponse1 :: IO (Maybe String) +-- f (Just old) -> do +-- putStrLn "The key is currently bound to " ++ show old ++ +-- ". Would you like to change or delete it?" +-- getUserresponse2 :: IO (Maybe String) +-- @ +-- +-- 'alterF' is the most general operation for working with an individual +-- key that may or may not be in a given map. When used with trivial +-- functors like 'Identity' and 'Const', it is often slightly slower than +-- more specialized combinators like 'lookup' and 'insert'. However, when +-- the functor is non-trivial and key comparison is not particularly cheap, +-- it is the fastest way. +-- +-- Note on rewrite rules: +-- +-- This module includes GHC rewrite rules to optimize 'alterF' for +-- the 'Const' and 'Identity' functors. In general, these rules +-- improve performance. The sole exception is that when using +-- 'Identity', deleting a key that is already absent takes longer +-- than it would without the rules. If you expect this to occur +-- a very large fraction of the time, you might consider using a +-- private copy of the 'Identity' type. +-- +-- Note: 'alterF' is a flipped version of the 'at' combinator from +-- 'Control.Lens.At'. +alterF :: (Functor f, Ord k) + => (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a) +alterF f k m = atKeyImpl Strict k f m + +#ifndef __GLASGOW_HASKELL__ +{-# INLINE alterF #-} +#else +{-# INLINABLE [2] alterF #-} + +-- We can save a little time by recognizing the special case of +-- `Control.Applicative.Const` and just doing a lookup. +{-# RULES +"alterF/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterF f k = \m -> Const . getConst . f $ lookup k m + #-} +#if MIN_VERSION_base(4,8,0) +-- base 4.8 and above include Data.Functor.Identity, so we can +-- save a pretty decent amount of time by handling it specially. +{-# RULES +"alterF/Identity" forall k f . alterF f k = atKeyIdentity k f + #-} + +atKeyIdentity :: Ord k => k -> (Maybe a -> Identity (Maybe a)) -> Map k a -> Identity (Map k a) +atKeyIdentity k f t = Identity $ atKeyPlain Strict k (coerce f) t +{-# INLINABLE atKeyIdentity #-} +#endif +#endif + +{-------------------------------------------------------------------- + Indexing +--------------------------------------------------------------------} + +-- | /O(log n)/. 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 + +updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a +updateAt f i t = i `seq` + case t of + Tip -> error "Map.updateAt: index out of range" + Bin sx kx x l r -> case compare i sizeL of + LT -> balanceR kx x (updateAt f i l) r + GT -> balanceL kx x l (updateAt f (i-sizeL-1) r) + EQ -> case f kx x of + Just x' -> x' `seq` Bin sx kx x' l r + Nothing -> glue l r + where + sizeL = size l + +{-------------------------------------------------------------------- + Minimal, Maximal +--------------------------------------------------------------------} + +-- | /O(log n)/. 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" + +updateMin :: (a -> Maybe a) -> Map k a -> Map k a +updateMin f m + = updateMinWithKey (\_ x -> f x) m + +-- | /O(log n)/. 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" + +updateMax :: (a -> Maybe a) -> Map k a -> Map k a +updateMax f m + = updateMaxWithKey (\_ x -> f x) m + + +-- | /O(log n)/. 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" + +updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a +updateMinWithKey _ Tip = Tip +updateMinWithKey f (Bin sx kx x Tip r) = case f kx x of + Nothing -> r + Just x' -> x' `seq` Bin sx kx x' Tip r +updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r + +-- | /O(log n)/. 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" + +updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a +updateMaxWithKey _ Tip = Tip +updateMaxWithKey f (Bin sx kx x l Tip) = case f kx x of + Nothing -> l + Just x' -> x' `seq` Bin sx kx x' l Tip +updateMaxWithKey f (Bin _ kx x l r) = balanceL kx x l (updateMaxWithKey f r) + +{-------------------------------------------------------------------- + Union. +--------------------------------------------------------------------} + +-- | The union of a list of maps, with a combining operation: +-- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@). +-- +-- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] +-- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] + +unionsWith :: Ord k => (a->a->a) -> [Map k a] -> Map k a +unionsWith f ts + = foldlStrict (unionWith f) empty ts +#if __GLASGOW_HASKELL__ +{-# INLINABLE unionsWith #-} +#endif + +{-------------------------------------------------------------------- + Union with a combining function +--------------------------------------------------------------------} +-- | /O(m*log(n\/m + 1)), m <= n/. Union with a combining function. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] + +unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a +unionWith _f t1 Tip = t1 +unionWith f t1 (Bin _ k x Tip Tip) = insertWithR f k x t1 +unionWith f (Bin _ k x Tip Tip) t2 = insertWith f k x t2 +unionWith _f Tip t2 = t2 +unionWith f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of + (l2, mb, r2) -> link k1 x1' (unionWith f l1 l2) (unionWith f r1 r2) + where !x1' = maybe x1 (f x1) mb +#if __GLASGOW_HASKELL__ +{-# INLINABLE unionWith #-} +#endif + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- Union with a combining function. +-- +-- > 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")] + +unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a +unionWithKey _f t1 Tip = t1 +unionWithKey f t1 (Bin _ k x Tip Tip) = insertWithKeyR f k x t1 +unionWithKey f (Bin _ k x Tip Tip) t2 = insertWithKey f k x t2 +unionWithKey _f Tip t2 = t2 +unionWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of + (l2, mb, r2) -> link k1 x1' (unionWithKey f l1 l2) (unionWithKey f r1 r2) + where !x1' = maybe x1 (f k1 x1) mb +#if __GLASGOW_HASKELL__ +{-# INLINABLE unionWithKey #-} +#endif + +{-------------------------------------------------------------------- + Difference +--------------------------------------------------------------------} + +-- | /O(n+m)/. 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@. +-- +-- > 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" + +differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a +differenceWith f = merge preserveMissing dropMissing (zipWithMaybeMatched $ \_ x1 x2 -> f x1 x2) +#if __GLASGOW_HASKELL__ +{-# INLINABLE differenceWith #-} +#endif + +-- | /O(n+m)/. 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@. +-- +-- > 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" + +differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a +differenceWithKey f = merge preserveMissing dropMissing (zipWithMaybeMatched f) +#if __GLASGOW_HASKELL__ +{-# INLINABLE differenceWithKey #-} +#endif + + +{-------------------------------------------------------------------- + Intersection +--------------------------------------------------------------------} + +-- | /O(m*log(n\/m + 1)), m <= n/. Intersection with a combining function. +-- +-- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" + +intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c +intersectionWith _f Tip _ = Tip +intersectionWith _f _ Tip = Tip +intersectionWith f (Bin _ k x1 l1 r1) t2 = case mb of + Just x2 -> let !x1' = f x1 x2 in link k x1' l1l2 r1r2 + Nothing -> link2 l1l2 r1r2 + where + !(l2, mb, r2) = splitLookup k t2 + !l1l2 = intersectionWith f l1 l2 + !r1r2 = intersectionWith f r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE intersectionWith #-} +#endif + +-- | /O(m*log(n\/m + 1)), m <= n/. Intersection with a combining function. +-- +-- > 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" + +intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c +intersectionWithKey _f Tip _ = Tip +intersectionWithKey _f _ Tip = Tip +intersectionWithKey f (Bin _ k x1 l1 r1) t2 = case mb of + Just x2 -> let !x1' = f k x1 x2 in link k x1' l1l2 r1r2 + Nothing -> link2 l1l2 r1r2 + where + !(l2, mb, r2) = splitLookup k t2 + !l1l2 = intersectionWithKey f l1 l2 + !r1r2 = intersectionWithKey f r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE intersectionWithKey #-} +#endif + +-- | Map covariantly over a @'WhenMissing' f k x@. +mapWhenMissing :: Functor f => (a -> b) -> WhenMissing f k x a -> WhenMissing f k x b +mapWhenMissing f q = WhenMissing + { missingSubtree = fmap (map f) . missingSubtree q + , missingKey = \k x -> fmap (forceMaybe . fmap f) $ missingKey q k x} + +-- | Map covariantly over a @'WhenMatched' f k x y@. +mapWhenMatched :: Functor f => (a -> b) -> WhenMatched f k x y a -> WhenMatched f k x y b +mapWhenMatched f q = WhenMatched + { matchedKey = \k x y -> fmap (forceMaybe . fmap f) $ runWhenMatched q k x y } + +-- | When a key is found in both maps, apply a function to the +-- key and values and maybe use the result in the merged map. +-- +-- @ +-- zipWithMaybeMatched :: (k -> x -> y -> Maybe z) +-- -> SimpleWhenMatched k x y z +-- @ +zipWithMaybeMatched :: Applicative f + => (k -> x -> y -> Maybe z) + -> WhenMatched f k x y z +zipWithMaybeMatched f = WhenMatched $ + \k x y -> pure $! forceMaybe $! f k x y +{-# INLINE zipWithMaybeMatched #-} + +-- | When a key is found in both maps, apply a function to the +-- key and values, perform the resulting action, and maybe use +-- the result in the merged map. +-- +-- This is the fundamental 'WhenMatched' tactic. +zipWithMaybeAMatched :: Applicative f + => (k -> x -> y -> f (Maybe z)) + -> WhenMatched f k x y z +zipWithMaybeAMatched f = WhenMatched $ + \ k x y -> forceMaybe <$> f k x y +{-# INLINE zipWithMaybeAMatched #-} + +-- | When a key is found in both maps, apply a function to the +-- key and values to produce an action and use its result in the merged map. +zipWithAMatched :: Applicative f + => (k -> x -> y -> f z) + -> WhenMatched f k x y z +zipWithAMatched f = WhenMatched $ + \ k x y -> (Just $!) <$> f k x y +{-# INLINE zipWithAMatched #-} + +-- | When a key is found in both maps, apply a function to the +-- key and values and use the result in the merged map. +-- +-- @ +-- zipWithMatched :: (k -> x -> y -> z) +-- -> SimpleWhenMatched k x y z +-- @ +zipWithMatched :: Applicative f + => (k -> x -> y -> z) -> WhenMatched f k x y z +zipWithMatched f = WhenMatched $ + \k x y -> pure $! Just $! f k x y +{-# INLINE zipWithMatched #-} + +-- | Map over the entries whose keys are missing from the other map, +-- optionally removing some. This is the most powerful 'SimpleWhenMissing' +-- tactic, but others are usually more efficient. +-- +-- @ +-- mapMaybeMissing :: (k -> x -> Maybe y) -> SimpleWhenMissing k x y +-- @ +-- +-- prop> mapMaybeMissing f = traverseMaybeMissing (\k x -> pure (f k x)) +-- +-- but @mapMaybeMissing@ uses fewer unnecessary 'Applicative' operations. +mapMaybeMissing :: Applicative f => (k -> x -> Maybe y) -> WhenMissing f k x y +mapMaybeMissing f = WhenMissing + { missingSubtree = \m -> pure $! mapMaybeWithKey f m + , missingKey = \k x -> pure $! forceMaybe $! f k x } +{-# INLINE mapMaybeMissing #-} + +-- | Map over the entries whose keys are missing from the other map. +-- +-- @ +-- mapMissing :: (k -> x -> y) -> SimpleWhenMissing k x y +-- @ +-- +-- prop> mapMissing f = mapMaybeMissing (\k x -> Just $ f k x) +-- +-- but @mapMissing@ is somewhat faster. +mapMissing :: Applicative f => (k -> x -> y) -> WhenMissing f k x y +mapMissing f = WhenMissing + { missingSubtree = \m -> pure $! mapWithKey f m + , missingKey = \k x -> pure $! Just $! f k x } +{-# INLINE mapMissing #-} + +-- | Traverse over the entries whose keys are missing from the other map, +-- optionally producing values to put in the result. +-- This is the most powerful 'WhenMissing' tactic, but others are usually +-- more efficient. +traverseMaybeMissing :: Applicative f + => (k -> x -> f (Maybe y)) -> WhenMissing f k x y +traverseMaybeMissing f = WhenMissing + { missingSubtree = traverseMaybeWithKey f + , missingKey = \k x -> forceMaybe <$> f k x } +{-# INLINE traverseMaybeMissing #-} + +-- | Traverse over the entries whose keys are missing from the other map. +traverseMissing :: Applicative f + => (k -> x -> f y) -> WhenMissing f k x y +traverseMissing f = WhenMissing + { missingSubtree = traverseWithKey f + , missingKey = \k x -> (Just $!) <$> f k x } +{-# INLINE traverseMissing #-} + +forceMaybe :: Maybe a -> Maybe a +forceMaybe Nothing = Nothing +forceMaybe m@(Just !_) = m +{-# INLINE forceMaybe #-} + +{-------------------------------------------------------------------- + MergeWithKey +--------------------------------------------------------------------} + +-- | /O(n+m)/. An unsafe universal combining function. +-- +-- WARNING: This function can produce corrupt maps and its results +-- may depend on the internal structures of its inputs. Users should +-- prefer 'Data.Map.Strict.Merge.merge' or +-- 'Data.Map.Strict.Merge.mergeA'. +-- +-- When 'mergeWithKey' is given three arguments, it is inlined to the call +-- site. You should therefore use 'mergeWithKey' only to define custom +-- combining functions. For example, you could define 'unionWithKey', +-- 'differenceWithKey' and 'intersectionWithKey' as +-- +-- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2 +-- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2 +-- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2 +-- +-- When calling @'mergeWithKey' combine only1 only2@, a function combining two +-- 'Map's is created, such that +-- +-- * if a key is present in both maps, it is passed with both corresponding +-- values to the @combine@ function. Depending on the result, the key is either +-- present in the result with specified value, or is left out; +-- +-- * a nonempty subtree present only in the first map is passed to @only1@ and +-- the output is added to the result; +-- +-- * a nonempty subtree present only in the second map is passed to @only2@ and +-- the output is added to the result. +-- +-- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/. +-- The values can be modified arbitrarily. Most common variants of @only1@ and +-- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or +-- @'filterWithKey' f@ could be used for any @f@. + +mergeWithKey :: Ord k + => (k -> a -> b -> Maybe c) + -> (Map k a -> Map k c) + -> (Map k b -> Map k c) + -> Map k a -> Map k b -> Map k c +mergeWithKey f g1 g2 = go + where + go Tip t2 = g2 t2 + go t1 Tip = g1 t1 + go (Bin _ kx x l1 r1) t2 = + case found of + Nothing -> case g1 (singleton kx x) of + Tip -> link2 l' r' + (Bin _ _ x' Tip Tip) -> link kx x' l' r' + _ -> error "mergeWithKey: Given function only1 does not fulfill required conditions (see documentation)" + Just x2 -> case f kx x x2 of + Nothing -> link2 l' r' + Just x' -> link kx x' l' r' + where + (l2, found, r2) = splitLookup kx t2 + l' = go l1 l2 + r' = go r1 r2 +{-# INLINE mergeWithKey #-} + +{-------------------------------------------------------------------- + Filter and partition +--------------------------------------------------------------------} + +-- | /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" + +mapMaybe :: (a -> Maybe b) -> Map k a -> Map k b +mapMaybe f = mapMaybeWithKey (\_ x -> f x) + +-- | /O(n)/. 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" + +mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b +mapMaybeWithKey _ Tip = Tip +mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of + Just y -> y `seq` link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r) + Nothing -> link2 (mapMaybeWithKey f l) (mapMaybeWithKey f r) + +-- | /O(n)/. Traverse keys\/values and collect the 'Just' results. + +traverseMaybeWithKey :: Applicative f + => (k -> a -> f (Maybe b)) -> Map k a -> f (Map k b) +traverseMaybeWithKey = go + where + go _ Tip = pure Tip + go f (Bin _ kx x Tip Tip) = maybe Tip (\ !x' -> Bin 1 kx x' Tip Tip) <$> f kx x + go f (Bin _ kx x l r) = combine <$> go f l <*> f kx x <*> go f r + where + combine !l' mx !r' = case mx of + Nothing -> link2 l' r' + Just !x' -> link kx x' l' r' + +-- | /O(n)/. 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")]) + +mapEither :: (a -> Either b c) -> Map k a -> (Map k b, Map k c) +mapEither f m + = mapEitherWithKey (\_ x -> f x) m + +-- | /O(n)/. 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")]) + +mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) +mapEitherWithKey f0 t0 = toPair $ go f0 t0 + where + go _ Tip = (Tip :*: Tip) + go f (Bin _ kx x l r) = case f kx x of + Left y -> y `seq` (link kx y l1 r1 :*: link2 l2 r2) + Right z -> z `seq` (link2 l1 r1 :*: link kx z l2 r2) + where + (l1 :*: l2) = go f l + (r1 :*: r2) = go f r + +{-------------------------------------------------------------------- + Mapping +--------------------------------------------------------------------} +-- | /O(n)/. Map a function over all values in the map. +-- +-- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] + +map :: (a -> b) -> Map k a -> Map k b +map f = go + where + go Tip = Tip + go (Bin sx kx x l r) = let !x' = f x in Bin sx kx x' (go l) (go r) +-- We use `go` to let `map` inline. This is important if `f` is a constant +-- function. + +#ifdef __GLASGOW_HASKELL__ +{-# NOINLINE [1] map #-} +{-# RULES +"map/map" forall f g xs . map f (map g xs) = map (f . g) xs + #-} +#endif +#if __GLASGOW_HASKELL__ >= 709 +-- Safe coercions were introduced in 7.8, but did not work well with RULES yet. +{-# RULES +"mapSeq/coerce" map coerce = coerce + #-} +#endif + +-- | /O(n)/. 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")] + +mapWithKey :: (k -> a -> b) -> Map k a -> Map k b +mapWithKey _ Tip = Tip +mapWithKey f (Bin sx kx x l r) = + let x' = f kx x + in x' `seq` Bin sx kx x' (mapWithKey f l) (mapWithKey f r) + +#ifdef __GLASGOW_HASKELL__ +{-# NOINLINE [1] mapWithKey #-} +{-# RULES +"mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) = + mapWithKey (\k a -> f k (g k a)) xs +"mapWithKey/map" forall f g xs . mapWithKey f (map g xs) = + mapWithKey (\k a -> f k (g a)) xs +"map/mapWithKey" forall f g xs . map f (mapWithKey g xs) = + mapWithKey (\k a -> f (g k a)) xs + #-} +#endif + +-- | /O(n)/. +-- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (\v' -> v' `seq` (k,v')) <$> f k v) ('toList' m)@ +-- That is, it behaves much like a regular 'traverse' except that the traversing +-- function also has access to the key associated with a value and the values are +-- forced before they are installed in the result map. +-- +-- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')]) +-- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')]) == Nothing +traverseWithKey :: Applicative t => (k -> a -> t b) -> Map k a -> t (Map k b) +traverseWithKey f = go + where + go Tip = pure Tip + go (Bin 1 k v _ _) = (\ !v' -> Bin 1 k v' Tip Tip) <$> f k v + go (Bin s k v l r) = (\ l' !v' r' -> Bin s k v' l' r') <$> go l <*> f k v <*> go r +{-# INLINE traverseWithKey #-} + +-- | /O(n)/. The function 'mapAccum' threads an accumulating +-- argument through the map in ascending order of keys. +-- +-- > let f a b = (a ++ b, b ++ "X") +-- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) + +mapAccum :: (a -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccum f a m + = mapAccumWithKey (\a' _ x' -> f a' x') a m + +-- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating +-- argument through the map in ascending order of keys. +-- +-- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") +-- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) + +mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumWithKey f a t + = mapAccumL f a t + +-- | /O(n)/. The function 'mapAccumL' threads an accumulating +-- argument through the map in ascending order of keys. +mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumL _ a Tip = (a,Tip) +mapAccumL f a (Bin sx kx x l r) = + let (a1,l') = mapAccumL f a l + (a2,x') = f a1 kx x + (a3,r') = mapAccumL f a2 r + in x' `seq` (a3,Bin sx kx x' l' r') + +-- | /O(n)/. The function 'mapAccumR' threads an accumulating +-- argument through the map in descending order of keys. +mapAccumRWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumRWithKey _ a Tip = (a,Tip) +mapAccumRWithKey f a (Bin sx kx x l r) = + let (a1,r') = mapAccumRWithKey f a r + (a2,x') = f a1 kx x + (a3,l') = mapAccumRWithKey f a2 l + in x' `seq` (a3,Bin sx kx x' l' r') + +-- | /O(n*log n)/. +-- @'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" + +mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a +mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] +#if __GLASGOW_HASKELL__ +{-# INLINABLE mapKeysWith #-} +#endif + +{-------------------------------------------------------------------- + Conversions +--------------------------------------------------------------------} + +-- | /O(n)/. Build a map from a set of keys and a function which for each key +-- computes its value. +-- +-- > fromSet (\k -> replicate k 'a') (Data.Set.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")] +-- > fromSet undefined Data.Set.empty == empty + +fromSet :: (k -> a) -> Set.Set k -> Map k a +fromSet _ Set.Tip = Tip +fromSet f (Set.Bin sz x l r) = case f x of v -> v `seq` Bin sz x v (fromSet f l) (fromSet f r) + +{-------------------------------------------------------------------- + Lists + use [foldlStrict] to reduce demand on the control-stack +--------------------------------------------------------------------} +-- | /O(n*log n)/. 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. +-- +-- If the keys of the list are ordered, linear-time implementation is used, +-- with the performance equal to 'fromDistinctAscList'. +-- +-- > 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")] + +-- For some reason, when 'singleton' is used in fromList or in +-- create, it is not inlined, so we inline it manually. +fromList :: Ord k => [(k,a)] -> Map k a +fromList [] = Tip +fromList [(kx, x)] = x `seq` Bin 1 kx x Tip Tip +fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = x0 `seq` fromList' (Bin 1 kx0 x0 Tip Tip) xs0 + | otherwise = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 + where + not_ordered _ [] = False + not_ordered kx ((ky,_) : _) = kx >= ky + {-# INLINE not_ordered #-} + + fromList' t0 xs = foldlStrict ins t0 xs + where ins t (k,x) = insert k x t + + go !_ t [] = t + go _ t [(kx, x)] = x `seq` insertMax kx x t + go s l xs@((kx, x) : xss) | not_ordered kx xss = fromList' l xs + | otherwise = case create s xss of + (r, ys, []) -> x `seq` go (s `shiftL` 1) (link kx x l r) ys + (r, _, ys) -> x `seq` fromList' (link kx x l r) ys + + -- The create is returning a triple (tree, xs, ys). Both xs and ys + -- represent not yet processed elements and only one of them can be nonempty. + -- If ys is nonempty, the keys in ys are not ordered with respect to tree + -- and must be inserted using fromList'. Otherwise the keys have been + -- ordered so far. + create !_ [] = (Tip, [], []) + create s xs@(xp : xss) + | s == 1 = case xp of (kx, x) | not_ordered kx xss -> x `seq` (Bin 1 kx x Tip Tip, [], xss) + | otherwise -> x `seq` (Bin 1 kx x Tip Tip, xss, []) + | otherwise = case create (s `shiftR` 1) xs of + res@(_, [], _) -> res + (l, [(ky, y)], zs) -> y `seq` (insertMax ky y l, [], zs) + (l, ys@((ky, y):yss), _) | not_ordered ky yss -> (l, [], ys) + | otherwise -> case create (s `shiftR` 1) yss of + (r, zs, ws) -> y `seq` (link ky y l r, zs, ws) +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromList #-} +#endif + +-- | /O(n*log n)/. 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 + +fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a +fromListWith f xs + = fromListWithKey (\_ x y -> f x y) xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromListWith #-} +#endif + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. +-- +-- > let f k a1 a2 = (show k) ++ a1 ++ a2 +-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "3ab"), (5, "5a5ba")] +-- > fromListWithKey f [] == empty + +fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromListWithKey f xs + = foldlStrict ins empty xs + where + ins t (k,x) = insertWithKey f k x t +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromListWithKey #-} +#endif + +{-------------------------------------------------------------------- + Building trees from ascending/descending lists can be done in linear time. + + Note that if [xs] is ascending then: + fromAscList xs == fromList xs + fromAscListWith f xs == fromListWith f xs + + If [xs] is descending then: + fromDescList xs == fromList xs + fromDescListWith f xs == fromListWith f xs +--------------------------------------------------------------------} + +-- | /O(n)/. 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")] +-- > valid (fromAscList [(3,"b"), (5,"a"), (5,"b")]) == True +-- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False +fromAscList :: Eq k => [(k,a)] -> Map k a +fromAscList xs + = fromAscListWithKey (\_ x _ -> x) xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromAscList #-} +#endif + +-- | /O(n)/. Build a map from a descending list in linear time. +-- /The precondition (input list is descending) is not checked./ +-- +-- > fromDescList [(5,"a"), (3,"b")] == fromList [(3, "b"), (5, "a")] +-- > fromDescList [(5,"a"), (5,"b"), (3,"a")] == fromList [(3, "b"), (5, "b")] +-- > valid (fromDescList [(5,"a"), (5,"b"), (3,"b")]) == True +-- > valid (fromDescList [(5,"a"), (3,"b"), (5,"b")]) == False +fromDescList :: Eq k => [(k,a)] -> Map k a +fromDescList xs + = fromDescListWithKey (\_ x _ -> x) xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromDescList #-} +#endif + +-- | /O(n)/. 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")] +-- > valid (fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")]) == True +-- > valid (fromAscListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False + +fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a +fromAscListWith f xs + = fromAscListWithKey (\_ x y -> f x y) xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromAscListWith #-} +#endif + +-- | /O(n)/. Build a map from a descending list in linear time with a combining function for equal keys. +-- /The precondition (input list is descending) is not checked./ +-- +-- > fromDescListWith (++) [(5,"a"), (5,"b"), (3,"b")] == fromList [(3, "b"), (5, "ba")] +-- > valid (fromDescListWith (++) [(5,"a"), (5,"b"), (3,"b")]) == True +-- > valid (fromDescListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False + +fromDescListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a +fromDescListWith f xs + = fromDescListWithKey (\_ x y -> f x y) xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromDescListWith #-} +#endif + +-- | /O(n)/. 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./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] == fromList [(3, "b"), (5, "5:b5:ba")] +-- > valid (fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")]) == True +-- > valid (fromAscListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False + +fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromAscListWithKey f xs + = fromDistinctAscList (combineEq f xs) + where + -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] + combineEq _ xs' + = case xs' of + [] -> [] + [x] -> [x] + (x:xx) -> combineEq' x xx + + combineEq' z [] = [z] + combineEq' z@(kz,zz) (x@(kx,xx):xs') + | kx==kz = let yy = f kx xx zz in yy `seq` combineEq' (kx,yy) xs' + | otherwise = z:combineEq' x xs' +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromAscListWithKey #-} +#endif + +-- | /O(n)/. Build a map from a descending list in linear time with a +-- combining function for equal keys. +-- /The precondition (input list is descending) is not checked./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromDescListWithKey f [(5,"a"), (5,"b"), (5,"b"), (3,"b")] == fromList [(3, "b"), (5, "5:b5:ba")] +-- > valid (fromDescListWithKey f [(5,"a"), (5,"b"), (5,"b"), (3,"b")]) == True +-- > valid (fromDescListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False + +fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromDescListWithKey f xs + = fromDistinctDescList (combineEq f xs) + where + -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] + combineEq _ xs' + = case xs' of + [] -> [] + [x] -> [x] + (x:xx) -> combineEq' x xx + + combineEq' z [] = [z] + combineEq' z@(kz,zz) (x@(kx,xx):xs') + | kx==kz = let yy = f kx xx zz in yy `seq` combineEq' (kx,yy) xs' + | otherwise = z:combineEq' x xs' +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromDescListWithKey #-} +#endif + +-- | /O(n)/. 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")] +-- > valid (fromDistinctAscList [(3,"b"), (5,"a")]) == True +-- > valid (fromDistinctAscList [(3,"b"), (5,"a"), (5,"b")]) == False + +-- For some reason, when 'singleton' is used in fromDistinctAscList or in +-- create, it is not inlined, so we inline it manually. +fromDistinctAscList :: [(k,a)] -> Map k a +fromDistinctAscList [] = Tip +fromDistinctAscList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 + where + go !_ t [] = t + go s l ((kx, x) : xs) = case create s xs of + (r, ys) -> x `seq` go (s `shiftL` 1) (link kx x l r) ys + + create !_ [] = (Tip, []) + create s xs@(x' : xs') + | s == 1 = case x' of (kx, x) -> x `seq` (Bin 1 kx x Tip Tip, xs') + | otherwise = case create (s `shiftR` 1) xs of + res@(_, []) -> res + (l, (ky, y):ys) -> case create (s `shiftR` 1) ys of + (r, zs) -> y `seq` (link ky y l r, zs) + +-- | /O(n)/. Build a map from a descending list of distinct elements in linear time. +-- /The precondition is not checked./ +-- +-- > fromDistinctDescList [(5,"a"), (3,"b")] == fromList [(3, "b"), (5, "a")] +-- > valid (fromDistinctDescList [(5,"a"), (3,"b")]) == True +-- > valid (fromDistinctDescList [(5,"a"), (3,"b"), (3,"a")]) == False + +-- For some reason, when 'singleton' is used in fromDistinctDescList or in +-- create, it is not inlined, so we inline it manually. +fromDistinctDescList :: [(k,a)] -> Map k a +fromDistinctDescList [] = Tip +fromDistinctDescList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 + where + go !_ t [] = t + go s r ((kx, x) : xs) = case create s xs of + (l, ys) -> x `seq` go (s `shiftL` 1) (link kx x l r) ys + + create !_ [] = (Tip, []) + create s xs@(x' : xs') + | s == 1 = case x' of (kx, x) -> x `seq` (Bin 1 kx x Tip Tip, xs') + | otherwise = case create (s `shiftR` 1) xs of + res@(_, []) -> res + (r, (ky, y):ys) -> case create (s `shiftR` 1) ys of + (l, zs) -> y `seq` (link ky y l r, zs) diff --git a/Data/Map/Strict/Merge.hs b/Data/Map/Strict/Merge.hs new file mode 100644 --- /dev/null +++ b/Data/Map/Strict/Merge.hs @@ -0,0 +1,99 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Safe #-} +#endif +#if __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE RoleAnnotations #-} +{-# LANGUAGE TypeFamilies #-} +#define USE_MAGIC_PROXY 1 +#endif + +#if USE_MAGIC_PROXY +{-# LANGUAGE MagicHash #-} +#endif + +#include "containers.h" + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Map.Strict.Merge +-- Copyright : (c) David Feuer 2016 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- This module defines an API for writing functions that merge two +-- maps. The key functions are 'merge' and 'mergeA'. +-- Each of these can be used with several different "merge tactics". +-- +-- The 'merge' and 'mergeA' functions are shared by +-- the lazy and strict modules. Only the choice of merge tactics +-- determines strictness. If you use 'Data.Map.Strict.Merge.mapMissing' +-- from this module then the results will be forced before they are +-- inserted. If you use 'Data.Map.Lazy.Merge.mapMissing' from +-- "Data.Map.Lazy.Merge" then they will not. +-- +-- == Efficiency note +-- +-- The 'Category', 'Applicative', and 'Monad' instances for 'WhenMissing' +-- tactics are included because they are valid. However, they are +-- inefficient in many cases and should usually be avoided. The instances +-- for 'WhenMatched' tactics should not pose any major efficiency problems. + +module Data.Map.Strict.Merge ( + -- ** Simple merge tactic types + SimpleWhenMissing + , SimpleWhenMatched + + -- ** General combining function + , merge + + -- *** @WhenMatched@ tactics + , zipWithMaybeMatched + , zipWithMatched + + -- *** @WhenMissing@ tactics + , mapMaybeMissing + , dropMissing + , preserveMissing + , mapMissing + , filterMissing + + -- ** Applicative merge tactic types + , WhenMissing + , WhenMatched + + -- ** Applicative general combining function + , mergeA + + -- *** @WhenMatched@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , zipWithMaybeAMatched + , zipWithAMatched + + -- *** @WhenMissing@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , traverseMaybeMissing + , traverseMissing + , filterAMissing + + -- ** Covariant maps for tactics + , mapWhenMissing + , mapWhenMatched + + -- ** Miscellaneous functions on tactics + + , runWhenMatched + , runWhenMissing + ) where + +import Data.Map.Strict.Internal diff --git a/Data/Sequence.hs b/Data/Sequence.hs --- a/Data/Sequence.hs +++ b/Data/Sequence.hs @@ -1,2492 +1,153 @@ {-# LANGUAGE CPP #-} -#if __GLASGOW_HASKELL__ -{-# LANGUAGE DeriveDataTypeable #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE FlexibleInstances #-} -#endif -#if __GLASGOW_HASKELL__ >= 703 -{-# LANGUAGE Trustworthy #-} -#endif -#if __GLASGOW_HASKELL__ >= 708 -{-# LANGUAGE TypeFamilies #-} -#endif - -#include "containers.h" - ------------------------------------------------------------------------------ --- | --- Module : Data.Sequence --- Copyright : (c) Ross Paterson 2005 --- (c) Louis Wasserman 2009 --- (c) Bertram Felgenhauer, David Feuer, Ross Paterson, and --- Milan Straka 2014 --- License : BSD-style --- Maintainer : libraries@haskell.org --- Stability : experimental --- Portability : portable --- --- General purpose finite sequences. --- Apart from being finite and having strict operations, sequences --- also differ from lists in supporting a wider variety of operations --- efficiently. --- --- An amortized running time is given for each operation, with /n/ referring --- to the length of the sequence and /i/ being the integral index used by --- some operations. These bounds hold even in a persistent (shared) setting. --- --- The implementation uses 2-3 finger trees annotated with sizes, --- as described in section 4.2 of --- --- * Ralf Hinze and Ross Paterson, --- \"Finger trees: a simple general-purpose data structure\", --- /Journal of Functional Programming/ 16:2 (2006) pp 197-217. --- --- --- /Note/: Many of these operations have the same names as similar --- operations on lists in the "Prelude". The ambiguity may be resolved --- using either qualification or the @hiding@ clause. --- --- /Warning/: The size of a 'Seq' must not exceed @maxBound::Int@. Violation --- of this condition is not detected and if the size limit is exceeded, the --- behaviour of the sequence is undefined. This is unlikely to occur in most --- applications, but some care may be required when using '><', '<*>', '*>', or --- '>>', particularly repeatedly and particularly in combination with --- 'replicate' or 'fromFunction'. --- ------------------------------------------------------------------------------ - -module Data.Sequence ( -#if !defined(TESTING) - Seq, -#else - Seq(..), Elem(..), FingerTree(..), Node(..), Digit(..), -#endif - -- * Construction - empty, -- :: Seq a - singleton, -- :: a -> Seq a - (<|), -- :: a -> Seq a -> Seq a - (|>), -- :: Seq a -> a -> Seq a - (><), -- :: Seq a -> Seq a -> Seq a - fromList, -- :: [a] -> Seq a - fromFunction, -- :: Int -> (Int -> a) -> Seq a - fromArray, -- :: Ix i => Array i a -> Seq a - -- ** Repetition - replicate, -- :: Int -> a -> Seq a - replicateA, -- :: Applicative f => Int -> f a -> f (Seq a) - replicateM, -- :: Monad m => Int -> m a -> m (Seq a) - -- ** Iterative construction - iterateN, -- :: Int -> (a -> a) -> a -> Seq a - unfoldr, -- :: (b -> Maybe (a, b)) -> b -> Seq a - unfoldl, -- :: (b -> Maybe (b, a)) -> b -> Seq a - -- * Deconstruction - -- | Additional functions for deconstructing sequences are available - -- via the 'Foldable' instance of 'Seq'. - - -- ** Queries - null, -- :: Seq a -> Bool - length, -- :: Seq a -> Int - -- ** Views - ViewL(..), - viewl, -- :: Seq a -> ViewL a - ViewR(..), - viewr, -- :: Seq a -> ViewR a - -- * Scans - scanl, -- :: (a -> b -> a) -> a -> Seq b -> Seq a - scanl1, -- :: (a -> a -> a) -> Seq a -> Seq a - scanr, -- :: (a -> b -> b) -> b -> Seq a -> Seq b - scanr1, -- :: (a -> a -> a) -> Seq a -> Seq a - -- * Sublists - tails, -- :: Seq a -> Seq (Seq a) - inits, -- :: Seq a -> Seq (Seq a) - -- ** Sequential searches - takeWhileL, -- :: (a -> Bool) -> Seq a -> Seq a - takeWhileR, -- :: (a -> Bool) -> Seq a -> Seq a - dropWhileL, -- :: (a -> Bool) -> Seq a -> Seq a - dropWhileR, -- :: (a -> Bool) -> Seq a -> Seq a - spanl, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) - spanr, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) - breakl, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) - breakr, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) - partition, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) - filter, -- :: (a -> Bool) -> Seq a -> Seq a - -- * Sorting - sort, -- :: Ord a => Seq a -> Seq a - sortBy, -- :: (a -> a -> Ordering) -> Seq a -> Seq a - unstableSort, -- :: Ord a => Seq a -> Seq a - unstableSortBy, -- :: (a -> a -> Ordering) -> Seq a -> Seq a - -- * Indexing - index, -- :: Seq a -> Int -> a - adjust, -- :: (a -> a) -> Int -> Seq a -> Seq a - update, -- :: Int -> a -> Seq a -> Seq a - take, -- :: Int -> Seq a -> Seq a - drop, -- :: Int -> Seq a -> Seq a - splitAt, -- :: Int -> Seq a -> (Seq a, Seq a) - -- ** Indexing with predicates - -- | These functions perform sequential searches from the left - -- or right ends of the sequence, returning indices of matching - -- elements. - elemIndexL, -- :: Eq a => a -> Seq a -> Maybe Int - elemIndicesL, -- :: Eq a => a -> Seq a -> [Int] - elemIndexR, -- :: Eq a => a -> Seq a -> Maybe Int - elemIndicesR, -- :: Eq a => a -> Seq a -> [Int] - findIndexL, -- :: (a -> Bool) -> Seq a -> Maybe Int - findIndicesL, -- :: (a -> Bool) -> Seq a -> [Int] - findIndexR, -- :: (a -> Bool) -> Seq a -> Maybe Int - findIndicesR, -- :: (a -> Bool) -> Seq a -> [Int] - -- * Folds - -- | General folds are available via the 'Foldable' instance of 'Seq'. - foldlWithIndex, -- :: (b -> Int -> a -> b) -> b -> Seq a -> b - foldrWithIndex, -- :: (Int -> a -> b -> b) -> b -> Seq a -> b - -- * Transformations - mapWithIndex, -- :: (Int -> a -> b) -> Seq a -> Seq b - reverse, -- :: Seq a -> Seq a - -- ** Zips - zip, -- :: Seq a -> Seq b -> Seq (a, b) - zipWith, -- :: (a -> b -> c) -> Seq a -> Seq b -> Seq c - zip3, -- :: Seq a -> Seq b -> Seq c -> Seq (a, b, c) - zipWith3, -- :: (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d - zip4, -- :: Seq a -> Seq b -> Seq c -> Seq d -> Seq (a, b, c, d) - zipWith4, -- :: (a -> b -> c -> d -> e) -> Seq a -> Seq b -> Seq c -> Seq d -> Seq e -#if TESTING - Sized(..), - deep, - node2, - node3, -#endif - ) where - -import Prelude hiding ( - Functor(..), -#if MIN_VERSION_base(4,8,0) - Applicative, (<$>), foldMap, Monoid, -#endif - null, length, take, drop, splitAt, foldl, foldl1, foldr, foldr1, - scanl, scanl1, scanr, scanr1, replicate, zip, zipWith, zip3, zipWith3, - takeWhile, dropWhile, iterate, reverse, filter, mapM, sum, all) -import qualified Data.List -import Control.Applicative (Applicative(..), (<$>), Alternative, - WrappedMonad(..), liftA, liftA2, liftA3) -import qualified Control.Applicative as Applicative (Alternative(..)) -import Control.DeepSeq (NFData(rnf)) -import Control.Monad (MonadPlus(..), ap) -import Data.Monoid (Monoid(..)) -import Data.Functor (Functor(..)) -import Data.Foldable (Foldable(foldl, foldl1, foldr, foldr1, foldMap), foldl', toList) -#if MIN_VERSION_base(4,8,0) -import Data.Foldable (foldr') -#endif -#if MIN_VERSION_base(4,9,0) -import Data.Semigroup (Semigroup((<>))) -#endif -import Data.Traversable -import Data.Typeable - --- GHC specific stuff -#ifdef __GLASGOW_HASKELL__ -import GHC.Exts (build) -import Text.Read (Lexeme(Ident), lexP, parens, prec, - readPrec, readListPrec, readListPrecDefault) -import Data.Data -import Data.String (IsString(..)) -#endif - --- Array stuff, with GHC.Arr on GHC -import Data.Array (Ix, Array) -import qualified Data.Array -#ifdef __GLASGOW_HASKELL__ -import qualified GHC.Arr -#endif - --- Coercion on GHC 7.8+ -#if __GLASGOW_HASKELL__ >= 708 -import Data.Coerce -import qualified GHC.Exts -#else -#endif - --- Identity functor on base 4.8 (GHC 7.10+) -#if MIN_VERSION_base(4,8,0) -import Data.Functor.Identity (Identity(..)) -#endif - - -infixr 5 `consTree` -infixl 5 `snocTree` -infixr 5 `appendTree0` - -infixr 5 >< -infixr 5 <|, :< -infixl 5 |>, :> - -class Sized a where - size :: a -> Int - --- | General-purpose finite sequences. -newtype Seq a = Seq (FingerTree (Elem a)) - -instance Functor Seq where - fmap = fmapSeq -#ifdef __GLASGOW_HASKELL__ - x <$ s = replicate (length s) x -#endif - -fmapSeq :: (a -> b) -> Seq a -> Seq b -fmapSeq f (Seq xs) = Seq (fmap (fmap f) xs) -#ifdef __GLASGOW_HASKELL__ -{-# NOINLINE [1] fmapSeq #-} -{-# RULES -"fmapSeq/fmapSeq" forall f g xs . fmapSeq f (fmapSeq g xs) = fmapSeq (f . g) xs - #-} -#endif -#if __GLASGOW_HASKELL__ >= 709 --- Safe coercions were introduced in 7.8, but did not work well with RULES yet. -{-# RULES -"fmapSeq/coerce" fmapSeq coerce = coerce - #-} -#endif - -instance Foldable Seq where - foldMap f (Seq xs) = foldMap (foldMap f) xs - foldr f z (Seq xs) = foldr (flip (foldr f)) z xs - foldl f z (Seq xs) = foldl (foldl f) z xs - - foldr1 f (Seq xs) = getElem (foldr1 f' xs) - where f' (Elem x) (Elem y) = Elem (f x y) - - foldl1 f (Seq xs) = getElem (foldl1 f' xs) - where f' (Elem x) (Elem y) = Elem (f x y) - -#if MIN_VERSION_base(4,8,0) - length = length - {-# INLINE length #-} - null = null - {-# INLINE null #-} -#endif - -instance Traversable Seq where - traverse f (Seq xs) = Seq <$> traverse (traverse f) xs - -instance NFData a => NFData (Seq a) where - rnf (Seq xs) = rnf xs - -instance Monad Seq where - return = pure - xs >>= f = foldl' add empty xs - where add ys x = ys >< f x - (>>) = (*>) - -instance Applicative Seq where - pure = singleton - xs *> ys = cycleN (length xs) ys - - fs <*> xs@(Seq xsFT) = case viewl fs of - EmptyL -> empty - firstf :< fs' -> case viewr fs' of - EmptyR -> fmap firstf xs - Seq fs''FT :> lastf -> case rigidify xsFT of - RigidEmpty -> empty - RigidOne (Elem x) -> fmap ($x) fs - RigidTwo (Elem x1) (Elem x2) -> - Seq $ ap2FT firstf fs''FT lastf (x1, x2) - RigidThree (Elem x1) (Elem x2) (Elem x3) -> - Seq $ ap3FT firstf fs''FT lastf (x1, x2, x3) - RigidFull r@(Rigid s pr _m sf) -> Seq $ - Deep (s * length fs) - (fmap (fmap firstf) (nodeToDigit pr)) - (aptyMiddle (fmap firstf) (fmap lastf) fmap fs''FT r) - (fmap (fmap lastf) (nodeToDigit sf)) - - -ap2FT :: (a -> b) -> FingerTree (Elem (a->b)) -> (a -> b) -> (a,a) -> FingerTree (Elem b) -ap2FT firstf fs lastf (x,y) = - Deep (size fs * 2 + 4) - (Two (Elem $ firstf x) (Elem $ firstf y)) - (mapMulFT 2 (\(Elem f) -> Node2 2 (Elem (f x)) (Elem (f y))) fs) - (Two (Elem $ lastf x) (Elem $ lastf y)) - -ap3FT :: (a -> b) -> FingerTree (Elem (a->b)) -> (a -> b) -> (a,a,a) -> FingerTree (Elem b) -ap3FT firstf fs lastf (x,y,z) = Deep (size fs * 3 + 6) - (Three (Elem $ firstf x) (Elem $ firstf y) (Elem $ firstf z)) - (mapMulFT 3 (\(Elem f) -> Node3 3 (Elem (f x)) (Elem (f y)) (Elem (f z))) fs) - (Three (Elem $ lastf x) (Elem $ lastf y) (Elem $ lastf z)) - - -data Rigidified a = RigidEmpty - | RigidOne a - | RigidTwo a a - | RigidThree a a a - | RigidFull (Rigid a) -#ifdef TESTING - deriving Show -#endif - --- | A finger tree whose top level has only Two and/or Three digits, and whose --- other levels have only One and Two digits. A Rigid tree is precisely what one --- gets by unzipping/inverting a 2-3 tree, so it is precisely what we need to --- turn a finger tree into in order to transform it into a 2-3 tree. -data Rigid a = Rigid {-# UNPACK #-} !Int !(Digit23 a) (Thin (Node a)) !(Digit23 a) -#ifdef TESTING - deriving Show -#endif - --- | A finger tree whose digits are all ones and twos -data Thin a = EmptyTh - | SingleTh a - | DeepTh {-# UNPACK #-} !Int !(Digit12 a) (Thin (Node a)) !(Digit12 a) -#ifdef TESTING - deriving Show -#endif - -data Digit12 a = One12 a | Two12 a a -#ifdef TESTING - deriving Show -#endif - --- | Sometimes, we want to emphasize that we are viewing a node as a top-level --- digit of a 'Rigid' tree. -type Digit23 a = Node a - --- | 'aptyMiddle' does most of the hard work of computing @fs<*>xs@. It --- produces the center part of a finger tree, with a prefix corresponding to --- the prefix of @xs@ and a suffix corresponding to the suffix of @xs@ omitted; --- the missing suffix and prefix are added by the caller. For the recursive --- call, it squashes the prefix and the suffix into the center tree. Once it --- gets to the bottom, it turns the tree into a 2-3 tree, applies 'mapMulFT' to --- produce the main body, and glues all the pieces together. --- --- 'map23' itself is a bit horrifying because of the nested types involved. Its --- job is to map over the *elements* of a 2-3 tree, rather than the subtrees. --- If we used a higher-order nested type with MPTC, we could probably use a --- class, but as it is we have to build up 'map23' explicitly through the --- recursion. -aptyMiddle - :: (c -> d) - -> (c -> d) - -> ((a -> b) -> c -> d) - -> FingerTree (Elem (a -> b)) - -> Rigid c - -> FingerTree (Node d) - --- Not at the bottom yet - -aptyMiddle firstf - lastf - map23 - fs - (Rigid s pr (DeepTh sm prm mm sfm) sf) - = Deep (sm + s * (size fs + 1)) -- note: sm = s - size pr - size sf - (fmap (fmap firstf) (digit12ToDigit prm)) - (aptyMiddle (fmap firstf) - (fmap lastf) - (fmap . map23) - fs - (Rigid s (squashL pr prm) mm (squashR sfm sf))) - (fmap (fmap lastf) (digit12ToDigit sfm)) - --- At the bottom - -aptyMiddle firstf - lastf - map23 - fs - (Rigid s pr EmptyTh sf) - = deep - (One (fmap firstf sf)) - (mapMulFT s (\(Elem f) -> fmap (fmap (map23 f)) converted) fs) - (One (fmap lastf pr)) - where converted = node2 pr sf - -aptyMiddle firstf - lastf - map23 - fs - (Rigid s pr (SingleTh q) sf) - = deep - (Two (fmap firstf q) (fmap firstf sf)) - (mapMulFT s (\(Elem f) -> fmap (fmap (map23 f)) converted) fs) - (Two (fmap lastf pr) (fmap lastf q)) - where converted = node3 pr q sf - -digit12ToDigit :: Digit12 a -> Digit a -digit12ToDigit (One12 a) = One a -digit12ToDigit (Two12 a b) = Two a b - --- Squash the first argument down onto the left side of the second. -squashL :: Digit23 a -> Digit12 (Node a) -> Digit23 (Node a) -squashL m (One12 n) = node2 m n -squashL m (Two12 n1 n2) = node3 m n1 n2 - --- Squash the second argument down onto the right side of the first -squashR :: Digit12 (Node a) -> Digit23 a -> Digit23 (Node a) -squashR (One12 n) m = node2 n m -squashR (Two12 n1 n2) m = node3 n1 n2 m - - --- | /O(m*n)/ (incremental) Takes an /O(m)/ function and a finger tree of size --- /n/ and maps the function over the tree leaves. Unlike the usual 'fmap', the --- function is applied to the "leaves" of the 'FingerTree' (i.e., given a --- @FingerTree (Elem a)@, it applies the function to elements of type @Elem --- a@), replacing the leaves with subtrees of at least the same height, e.g., --- @Node(Node(Elem y))@. The multiplier argument serves to make the annotations --- match up properly. -mapMulFT :: Int -> (a -> b) -> FingerTree a -> FingerTree b -mapMulFT _ _ Empty = Empty -mapMulFT _mul f (Single a) = Single (f a) -mapMulFT mul f (Deep s pr m sf) = Deep (mul * s) (fmap f pr) (mapMulFT mul (mapMulNode mul f) m) (fmap f sf) - -mapMulNode :: Int -> (a -> b) -> Node a -> Node b -mapMulNode mul f (Node2 s a b) = Node2 (mul * s) (f a) (f b) -mapMulNode mul f (Node3 s a b c) = Node3 (mul * s) (f a) (f b) (f c) - --- | /O(log n)/ (incremental) Takes the extra flexibility out of a 'FingerTree' --- to make it a genuine 2-3 finger tree. The result of 'rigidify' will have --- only two and three digits at the top level and only one and two --- digits elsewhere. If the tree has fewer than four elements, 'rigidify' --- will simply extract them, and will not build a tree. -rigidify :: FingerTree (Elem a) -> Rigidified (Elem a) --- The patterns below just fix up the top level of the tree; 'rigidify' --- delegates the hard work to 'thin'. - -rigidify Empty = RigidEmpty - -rigidify (Single q) = RigidOne q - --- The left digit is Two or Three -rigidify (Deep s (Two a b) m sf) = rigidifyRight s (node2 a b) m sf -rigidify (Deep s (Three a b c) m sf) = rigidifyRight s (node3 a b c) m sf - --- The left digit is Four -rigidify (Deep s (Four a b c d) m sf) = rigidifyRight s (node2 a b) (node2 c d `consTree` m) sf - --- The left digit is One -rigidify (Deep s (One a) m sf) = case viewLTree m of - Just2 (Node2 _ b c) m' -> rigidifyRight s (node3 a b c) m' sf - Just2 (Node3 _ b c d) m' -> rigidifyRight s (node2 a b) (node2 c d `consTree` m') sf - Nothing2 -> case sf of - One b -> RigidTwo a b - Two b c -> RigidThree a b c - Three b c d -> RigidFull $ Rigid s (node2 a b) EmptyTh (node2 c d) - Four b c d e -> RigidFull $ Rigid s (node3 a b c) EmptyTh (node2 d e) - --- | /O(log n)/ (incremental) Takes a tree whose left side has been rigidified --- and finishes the job. -rigidifyRight :: Int -> Digit23 (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> Rigidified (Elem a) - --- The right digit is Two, Three, or Four -rigidifyRight s pr m (Two a b) = RigidFull $ Rigid s pr (thin m) (node2 a b) -rigidifyRight s pr m (Three a b c) = RigidFull $ Rigid s pr (thin m) (node3 a b c) -rigidifyRight s pr m (Four a b c d) = RigidFull $ Rigid s pr (thin $ m `snocTree` node2 a b) (node2 c d) - --- The right digit is One -rigidifyRight s pr m (One e) = case viewRTree m of - Just2 m' (Node2 _ a b) -> RigidFull $ Rigid s pr (thin m') (node3 a b e) - Just2 m' (Node3 _ a b c) -> RigidFull $ Rigid s pr (thin $ m' `snocTree` node2 a b) (node2 c e) - Nothing2 -> case pr of - Node2 _ a b -> RigidThree a b e - Node3 _ a b c -> RigidFull $ Rigid s (node2 a b) EmptyTh (node2 c e) - --- | /O(log n)/ (incremental) Rejigger a finger tree so the digits are all ones --- and twos. -thin :: Sized a => FingerTree a -> Thin a --- Note that 'thin12' will produce a 'DeepTh' constructor immediately before --- recursively calling 'thin'. -thin Empty = EmptyTh -thin (Single a) = SingleTh a -thin (Deep s pr m sf) = - case pr of - One a -> thin12 s (One12 a) m sf - Two a b -> thin12 s (Two12 a b) m sf - Three a b c -> thin12 s (One12 a) (node2 b c `consTree` m) sf - Four a b c d -> thin12 s (Two12 a b) (node2 c d `consTree` m) sf - -thin12 :: Sized a => Int -> Digit12 a -> FingerTree (Node a) -> Digit a -> Thin a -thin12 s pr m (One a) = DeepTh s pr (thin m) (One12 a) -thin12 s pr m (Two a b) = DeepTh s pr (thin m) (Two12 a b) -thin12 s pr m (Three a b c) = DeepTh s pr (thin $ m `snocTree` node2 a b) (One12 c) -thin12 s pr m (Four a b c d) = DeepTh s pr (thin $ m `snocTree` node2 a b) (Two12 c d) - - -instance MonadPlus Seq where - mzero = empty - mplus = (><) - -instance Alternative Seq where - empty = empty - (<|>) = (><) - -instance Eq a => Eq (Seq a) where - xs == ys = length xs == length ys && toList xs == toList ys - -instance Ord a => Ord (Seq a) where - compare xs ys = compare (toList xs) (toList ys) - -#if TESTING -instance Show a => Show (Seq a) where - showsPrec p (Seq x) = showsPrec p x -#else -instance Show a => Show (Seq a) where - showsPrec p xs = showParen (p > 10) $ - showString "fromList " . shows (toList xs) -#endif - -instance Read a => Read (Seq a) where -#ifdef __GLASGOW_HASKELL__ - readPrec = parens $ prec 10 $ do - Ident "fromList" <- lexP - xs <- readPrec - return (fromList xs) - - readListPrec = readListPrecDefault -#else - readsPrec p = readParen (p > 10) $ \ r -> do - ("fromList",s) <- lex r - (xs,t) <- reads s - return (fromList xs,t) -#endif - -instance Monoid (Seq a) where - mempty = empty -#if !(MIN_VERSION_base(4,9,0)) - mappend = (><) -#else - mappend = (<>) - -instance Semigroup (Seq a) where - (<>) = (><) -#endif - -INSTANCE_TYPEABLE1(Seq,seqTc,"Seq") - -#if __GLASGOW_HASKELL__ -instance Data a => Data (Seq a) where - gfoldl f z s = case viewl s of - EmptyL -> z empty - x :< xs -> z (<|) `f` x `f` xs - - gunfold k z c = case constrIndex c of - 1 -> z empty - 2 -> k (k (z (<|))) - _ -> error "gunfold" - - toConstr xs - | null xs = emptyConstr - | otherwise = consConstr - - dataTypeOf _ = seqDataType - - dataCast1 f = gcast1 f - -emptyConstr, consConstr :: Constr -emptyConstr = mkConstr seqDataType "empty" [] Prefix -consConstr = mkConstr seqDataType "<|" [] Infix - -seqDataType :: DataType -seqDataType = mkDataType "Data.Sequence.Seq" [emptyConstr, consConstr] -#endif - --- Finger trees - -data FingerTree a - = Empty - | Single a - | Deep {-# UNPACK #-} !Int !(Digit a) (FingerTree (Node a)) !(Digit a) -#if TESTING - deriving Show -#endif - -instance Sized a => Sized (FingerTree a) where - {-# SPECIALIZE instance Sized (FingerTree (Elem a)) #-} - {-# SPECIALIZE instance Sized (FingerTree (Node a)) #-} - size Empty = 0 - size (Single x) = size x - size (Deep v _ _ _) = v - -instance Foldable FingerTree where - foldMap _ Empty = mempty - foldMap f (Single x) = f x - foldMap f (Deep _ pr m sf) = - foldMap f pr `mappend` (foldMap (foldMap f) m `mappend` foldMap f sf) - - foldr _ z Empty = z - foldr f z (Single x) = x `f` z - foldr f z (Deep _ pr m sf) = - foldr f (foldr (flip (foldr f)) (foldr f z sf) m) pr - - foldl _ z Empty = z - foldl f z (Single x) = z `f` x - foldl f z (Deep _ pr m sf) = - foldl f (foldl (foldl f) (foldl f z pr) m) sf - - foldr1 _ Empty = error "foldr1: empty sequence" - foldr1 _ (Single x) = x - foldr1 f (Deep _ pr m sf) = - foldr f (foldr (flip (foldr f)) (foldr1 f sf) m) pr - - foldl1 _ Empty = error "foldl1: empty sequence" - foldl1 _ (Single x) = x - foldl1 f (Deep _ pr m sf) = - foldl f (foldl (foldl f) (foldl1 f pr) m) sf - -instance Functor FingerTree where - fmap _ Empty = Empty - fmap f (Single x) = Single (f x) - fmap f (Deep v pr m sf) = - Deep v (fmap f pr) (fmap (fmap f) m) (fmap f sf) - -instance Traversable FingerTree where - traverse _ Empty = pure Empty - traverse f (Single x) = Single <$> f x - traverse f (Deep v pr m sf) = - Deep v <$> traverse f pr <*> traverse (traverse f) m <*> - traverse f sf - -instance NFData a => NFData (FingerTree a) where - rnf (Empty) = () - rnf (Single x) = rnf x - rnf (Deep _ pr m sf) = rnf pr `seq` rnf sf `seq` rnf m - -{-# INLINE deep #-} -deep :: Sized a => Digit a -> FingerTree (Node a) -> Digit a -> FingerTree a -deep pr m sf = Deep (size pr + size m + size sf) pr m sf - -{-# INLINE pullL #-} -pullL :: Int -> FingerTree (Node a) -> Digit a -> FingerTree a -pullL s m sf = case viewLTree m of - Nothing2 -> digitToTree' s sf - Just2 pr m' -> Deep s (nodeToDigit pr) m' sf - -{-# INLINE pullR #-} -pullR :: Int -> Digit a -> FingerTree (Node a) -> FingerTree a -pullR s pr m = case viewRTree m of - Nothing2 -> digitToTree' s pr - Just2 m' sf -> Deep s pr m' (nodeToDigit sf) - -{-# SPECIALIZE deepL :: Maybe (Digit (Elem a)) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> FingerTree (Elem a) #-} -{-# SPECIALIZE deepL :: Maybe (Digit (Node a)) -> FingerTree (Node (Node a)) -> Digit (Node a) -> FingerTree (Node a) #-} -deepL :: Sized a => Maybe (Digit a) -> FingerTree (Node a) -> Digit a -> FingerTree a -deepL Nothing m sf = pullL (size m + size sf) m sf -deepL (Just pr) m sf = deep pr m sf - -{-# SPECIALIZE deepR :: Digit (Elem a) -> FingerTree (Node (Elem a)) -> Maybe (Digit (Elem a)) -> FingerTree (Elem a) #-} -{-# SPECIALIZE deepR :: Digit (Node a) -> FingerTree (Node (Node a)) -> Maybe (Digit (Node a)) -> FingerTree (Node a) #-} -deepR :: Sized a => Digit a -> FingerTree (Node a) -> Maybe (Digit a) -> FingerTree a -deepR pr m Nothing = pullR (size m + size pr) pr m -deepR pr m (Just sf) = deep pr m sf - --- Digits - -data Digit a - = One a - | Two a a - | Three a a a - | Four a a a a -#if TESTING - deriving Show -#endif - -instance Foldable Digit where - foldMap f (One a) = f a - foldMap f (Two a b) = f a `mappend` f b - foldMap f (Three a b c) = f a `mappend` (f b `mappend` f c) - foldMap f (Four a b c d) = f a `mappend` (f b `mappend` (f c `mappend` f d)) - - foldr f z (One a) = a `f` z - foldr f z (Two a b) = a `f` (b `f` z) - foldr f z (Three a b c) = a `f` (b `f` (c `f` z)) - foldr f z (Four a b c d) = a `f` (b `f` (c `f` (d `f` z))) - - foldl f z (One a) = z `f` a - foldl f z (Two a b) = (z `f` a) `f` b - foldl f z (Three a b c) = ((z `f` a) `f` b) `f` c - foldl f z (Four a b c d) = (((z `f` a) `f` b) `f` c) `f` d - - foldr1 _ (One a) = a - foldr1 f (Two a b) = a `f` b - foldr1 f (Three a b c) = a `f` (b `f` c) - foldr1 f (Four a b c d) = a `f` (b `f` (c `f` d)) - - foldl1 _ (One a) = a - foldl1 f (Two a b) = a `f` b - foldl1 f (Three a b c) = (a `f` b) `f` c - foldl1 f (Four a b c d) = ((a `f` b) `f` c) `f` d - -instance Functor Digit where - {-# INLINE fmap #-} - fmap f (One a) = One (f a) - fmap f (Two a b) = Two (f a) (f b) - fmap f (Three a b c) = Three (f a) (f b) (f c) - fmap f (Four a b c d) = Four (f a) (f b) (f c) (f d) - -instance Traversable Digit where - {-# INLINE traverse #-} - traverse f (One a) = One <$> f a - traverse f (Two a b) = Two <$> f a <*> f b - traverse f (Three a b c) = Three <$> f a <*> f b <*> f c - traverse f (Four a b c d) = Four <$> f a <*> f b <*> f c <*> f d - -instance NFData a => NFData (Digit a) where - rnf (One a) = rnf a - rnf (Two a b) = rnf a `seq` rnf b - rnf (Three a b c) = rnf a `seq` rnf b `seq` rnf c - rnf (Four a b c d) = rnf a `seq` rnf b `seq` rnf c `seq` rnf d - -instance Sized a => Sized (Digit a) where - {-# INLINE size #-} - size = foldl1 (+) . fmap size - -{-# SPECIALIZE digitToTree :: Digit (Elem a) -> FingerTree (Elem a) #-} -{-# SPECIALIZE digitToTree :: Digit (Node a) -> FingerTree (Node a) #-} -digitToTree :: Sized a => Digit a -> FingerTree a -digitToTree (One a) = Single a -digitToTree (Two a b) = deep (One a) Empty (One b) -digitToTree (Three a b c) = deep (Two a b) Empty (One c) -digitToTree (Four a b c d) = deep (Two a b) Empty (Two c d) - --- | Given the size of a digit and the digit itself, efficiently converts --- it to a FingerTree. -digitToTree' :: Int -> Digit a -> FingerTree a -digitToTree' n (Four a b c d) = Deep n (Two a b) Empty (Two c d) -digitToTree' n (Three a b c) = Deep n (Two a b) Empty (One c) -digitToTree' n (Two a b) = Deep n (One a) Empty (One b) -digitToTree' n (One a) = n `seq` Single a - --- Nodes - -data Node a - = Node2 {-# UNPACK #-} !Int a a - | Node3 {-# UNPACK #-} !Int a a a -#if TESTING - deriving Show -#endif - -instance Foldable Node where - foldMap f (Node2 _ a b) = f a `mappend` f b - foldMap f (Node3 _ a b c) = f a `mappend` (f b `mappend` f c) - - foldr f z (Node2 _ a b) = a `f` (b `f` z) - foldr f z (Node3 _ a b c) = a `f` (b `f` (c `f` z)) - - foldl f z (Node2 _ a b) = (z `f` a) `f` b - foldl f z (Node3 _ a b c) = ((z `f` a) `f` b) `f` c - -instance Functor Node where - {-# INLINE fmap #-} - fmap f (Node2 v a b) = Node2 v (f a) (f b) - fmap f (Node3 v a b c) = Node3 v (f a) (f b) (f c) - -instance Traversable Node where - {-# INLINE traverse #-} - traverse f (Node2 v a b) = Node2 v <$> f a <*> f b - traverse f (Node3 v a b c) = Node3 v <$> f a <*> f b <*> f c - -instance NFData a => NFData (Node a) where - rnf (Node2 _ a b) = rnf a `seq` rnf b - rnf (Node3 _ a b c) = rnf a `seq` rnf b `seq` rnf c - -instance Sized (Node a) where - size (Node2 v _ _) = v - size (Node3 v _ _ _) = v - -{-# INLINE node2 #-} -node2 :: Sized a => a -> a -> Node a -node2 a b = Node2 (size a + size b) a b - -{-# INLINE node3 #-} -node3 :: Sized a => a -> a -> a -> Node a -node3 a b c = Node3 (size a + size b + size c) a b c - -nodeToDigit :: Node a -> Digit a -nodeToDigit (Node2 _ a b) = Two a b -nodeToDigit (Node3 _ a b c) = Three a b c - --- Elements - -newtype Elem a = Elem { getElem :: a } -#if TESTING - deriving Show -#endif - -instance Sized (Elem a) where - size _ = 1 - -instance Functor Elem where -#if __GLASGOW_HASKELL__ >= 708 --- This cuts the time for <*> by around a fifth. - fmap = coerce -#else - fmap f (Elem x) = Elem (f x) -#endif - -instance Foldable Elem where - foldMap f (Elem x) = f x - foldr f z (Elem x) = f x z - foldl f z (Elem x) = f z x - -instance Traversable Elem where - traverse f (Elem x) = Elem <$> f x - -instance NFData a => NFData (Elem a) where - rnf (Elem x) = rnf x - -------------------------------------------------------- --- Applicative construction -------------------------------------------------------- -#if !MIN_VERSION_base(4,8,0) -newtype Identity a = Identity {runIdentity :: a} - -instance Functor Identity where - fmap f (Identity x) = Identity (f x) - -instance Applicative Identity where - pure = Identity - Identity f <*> Identity x = Identity (f x) -#endif - --- | This is essentially a clone of Control.Monad.State.Strict. -newtype State s a = State {runState :: s -> (s, a)} - -instance Functor (State s) where - fmap = liftA - -instance Monad (State s) where - {-# INLINE return #-} - {-# INLINE (>>=) #-} - return = pure - m >>= k = State $ \ s -> case runState m s of - (s', x) -> runState (k x) s' - -instance Applicative (State s) where - {-# INLINE pure #-} - pure x = State $ \ s -> (s, x) - (<*>) = ap - -execState :: State s a -> s -> a -execState m x = snd (runState m x) - --- | 'applicativeTree' takes an Applicative-wrapped construction of a --- piece of a FingerTree, assumed to always have the same size (which --- is put in the second argument), and replicates it as many times as --- specified. This is a generalization of 'replicateA', which itself --- is a generalization of many Data.Sequence methods. -{-# SPECIALIZE applicativeTree :: Int -> Int -> State s a -> State s (FingerTree a) #-} -{-# SPECIALIZE applicativeTree :: Int -> Int -> Identity a -> Identity (FingerTree a) #-} --- Special note: the Identity specialization automatically does node sharing, --- reducing memory usage of the resulting tree to /O(log n)/. -applicativeTree :: Applicative f => Int -> Int -> f a -> f (FingerTree a) -applicativeTree n mSize m = mSize `seq` case n of - 0 -> pure Empty - 1 -> fmap Single m - 2 -> deepA one emptyTree one - 3 -> deepA two emptyTree one - 4 -> deepA two emptyTree two - 5 -> deepA three emptyTree two - 6 -> deepA three emptyTree three - _ -> case n `quotRem` 3 of - (q,0) -> deepA three (applicativeTree (q - 2) mSize' n3) three - (q,1) -> deepA two (applicativeTree (q - 1) mSize' n3) two - (q,_) -> deepA three (applicativeTree (q - 1) mSize' n3) two - where - one = fmap One m - two = liftA2 Two m m - three = liftA3 Three m m m - deepA = liftA3 (Deep (n * mSize)) - mSize' = 3 * mSize - n3 = liftA3 (Node3 mSize') m m m - emptyTree = pure Empty - ------------------------------------------------------------------------- --- Construction ------------------------------------------------------------------------- - --- | /O(1)/. The empty sequence. -empty :: Seq a -empty = Seq Empty - --- | /O(1)/. A singleton sequence. -singleton :: a -> Seq a -singleton x = Seq (Single (Elem x)) - --- | /O(log n)/. @replicate n x@ is a sequence consisting of @n@ copies of @x@. -replicate :: Int -> a -> Seq a -replicate n x - | n >= 0 = runIdentity (replicateA n (Identity x)) - | otherwise = error "replicate takes a nonnegative integer argument" - --- | 'replicateA' is an 'Applicative' version of 'replicate', and makes --- /O(log n)/ calls to '<*>' and 'pure'. --- --- > replicateA n x = sequenceA (replicate n x) -replicateA :: Applicative f => Int -> f a -> f (Seq a) -replicateA n x - | n >= 0 = Seq <$> applicativeTree n 1 (Elem <$> x) - | otherwise = error "replicateA takes a nonnegative integer argument" - --- | 'replicateM' is a sequence counterpart of 'Control.Monad.replicateM'. --- --- > replicateM n x = sequence (replicate n x) -replicateM :: Monad m => Int -> m a -> m (Seq a) -replicateM n x - | n >= 0 = unwrapMonad (replicateA n (WrapMonad x)) - | otherwise = error "replicateM takes a nonnegative integer argument" - --- | @'cycleN' n xs@ concatenates @n@ copies of @xs@. -cycleN :: Int -> Seq a -> Seq a -cycleN n xs - | n < 0 = error "cycleN takes a nonnegative integer argument" - | n == 0 = empty - | n == 1 = xs -cycleN n (Seq xsFT) = case rigidify xsFT of - RigidEmpty -> empty - RigidOne (Elem x) -> replicate n x - RigidTwo x1 x2 -> Seq $ - Deep (n*2) pair - (runIdentity $ applicativeTree (n-2) 2 (Identity (node2 x1 x2))) - pair - where pair = Two x1 x2 - RigidThree x1 x2 x3 -> Seq $ - Deep (n*3) triple - (runIdentity $ applicativeTree (n-2) 3 (Identity (node3 x1 x2 x3))) - triple - where triple = Three x1 x2 x3 - RigidFull r@(Rigid s pr _m sf) -> Seq $ - Deep (n*s) - (nodeToDigit pr) - (cycleNMiddle (n-2) r) - (nodeToDigit sf) - -cycleNMiddle - :: Int - -> Rigid c - -> FingerTree (Node c) - -STRICT_1_OF_2(cycleNMiddle) - --- Not at the bottom yet - -cycleNMiddle n - (Rigid s pr (DeepTh sm prm mm sfm) sf) - = Deep (sm + s * (n + 1)) -- note: sm = s - size pr - size sf - (digit12ToDigit prm) - (cycleNMiddle n - (Rigid s (squashL pr prm) mm (squashR sfm sf))) - (digit12ToDigit sfm) - --- At the bottom - -cycleNMiddle n - (Rigid s pr EmptyTh sf) - = deep - (One sf) - (runIdentity $ applicativeTree n s (Identity converted)) - (One pr) - where converted = node2 pr sf - -cycleNMiddle n - (Rigid s pr (SingleTh q) sf) - = deep - (Two q sf) - (runIdentity $ applicativeTree n s (Identity converted)) - (Two pr q) - where converted = node3 pr q sf - - --- | /O(1)/. Add an element to the left end of a sequence. --- Mnemonic: a triangle with the single element at the pointy end. -(<|) :: a -> Seq a -> Seq a -x <| Seq xs = Seq (Elem x `consTree` xs) - -{-# SPECIALIZE consTree :: Elem a -> FingerTree (Elem a) -> FingerTree (Elem a) #-} -{-# SPECIALIZE consTree :: Node a -> FingerTree (Node a) -> FingerTree (Node a) #-} -consTree :: Sized a => a -> FingerTree a -> FingerTree a -consTree a Empty = Single a -consTree a (Single b) = deep (One a) Empty (One b) -consTree a (Deep s (Four b c d e) m sf) = m `seq` - Deep (size a + s) (Two a b) (node3 c d e `consTree` m) sf -consTree a (Deep s (Three b c d) m sf) = - Deep (size a + s) (Four a b c d) m sf -consTree a (Deep s (Two b c) m sf) = - Deep (size a + s) (Three a b c) m sf -consTree a (Deep s (One b) m sf) = - Deep (size a + s) (Two a b) m sf - --- | /O(1)/. Add an element to the right end of a sequence. --- Mnemonic: a triangle with the single element at the pointy end. -(|>) :: Seq a -> a -> Seq a -Seq xs |> x = Seq (xs `snocTree` Elem x) - -{-# SPECIALIZE snocTree :: FingerTree (Elem a) -> Elem a -> FingerTree (Elem a) #-} -{-# SPECIALIZE snocTree :: FingerTree (Node a) -> Node a -> FingerTree (Node a) #-} -snocTree :: Sized a => FingerTree a -> a -> FingerTree a -snocTree Empty a = Single a -snocTree (Single a) b = deep (One a) Empty (One b) -snocTree (Deep s pr m (Four a b c d)) e = m `seq` - Deep (s + size e) pr (m `snocTree` node3 a b c) (Two d e) -snocTree (Deep s pr m (Three a b c)) d = - Deep (s + size d) pr m (Four a b c d) -snocTree (Deep s pr m (Two a b)) c = - Deep (s + size c) pr m (Three a b c) -snocTree (Deep s pr m (One a)) b = - Deep (s + size b) pr m (Two a b) - --- | /O(log(min(n1,n2)))/. Concatenate two sequences. -(><) :: Seq a -> Seq a -> Seq a -Seq xs >< Seq ys = Seq (appendTree0 xs ys) - --- The appendTree/addDigits gunk below is machine generated - -appendTree0 :: FingerTree (Elem a) -> FingerTree (Elem a) -> FingerTree (Elem a) -appendTree0 Empty xs = - xs -appendTree0 xs Empty = - xs -appendTree0 (Single x) xs = - x `consTree` xs -appendTree0 xs (Single x) = - xs `snocTree` x -appendTree0 (Deep s1 pr1 m1 sf1) (Deep s2 pr2 m2 sf2) = - Deep (s1 + s2) pr1 (addDigits0 m1 sf1 pr2 m2) sf2 - -{-# SPECIALIZE addDigits0 :: FingerTree (Node (Elem a)) -> Digit (Elem a) -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> FingerTree (Node (Elem a)) #-} -{-# SPECIALIZE addDigits0 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) #-} -addDigits0 :: Sized a => FingerTree (Node a) -> Digit a -> Digit a -> FingerTree (Node a) -> FingerTree (Node a) -addDigits0 m1 (One a) (One b) m2 = - appendTree1 m1 (node2 a b) m2 -addDigits0 m1 (One a) (Two b c) m2 = - appendTree1 m1 (node3 a b c) m2 -addDigits0 m1 (One a) (Three b c d) m2 = - appendTree2 m1 (node2 a b) (node2 c d) m2 -addDigits0 m1 (One a) (Four b c d e) m2 = - appendTree2 m1 (node3 a b c) (node2 d e) m2 -addDigits0 m1 (Two a b) (One c) m2 = - appendTree1 m1 (node3 a b c) m2 -addDigits0 m1 (Two a b) (Two c d) m2 = - appendTree2 m1 (node2 a b) (node2 c d) m2 -addDigits0 m1 (Two a b) (Three c d e) m2 = - appendTree2 m1 (node3 a b c) (node2 d e) m2 -addDigits0 m1 (Two a b) (Four c d e f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits0 m1 (Three a b c) (One d) m2 = - appendTree2 m1 (node2 a b) (node2 c d) m2 -addDigits0 m1 (Three a b c) (Two d e) m2 = - appendTree2 m1 (node3 a b c) (node2 d e) m2 -addDigits0 m1 (Three a b c) (Three d e f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits0 m1 (Three a b c) (Four d e f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits0 m1 (Four a b c d) (One e) m2 = - appendTree2 m1 (node3 a b c) (node2 d e) m2 -addDigits0 m1 (Four a b c d) (Two e f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits0 m1 (Four a b c d) (Three e f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits0 m1 (Four a b c d) (Four e f g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 - -appendTree1 :: FingerTree (Node a) -> Node a -> FingerTree (Node a) -> FingerTree (Node a) -appendTree1 Empty a xs = - a `consTree` xs -appendTree1 xs a Empty = - xs `snocTree` a -appendTree1 (Single x) a xs = - x `consTree` a `consTree` xs -appendTree1 xs a (Single x) = - xs `snocTree` a `snocTree` x -appendTree1 (Deep s1 pr1 m1 sf1) a (Deep s2 pr2 m2 sf2) = - Deep (s1 + size a + s2) pr1 (addDigits1 m1 sf1 a pr2 m2) sf2 - -addDigits1 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) -addDigits1 m1 (One a) b (One c) m2 = - appendTree1 m1 (node3 a b c) m2 -addDigits1 m1 (One a) b (Two c d) m2 = - appendTree2 m1 (node2 a b) (node2 c d) m2 -addDigits1 m1 (One a) b (Three c d e) m2 = - appendTree2 m1 (node3 a b c) (node2 d e) m2 -addDigits1 m1 (One a) b (Four c d e f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits1 m1 (Two a b) c (One d) m2 = - appendTree2 m1 (node2 a b) (node2 c d) m2 -addDigits1 m1 (Two a b) c (Two d e) m2 = - appendTree2 m1 (node3 a b c) (node2 d e) m2 -addDigits1 m1 (Two a b) c (Three d e f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits1 m1 (Two a b) c (Four d e f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits1 m1 (Three a b c) d (One e) m2 = - appendTree2 m1 (node3 a b c) (node2 d e) m2 -addDigits1 m1 (Three a b c) d (Two e f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits1 m1 (Three a b c) d (Three e f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits1 m1 (Three a b c) d (Four e f g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits1 m1 (Four a b c d) e (One f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits1 m1 (Four a b c d) e (Two f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits1 m1 (Four a b c d) e (Three f g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits1 m1 (Four a b c d) e (Four f g h i) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 - -appendTree2 :: FingerTree (Node a) -> Node a -> Node a -> FingerTree (Node a) -> FingerTree (Node a) -appendTree2 Empty a b xs = - a `consTree` b `consTree` xs -appendTree2 xs a b Empty = - xs `snocTree` a `snocTree` b -appendTree2 (Single x) a b xs = - x `consTree` a `consTree` b `consTree` xs -appendTree2 xs a b (Single x) = - xs `snocTree` a `snocTree` b `snocTree` x -appendTree2 (Deep s1 pr1 m1 sf1) a b (Deep s2 pr2 m2 sf2) = - Deep (s1 + size a + size b + s2) pr1 (addDigits2 m1 sf1 a b pr2 m2) sf2 - -addDigits2 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) -addDigits2 m1 (One a) b c (One d) m2 = - appendTree2 m1 (node2 a b) (node2 c d) m2 -addDigits2 m1 (One a) b c (Two d e) m2 = - appendTree2 m1 (node3 a b c) (node2 d e) m2 -addDigits2 m1 (One a) b c (Three d e f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits2 m1 (One a) b c (Four d e f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits2 m1 (Two a b) c d (One e) m2 = - appendTree2 m1 (node3 a b c) (node2 d e) m2 -addDigits2 m1 (Two a b) c d (Two e f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits2 m1 (Two a b) c d (Three e f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits2 m1 (Two a b) c d (Four e f g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits2 m1 (Three a b c) d e (One f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits2 m1 (Three a b c) d e (Two f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits2 m1 (Three a b c) d e (Three f g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits2 m1 (Three a b c) d e (Four f g h i) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 -addDigits2 m1 (Four a b c d) e f (One g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits2 m1 (Four a b c d) e f (Two g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits2 m1 (Four a b c d) e f (Three g h i) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 -addDigits2 m1 (Four a b c d) e f (Four g h i j) m2 = - appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 - -appendTree3 :: FingerTree (Node a) -> Node a -> Node a -> Node a -> FingerTree (Node a) -> FingerTree (Node a) -appendTree3 Empty a b c xs = - a `consTree` b `consTree` c `consTree` xs -appendTree3 xs a b c Empty = - xs `snocTree` a `snocTree` b `snocTree` c -appendTree3 (Single x) a b c xs = - x `consTree` a `consTree` b `consTree` c `consTree` xs -appendTree3 xs a b c (Single x) = - xs `snocTree` a `snocTree` b `snocTree` c `snocTree` x -appendTree3 (Deep s1 pr1 m1 sf1) a b c (Deep s2 pr2 m2 sf2) = - Deep (s1 + size a + size b + size c + s2) pr1 (addDigits3 m1 sf1 a b c pr2 m2) sf2 - -addDigits3 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Node a -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) -addDigits3 m1 (One a) b c d (One e) m2 = - appendTree2 m1 (node3 a b c) (node2 d e) m2 -addDigits3 m1 (One a) b c d (Two e f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits3 m1 (One a) b c d (Three e f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits3 m1 (One a) b c d (Four e f g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits3 m1 (Two a b) c d e (One f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits3 m1 (Two a b) c d e (Two f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits3 m1 (Two a b) c d e (Three f g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits3 m1 (Two a b) c d e (Four f g h i) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 -addDigits3 m1 (Three a b c) d e f (One g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits3 m1 (Three a b c) d e f (Two g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits3 m1 (Three a b c) d e f (Three g h i) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 -addDigits3 m1 (Three a b c) d e f (Four g h i j) m2 = - appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 -addDigits3 m1 (Four a b c d) e f g (One h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits3 m1 (Four a b c d) e f g (Two h i) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 -addDigits3 m1 (Four a b c d) e f g (Three h i j) m2 = - appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 -addDigits3 m1 (Four a b c d) e f g (Four h i j k) m2 = - appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node2 j k) m2 - -appendTree4 :: FingerTree (Node a) -> Node a -> Node a -> Node a -> Node a -> FingerTree (Node a) -> FingerTree (Node a) -appendTree4 Empty a b c d xs = - a `consTree` b `consTree` c `consTree` d `consTree` xs -appendTree4 xs a b c d Empty = - xs `snocTree` a `snocTree` b `snocTree` c `snocTree` d -appendTree4 (Single x) a b c d xs = - x `consTree` a `consTree` b `consTree` c `consTree` d `consTree` xs -appendTree4 xs a b c d (Single x) = - xs `snocTree` a `snocTree` b `snocTree` c `snocTree` d `snocTree` x -appendTree4 (Deep s1 pr1 m1 sf1) a b c d (Deep s2 pr2 m2 sf2) = - Deep (s1 + size a + size b + size c + size d + s2) pr1 (addDigits4 m1 sf1 a b c d pr2 m2) sf2 - -addDigits4 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Node a -> Node a -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) -addDigits4 m1 (One a) b c d e (One f) m2 = - appendTree2 m1 (node3 a b c) (node3 d e f) m2 -addDigits4 m1 (One a) b c d e (Two f g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits4 m1 (One a) b c d e (Three f g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits4 m1 (One a) b c d e (Four f g h i) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 -addDigits4 m1 (Two a b) c d e f (One g) m2 = - appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 -addDigits4 m1 (Two a b) c d e f (Two g h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits4 m1 (Two a b) c d e f (Three g h i) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 -addDigits4 m1 (Two a b) c d e f (Four g h i j) m2 = - appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 -addDigits4 m1 (Three a b c) d e f g (One h) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 -addDigits4 m1 (Three a b c) d e f g (Two h i) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 -addDigits4 m1 (Three a b c) d e f g (Three h i j) m2 = - appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 -addDigits4 m1 (Three a b c) d e f g (Four h i j k) m2 = - appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node2 j k) m2 -addDigits4 m1 (Four a b c d) e f g h (One i) m2 = - appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 -addDigits4 m1 (Four a b c d) e f g h (Two i j) m2 = - appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 -addDigits4 m1 (Four a b c d) e f g h (Three i j k) m2 = - appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node2 j k) m2 -addDigits4 m1 (Four a b c d) e f g h (Four i j k l) m2 = - appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node3 j k l) m2 - --- | Builds a sequence from a seed value. Takes time linear in the --- number of generated elements. /WARNING:/ If the number of generated --- elements is infinite, this method will not terminate. -unfoldr :: (b -> Maybe (a, b)) -> b -> Seq a -unfoldr f = unfoldr' empty - -- uses tail recursion rather than, for instance, the List implementation. - where unfoldr' as b = maybe as (\ (a, b') -> unfoldr' (as |> a) b') (f b) - --- | @'unfoldl' f x@ is equivalent to @'reverse' ('unfoldr' ('fmap' swap . f) x)@. -unfoldl :: (b -> Maybe (b, a)) -> b -> Seq a -unfoldl f = unfoldl' empty - where unfoldl' as b = maybe as (\ (b', a) -> unfoldl' (a <| as) b') (f b) - --- | /O(n)/. Constructs a sequence by repeated application of a function --- to a seed value. --- --- > iterateN n f x = fromList (Prelude.take n (Prelude.iterate f x)) -iterateN :: Int -> (a -> a) -> a -> Seq a -iterateN n f x - | n >= 0 = replicateA n (State (\ y -> (f y, y))) `execState` x - | otherwise = error "iterateN takes a nonnegative integer argument" - ------------------------------------------------------------------------- --- Deconstruction ------------------------------------------------------------------------- - --- | /O(1)/. Is this the empty sequence? -null :: Seq a -> Bool -null (Seq Empty) = True -null _ = False - --- | /O(1)/. The number of elements in the sequence. -length :: Seq a -> Int -length (Seq xs) = size xs - --- Views - -data Maybe2 a b = Nothing2 | Just2 a b - --- | View of the left end of a sequence. -data ViewL a - = EmptyL -- ^ empty sequence - | a :< Seq a -- ^ leftmost element and the rest of the sequence -#if __GLASGOW_HASKELL__ - deriving (Eq, Ord, Show, Read, Data) -#else - deriving (Eq, Ord, Show, Read) -#endif - -INSTANCE_TYPEABLE1(ViewL,viewLTc,"ViewL") - -instance Functor ViewL where - {-# INLINE fmap #-} - fmap _ EmptyL = EmptyL - fmap f (x :< xs) = f x :< fmap f xs - -instance Foldable ViewL where - foldr _ z EmptyL = z - foldr f z (x :< xs) = f x (foldr f z xs) - - foldl _ z EmptyL = z - foldl f z (x :< xs) = foldl f (f z x) xs - - foldl1 _ EmptyL = error "foldl1: empty view" - foldl1 f (x :< xs) = foldl f x xs - -instance Traversable ViewL where - traverse _ EmptyL = pure EmptyL - traverse f (x :< xs) = (:<) <$> f x <*> traverse f xs - --- | /O(1)/. Analyse the left end of a sequence. -viewl :: Seq a -> ViewL a -viewl (Seq xs) = case viewLTree xs of - Nothing2 -> EmptyL - Just2 (Elem x) xs' -> x :< Seq xs' - -{-# SPECIALIZE viewLTree :: FingerTree (Elem a) -> Maybe2 (Elem a) (FingerTree (Elem a)) #-} -{-# SPECIALIZE viewLTree :: FingerTree (Node a) -> Maybe2 (Node a) (FingerTree (Node a)) #-} -viewLTree :: Sized a => FingerTree a -> Maybe2 a (FingerTree a) -viewLTree Empty = Nothing2 -viewLTree (Single a) = Just2 a Empty -viewLTree (Deep s (One a) m sf) = Just2 a (pullL (s - size a) m sf) -viewLTree (Deep s (Two a b) m sf) = - Just2 a (Deep (s - size a) (One b) m sf) -viewLTree (Deep s (Three a b c) m sf) = - Just2 a (Deep (s - size a) (Two b c) m sf) -viewLTree (Deep s (Four a b c d) m sf) = - Just2 a (Deep (s - size a) (Three b c d) m sf) - --- | View of the right end of a sequence. -data ViewR a - = EmptyR -- ^ empty sequence - | Seq a :> a -- ^ the sequence minus the rightmost element, - -- and the rightmost element -#if __GLASGOW_HASKELL__ - deriving (Eq, Ord, Show, Read, Data) -#else - deriving (Eq, Ord, Show, Read) -#endif - -INSTANCE_TYPEABLE1(ViewR,viewRTc,"ViewR") - -instance Functor ViewR where - {-# INLINE fmap #-} - fmap _ EmptyR = EmptyR - fmap f (xs :> x) = fmap f xs :> f x - -instance Foldable ViewR where - foldMap _ EmptyR = mempty - foldMap f (xs :> x) = foldMap f xs `mappend` f x - - foldr _ z EmptyR = z - foldr f z (xs :> x) = foldr f (f x z) xs - - foldl _ z EmptyR = z - foldl f z (xs :> x) = foldl f z xs `f` x - - foldr1 _ EmptyR = error "foldr1: empty view" - foldr1 f (xs :> x) = foldr f x xs -#if MIN_VERSION_base(4,8,0) - -- The default definitions are sensible for ViewL, but not so much for - -- ViewR. - null EmptyR = True - null (_ :> _) = False - - length = foldr' (\_ k -> k+1) 0 -#endif - -instance Traversable ViewR where - traverse _ EmptyR = pure EmptyR - traverse f (xs :> x) = (:>) <$> traverse f xs <*> f x - --- | /O(1)/. Analyse the right end of a sequence. -viewr :: Seq a -> ViewR a -viewr (Seq xs) = case viewRTree xs of - Nothing2 -> EmptyR - Just2 xs' (Elem x) -> Seq xs' :> x - -{-# SPECIALIZE viewRTree :: FingerTree (Elem a) -> Maybe2 (FingerTree (Elem a)) (Elem a) #-} -{-# SPECIALIZE viewRTree :: FingerTree (Node a) -> Maybe2 (FingerTree (Node a)) (Node a) #-} -viewRTree :: Sized a => FingerTree a -> Maybe2 (FingerTree a) a -viewRTree Empty = Nothing2 -viewRTree (Single z) = Just2 Empty z -viewRTree (Deep s pr m (One z)) = Just2 (pullR (s - size z) pr m) z -viewRTree (Deep s pr m (Two y z)) = - Just2 (Deep (s - size z) pr m (One y)) z -viewRTree (Deep s pr m (Three x y z)) = - Just2 (Deep (s - size z) pr m (Two x y)) z -viewRTree (Deep s pr m (Four w x y z)) = - Just2 (Deep (s - size z) pr m (Three w x y)) z - ------------------------------------------------------------------------- --- Scans --- --- These are not particularly complex applications of the Traversable --- functor, though making the correspondence with Data.List exact --- requires the use of (<|) and (|>). --- --- Note that save for the single (<|) or (|>), we maintain the original --- structure of the Seq, not having to do any restructuring of our own. --- --- wasserman.louis@gmail.com, 5/23/09 ------------------------------------------------------------------------- - --- | 'scanl' is similar to 'foldl', but returns a sequence of reduced --- values from the left: --- --- > scanl f z (fromList [x1, x2, ...]) = fromList [z, z `f` x1, (z `f` x1) `f` x2, ...] -scanl :: (a -> b -> a) -> a -> Seq b -> Seq a -scanl f z0 xs = z0 <| snd (mapAccumL (\ x z -> let x' = f x z in (x', x')) z0 xs) - --- | 'scanl1' is a variant of 'scanl' that has no starting value argument: --- --- > scanl1 f (fromList [x1, x2, ...]) = fromList [x1, x1 `f` x2, ...] -scanl1 :: (a -> a -> a) -> Seq a -> Seq a -scanl1 f xs = case viewl xs of - EmptyL -> error "scanl1 takes a nonempty sequence as an argument" - x :< xs' -> scanl f x xs' - --- | 'scanr' is the right-to-left dual of 'scanl'. -scanr :: (a -> b -> b) -> b -> Seq a -> Seq b -scanr f z0 xs = snd (mapAccumR (\ z x -> let z' = f x z in (z', z')) z0 xs) |> z0 - --- | 'scanr1' is a variant of 'scanr' that has no starting value argument. -scanr1 :: (a -> a -> a) -> Seq a -> Seq a -scanr1 f xs = case viewr xs of - EmptyR -> error "scanr1 takes a nonempty sequence as an argument" - xs' :> x -> scanr f x xs' - --- Indexing - --- | /O(log(min(i,n-i)))/. The element at the specified position, --- counting from 0. The argument should thus be a non-negative --- integer less than the size of the sequence. --- If the position is out of range, 'index' fails with an error. -index :: Seq a -> Int -> a -index (Seq xs) i - | 0 <= i && i < size xs = case lookupTree i xs of - Place _ (Elem x) -> x - | otherwise = error "index out of bounds" - -data Place a = Place {-# UNPACK #-} !Int a -#if TESTING - deriving Show -#endif - -{-# SPECIALIZE lookupTree :: Int -> FingerTree (Elem a) -> Place (Elem a) #-} -{-# SPECIALIZE lookupTree :: Int -> FingerTree (Node a) -> Place (Node a) #-} -lookupTree :: Sized a => Int -> FingerTree a -> Place a -lookupTree _ Empty = error "lookupTree of empty tree" -lookupTree i (Single x) = Place i x -lookupTree i (Deep totalSize pr m sf) - | i < spr = lookupDigit i pr - | i < spm = case lookupTree (i - spr) m of - Place i' xs -> lookupNode i' xs - | otherwise = lookupDigit (i - spm) sf - where - spr = size pr - spm = totalSize - size sf - -{-# SPECIALIZE lookupNode :: Int -> Node (Elem a) -> Place (Elem a) #-} -{-# SPECIALIZE lookupNode :: Int -> Node (Node a) -> Place (Node a) #-} -lookupNode :: Sized a => Int -> Node a -> Place a -lookupNode i (Node2 _ a b) - | i < sa = Place i a - | otherwise = Place (i - sa) b - where - sa = size a -lookupNode i (Node3 _ a b c) - | i < sa = Place i a - | i < sab = Place (i - sa) b - | otherwise = Place (i - sab) c - where - sa = size a - sab = sa + size b - -{-# SPECIALIZE lookupDigit :: Int -> Digit (Elem a) -> Place (Elem a) #-} -{-# SPECIALIZE lookupDigit :: Int -> Digit (Node a) -> Place (Node a) #-} -lookupDigit :: Sized a => Int -> Digit a -> Place a -lookupDigit i (One a) = Place i a -lookupDigit i (Two a b) - | i < sa = Place i a - | otherwise = Place (i - sa) b - where - sa = size a -lookupDigit i (Three a b c) - | i < sa = Place i a - | i < sab = Place (i - sa) b - | otherwise = Place (i - sab) c - where - sa = size a - sab = sa + size b -lookupDigit i (Four a b c d) - | i < sa = Place i a - | i < sab = Place (i - sa) b - | i < sabc = Place (i - sab) c - | otherwise = Place (i - sabc) d - where - sa = size a - sab = sa + size b - sabc = sab + size c - --- | /O(log(min(i,n-i)))/. Replace the element at the specified position. --- If the position is out of range, the original sequence is returned. -update :: Int -> a -> Seq a -> Seq a -update i x = adjust (const x) i - --- | /O(log(min(i,n-i)))/. Update the element at the specified position. --- If the position is out of range, the original sequence is returned. -adjust :: (a -> a) -> Int -> Seq a -> Seq a -adjust f i (Seq xs) - | 0 <= i && i < size xs = Seq (adjustTree (const (fmap f)) i xs) - | otherwise = Seq xs - -{-# SPECIALIZE adjustTree :: (Int -> Elem a -> Elem a) -> Int -> FingerTree (Elem a) -> FingerTree (Elem a) #-} -{-# SPECIALIZE adjustTree :: (Int -> Node a -> Node a) -> Int -> FingerTree (Node a) -> FingerTree (Node a) #-} -adjustTree :: Sized a => (Int -> a -> a) -> - Int -> FingerTree a -> FingerTree a -adjustTree _ _ Empty = error "adjustTree of empty tree" -adjustTree f i (Single x) = Single (f i x) -adjustTree f i (Deep s pr m sf) - | i < spr = Deep s (adjustDigit f i pr) m sf - | i < spm = Deep s pr (adjustTree (adjustNode f) (i - spr) m) sf - | otherwise = Deep s pr m (adjustDigit f (i - spm) sf) - where - spr = size pr - spm = spr + size m - -{-# SPECIALIZE adjustNode :: (Int -> Elem a -> Elem a) -> Int -> Node (Elem a) -> Node (Elem a) #-} -{-# SPECIALIZE adjustNode :: (Int -> Node a -> Node a) -> Int -> Node (Node a) -> Node (Node a) #-} -adjustNode :: Sized a => (Int -> a -> a) -> Int -> Node a -> Node a -adjustNode f i (Node2 s a b) - | i < sa = Node2 s (f i a) b - | otherwise = Node2 s a (f (i - sa) b) - where - sa = size a -adjustNode f i (Node3 s a b c) - | i < sa = Node3 s (f i a) b c - | i < sab = Node3 s a (f (i - sa) b) c - | otherwise = Node3 s a b (f (i - sab) c) - where - sa = size a - sab = sa + size b - -{-# SPECIALIZE adjustDigit :: (Int -> Elem a -> Elem a) -> Int -> Digit (Elem a) -> Digit (Elem a) #-} -{-# SPECIALIZE adjustDigit :: (Int -> Node a -> Node a) -> Int -> Digit (Node a) -> Digit (Node a) #-} -adjustDigit :: Sized a => (Int -> a -> a) -> Int -> Digit a -> Digit a -adjustDigit f i (One a) = One (f i a) -adjustDigit f i (Two a b) - | i < sa = Two (f i a) b - | otherwise = Two a (f (i - sa) b) - where - sa = size a -adjustDigit f i (Three a b c) - | i < sa = Three (f i a) b c - | i < sab = Three a (f (i - sa) b) c - | otherwise = Three a b (f (i - sab) c) - where - sa = size a - sab = sa + size b -adjustDigit f i (Four a b c d) - | i < sa = Four (f i a) b c d - | i < sab = Four a (f (i - sa) b) c d - | i < sabc = Four a b (f (i - sab) c) d - | otherwise = Four a b c (f (i- sabc) d) - where - sa = size a - sab = sa + size b - sabc = sab + size c - --- | /O(n)/. A generalization of 'fmap', 'mapWithIndex' takes a mapping --- function that also depends on the element's index, and applies it to every --- element in the sequence. -mapWithIndex :: (Int -> a -> b) -> Seq a -> Seq b -mapWithIndex f' (Seq xs') = Seq $ mapWithIndexTree (\s (Elem a) -> Elem (f' s a)) 0 xs' - where - {-# SPECIALIZE mapWithIndexTree :: (Int -> Elem y -> b) -> Int -> FingerTree (Elem y) -> FingerTree b #-} - {-# SPECIALIZE mapWithIndexTree :: (Int -> Node y -> b) -> Int -> FingerTree (Node y) -> FingerTree b #-} - mapWithIndexTree :: Sized a => (Int -> a -> b) -> Int -> FingerTree a -> FingerTree b - mapWithIndexTree _ s Empty = s `seq` Empty - mapWithIndexTree f s (Single xs) = Single $ f s xs - mapWithIndexTree f s (Deep n pr m sf) = sPspr `seq` sPsprm `seq` - Deep n - (mapWithIndexDigit f s pr) - (mapWithIndexTree (mapWithIndexNode f) sPspr m) - (mapWithIndexDigit f sPsprm sf) - where - sPspr = s + size pr - sPsprm = s + n - size sf - - {-# SPECIALIZE mapWithIndexDigit :: (Int -> Elem y -> b) -> Int -> Digit (Elem y) -> Digit b #-} - {-# SPECIALIZE mapWithIndexDigit :: (Int -> Node y -> b) -> Int -> Digit (Node y) -> Digit b #-} - mapWithIndexDigit :: Sized a => (Int -> a -> b) -> Int -> Digit a -> Digit b - mapWithIndexDigit f s (One a) = One (f s a) - mapWithIndexDigit f s (Two a b) = sPsa `seq` Two (f s a) (f sPsa b) - where - sPsa = s + size a - mapWithIndexDigit f s (Three a b c) = sPsa `seq` sPsab `seq` - Three (f s a) (f sPsa b) (f sPsab c) - where - sPsa = s + size a - sPsab = sPsa + size b - mapWithIndexDigit f s (Four a b c d) = sPsa `seq` sPsab `seq` sPsabc `seq` - Four (f s a) (f sPsa b) (f sPsab c) (f sPsabc d) - where - sPsa = s + size a - sPsab = sPsa + size b - sPsabc = sPsab + size c - - {-# SPECIALIZE mapWithIndexNode :: (Int -> Elem y -> b) -> Int -> Node (Elem y) -> Node b #-} - {-# SPECIALIZE mapWithIndexNode :: (Int -> Node y -> b) -> Int -> Node (Node y) -> Node b #-} - mapWithIndexNode :: Sized a => (Int -> a -> b) -> Int -> Node a -> Node b - mapWithIndexNode f s (Node2 ns a b) = sPsa `seq` Node2 ns (f s a) (f sPsa b) - where - sPsa = s + size a - mapWithIndexNode f s (Node3 ns a b c) = sPsa `seq` sPsab `seq` - Node3 ns (f s a) (f sPsa b) (f sPsab c) - where - sPsa = s + size a - sPsab = sPsa + size b - -#ifdef __GLASGOW_HASKELL__ -{-# NOINLINE [1] mapWithIndex #-} -{-# RULES -"mapWithIndex/mapWithIndex" forall f g xs . mapWithIndex f (mapWithIndex g xs) = - mapWithIndex (\k a -> f k (g k a)) xs -"mapWithIndex/fmapSeq" forall f g xs . mapWithIndex f (fmapSeq g xs) = - mapWithIndex (\k a -> f k (g a)) xs -"fmapSeq/mapWithIndex" forall f g xs . fmapSeq f (mapWithIndex g xs) = - mapWithIndex (\k a -> f (g k a)) xs - #-} -#endif - --- | /O(n)/. Convert a given sequence length and a function representing that --- sequence into a sequence. -fromFunction :: Int -> (Int -> a) -> Seq a -fromFunction len f | len < 0 = error "Data.Sequence.fromFunction called with negative len" - | len == 0 = empty - | otherwise = Seq $ create (lift_elem f) 1 0 len - where - create :: (Int -> a) -> Int -> Int -> Int -> FingerTree a - create b{-tree_builder-} s{-tree_size-} i{-start_index-} trees = i `seq` s `seq` case trees of - 1 -> Single $ b i - 2 -> Deep (2*s) (One (b i)) Empty (One (b (i+s))) - 3 -> Deep (3*s) (createTwo i) Empty (One (b (i+2*s))) - 4 -> Deep (4*s) (createTwo i) Empty (createTwo (i+2*s)) - 5 -> Deep (5*s) (createThree i) Empty (createTwo (i+3*s)) - 6 -> Deep (6*s) (createThree i) Empty (createThree (i+3*s)) - _ -> case trees `quotRem` 3 of - (trees', 1) -> Deep (trees*s) (createTwo i) - (create mb (3*s) (i+2*s) (trees'-1)) - (createTwo (i+(2+3*(trees'-1))*s)) - (trees', 2) -> Deep (trees*s) (createThree i) - (create mb (3*s) (i+3*s) (trees'-1)) - (createTwo (i+(3+3*(trees'-1))*s)) - (trees', _) -> Deep (trees*s) (createThree i) - (create mb (3*s) (i+3*s) (trees'-2)) - (createThree (i+(3+3*(trees'-2))*s)) - where - createTwo j = Two (b j) (b (j + s)) - {-# INLINE createTwo #-} - createThree j = Three (b j) (b (j + s)) (b (j + 2*s)) - {-# INLINE createThree #-} - mb j = Node3 (3*s) (b j) (b (j + s)) (b (j + 2*s)) - {-# INLINE mb #-} - - lift_elem :: (Int -> a) -> (Int -> Elem a) -#if __GLASGOW_HASKELL__ >= 708 - lift_elem g = coerce g -#else - lift_elem g = Elem . g -#endif - {-# INLINE lift_elem #-} - --- | /O(n)/. Create a sequence consisting of the elements of an 'Array'. --- Note that the resulting sequence elements may be evaluated lazily (as on GHC), --- so you must force the entire structure to be sure that the original array --- can be garbage-collected. -fromArray :: Ix i => Array i a -> Seq a -#ifdef __GLASGOW_HASKELL__ -fromArray a = fromFunction (GHC.Arr.numElements a) (GHC.Arr.unsafeAt a) - where - -- The following definition uses (Ix i) constraing, which is needed for the - -- other fromArray definition. - _ = Data.Array.rangeSize (Data.Array.bounds a) -#else -fromArray a = fromList2 (Data.Array.rangeSize (Data.Array.bounds a)) (Data.Array.elems a) -#endif - --- Splitting - --- | /O(log(min(i,n-i)))/. The first @i@ elements of a sequence. --- If @i@ is negative, @'take' i s@ yields the empty sequence. --- If the sequence contains fewer than @i@ elements, the whole sequence --- is returned. -take :: Int -> Seq a -> Seq a -take i = fst . splitAt' i - --- | /O(log(min(i,n-i)))/. Elements of a sequence after the first @i@. --- If @i@ is negative, @'drop' i s@ yields the whole sequence. --- If the sequence contains fewer than @i@ elements, the empty sequence --- is returned. -drop :: Int -> Seq a -> Seq a -drop i = snd . splitAt' i - --- | /O(log(min(i,n-i)))/. Split a sequence at a given position. --- @'splitAt' i s = ('take' i s, 'drop' i s)@. -splitAt :: Int -> Seq a -> (Seq a, Seq a) -splitAt i (Seq xs) = (Seq l, Seq r) - where (l, r) = split i xs - --- | /O(log(min(i,n-i))) A strict version of 'splitAt'. -splitAt' :: Int -> Seq a -> (Seq a, Seq a) -splitAt' i (Seq xs) = case split i xs of - (l, r) -> (Seq l, Seq r) - -split :: Int -> FingerTree (Elem a) -> - (FingerTree (Elem a), FingerTree (Elem a)) -split i Empty = i `seq` (Empty, Empty) -split i xs - | size xs > i = case splitTree i xs of - Split l x r -> (l, consTree x r) - | otherwise = (xs, Empty) - -data Split t a = Split t a t -#if TESTING - deriving Show -#endif - -{-# SPECIALIZE splitTree :: Int -> FingerTree (Elem a) -> Split (FingerTree (Elem a)) (Elem a) #-} -{-# SPECIALIZE splitTree :: Int -> FingerTree (Node a) -> Split (FingerTree (Node a)) (Node a) #-} -splitTree :: Sized a => Int -> FingerTree a -> Split (FingerTree a) a -splitTree _ Empty = error "splitTree of empty tree" -splitTree i (Single x) = i `seq` Split Empty x Empty -splitTree i (Deep _ pr m sf) - | i < spr = case splitDigit i pr of - Split l x r -> Split (maybe Empty digitToTree l) x (deepL r m sf) - | i < spm = case splitTree im m of - Split ml xs mr -> case splitNode (im - size ml) xs of - Split l x r -> Split (deepR pr ml l) x (deepL r mr sf) - | otherwise = case splitDigit (i - spm) sf of - Split l x r -> Split (deepR pr m l) x (maybe Empty digitToTree r) - where - spr = size pr - spm = spr + size m - im = i - spr - -{-# SPECIALIZE splitNode :: Int -> Node (Elem a) -> Split (Maybe (Digit (Elem a))) (Elem a) #-} -{-# SPECIALIZE splitNode :: Int -> Node (Node a) -> Split (Maybe (Digit (Node a))) (Node a) #-} -splitNode :: Sized a => Int -> Node a -> Split (Maybe (Digit a)) a -splitNode i (Node2 _ a b) - | i < sa = Split Nothing a (Just (One b)) - | otherwise = Split (Just (One a)) b Nothing - where - sa = size a -splitNode i (Node3 _ a b c) - | i < sa = Split Nothing a (Just (Two b c)) - | i < sab = Split (Just (One a)) b (Just (One c)) - | otherwise = Split (Just (Two a b)) c Nothing - where - sa = size a - sab = sa + size b - -{-# SPECIALIZE splitDigit :: Int -> Digit (Elem a) -> Split (Maybe (Digit (Elem a))) (Elem a) #-} -{-# SPECIALIZE splitDigit :: Int -> Digit (Node a) -> Split (Maybe (Digit (Node a))) (Node a) #-} -splitDigit :: Sized a => Int -> Digit a -> Split (Maybe (Digit a)) a -splitDigit i (One a) = i `seq` Split Nothing a Nothing -splitDigit i (Two a b) - | i < sa = Split Nothing a (Just (One b)) - | otherwise = Split (Just (One a)) b Nothing - where - sa = size a -splitDigit i (Three a b c) - | i < sa = Split Nothing a (Just (Two b c)) - | i < sab = Split (Just (One a)) b (Just (One c)) - | otherwise = Split (Just (Two a b)) c Nothing - where - sa = size a - sab = sa + size b -splitDigit i (Four a b c d) - | i < sa = Split Nothing a (Just (Three b c d)) - | i < sab = Split (Just (One a)) b (Just (Two c d)) - | i < sabc = Split (Just (Two a b)) c (Just (One d)) - | otherwise = Split (Just (Three a b c)) d Nothing - where - sa = size a - sab = sa + size b - sabc = sab + size c - --- | /O(n)/. Returns a sequence of all suffixes of this sequence, --- longest first. For example, --- --- > tails (fromList "abc") = fromList [fromList "abc", fromList "bc", fromList "c", fromList ""] --- --- Evaluating the /i/th suffix takes /O(log(min(i, n-i)))/, but evaluating --- every suffix in the sequence takes /O(n)/ due to sharing. -tails :: Seq a -> Seq (Seq a) -tails (Seq xs) = Seq (tailsTree (Elem . Seq) xs) |> empty - --- | /O(n)/. Returns a sequence of all prefixes of this sequence, --- shortest first. For example, --- --- > inits (fromList "abc") = fromList [fromList "", fromList "a", fromList "ab", fromList "abc"] --- --- Evaluating the /i/th prefix takes /O(log(min(i, n-i)))/, but evaluating --- every prefix in the sequence takes /O(n)/ due to sharing. -inits :: Seq a -> Seq (Seq a) -inits (Seq xs) = empty <| Seq (initsTree (Elem . Seq) xs) - --- This implementation of tails (and, analogously, inits) has the --- following algorithmic advantages: --- Evaluating each tail in the sequence takes linear total time, --- which is better than we could say for --- @fromList [drop n xs | n <- [0..length xs]]@. --- Evaluating any individual tail takes logarithmic time, which is --- better than we can say for either --- @scanr (<|) empty xs@ or @iterateN (length xs + 1) (\ xs -> let _ :< xs' = viewl xs in xs') xs@. --- --- Moreover, if we actually look at every tail in the sequence, the --- following benchmarks demonstrate that this implementation is modestly --- faster than any of the above: --- --- Times (ms) --- min mean +/-sd median max --- Seq.tails: 21.986 24.961 10.169 22.417 86.485 --- scanr: 85.392 87.942 2.488 87.425 100.217 --- iterateN: 29.952 31.245 1.574 30.412 37.268 --- --- The algorithm for tails (and, analogously, inits) is as follows: --- --- A Node in the FingerTree of tails is constructed by evaluating the --- corresponding tail of the FingerTree of Nodes, considering the first --- Node in this tail, and constructing a Node in which each tail of this --- Node is made to be the prefix of the remaining tree. This ends up --- working quite elegantly, as the remainder of the tail of the FingerTree --- of Nodes becomes the middle of a new tail, the suffix of the Node is --- the prefix, and the suffix of the original tree is retained. --- --- In particular, evaluating the /i/th tail involves making as --- many partial evaluations as the Node depth of the /i/th element. --- In addition, when we evaluate the /i/th tail, and we also evaluate --- the /j/th tail, and /m/ Nodes are on the path to both /i/ and /j/, --- each of those /m/ evaluations are shared between the computation of --- the /i/th and /j/th tails. --- --- wasserman.louis@gmail.com, 7/16/09 - -tailsDigit :: Digit a -> Digit (Digit a) -tailsDigit (One a) = One (One a) -tailsDigit (Two a b) = Two (Two a b) (One b) -tailsDigit (Three a b c) = Three (Three a b c) (Two b c) (One c) -tailsDigit (Four a b c d) = Four (Four a b c d) (Three b c d) (Two c d) (One d) - -initsDigit :: Digit a -> Digit (Digit a) -initsDigit (One a) = One (One a) -initsDigit (Two a b) = Two (One a) (Two a b) -initsDigit (Three a b c) = Three (One a) (Two a b) (Three a b c) -initsDigit (Four a b c d) = Four (One a) (Two a b) (Three a b c) (Four a b c d) - -tailsNode :: Node a -> Node (Digit a) -tailsNode (Node2 s a b) = Node2 s (Two a b) (One b) -tailsNode (Node3 s a b c) = Node3 s (Three a b c) (Two b c) (One c) - -initsNode :: Node a -> Node (Digit a) -initsNode (Node2 s a b) = Node2 s (One a) (Two a b) -initsNode (Node3 s a b c) = Node3 s (One a) (Two a b) (Three a b c) - -{-# SPECIALIZE tailsTree :: (FingerTree (Elem a) -> Elem b) -> FingerTree (Elem a) -> FingerTree (Elem b) #-} -{-# SPECIALIZE tailsTree :: (FingerTree (Node a) -> Node b) -> FingerTree (Node a) -> FingerTree (Node b) #-} --- | Given a function to apply to tails of a tree, applies that function --- to every tail of the specified tree. -tailsTree :: Sized a => (FingerTree a -> b) -> FingerTree a -> FingerTree b -tailsTree _ Empty = Empty -tailsTree f (Single x) = Single (f (Single x)) -tailsTree f (Deep n pr m sf) = - Deep n (fmap (\ pr' -> f (deep pr' m sf)) (tailsDigit pr)) - (tailsTree f' m) - (fmap (f . digitToTree) (tailsDigit sf)) - where - f' ms = let Just2 node m' = viewLTree ms in - fmap (\ pr' -> f (deep pr' m' sf)) (tailsNode node) - -{-# SPECIALIZE initsTree :: (FingerTree (Elem a) -> Elem b) -> FingerTree (Elem a) -> FingerTree (Elem b) #-} -{-# SPECIALIZE initsTree :: (FingerTree (Node a) -> Node b) -> FingerTree (Node a) -> FingerTree (Node b) #-} --- | Given a function to apply to inits of a tree, applies that function --- to every init of the specified tree. -initsTree :: Sized a => (FingerTree a -> b) -> FingerTree a -> FingerTree b -initsTree _ Empty = Empty -initsTree f (Single x) = Single (f (Single x)) -initsTree f (Deep n pr m sf) = - Deep n (fmap (f . digitToTree) (initsDigit pr)) - (initsTree f' m) - (fmap (f . deep pr m) (initsDigit sf)) - where - f' ms = let Just2 m' node = viewRTree ms in - fmap (\ sf' -> f (deep pr m' sf')) (initsNode node) - -{-# INLINE foldlWithIndex #-} --- | 'foldlWithIndex' is a version of 'foldl' that also provides access --- to the index of each element. -foldlWithIndex :: (b -> Int -> a -> b) -> b -> Seq a -> b -foldlWithIndex f z xs = foldl (\ g x i -> i `seq` f (g (i - 1)) i x) (const z) xs (length xs - 1) - -{-# INLINE foldrWithIndex #-} --- | 'foldrWithIndex' is a version of 'foldr' that also provides access --- to the index of each element. -foldrWithIndex :: (Int -> a -> b -> b) -> b -> Seq a -> b -foldrWithIndex f z xs = foldr (\ x g i -> i `seq` f i x (g (i+1))) (const z) xs 0 - -{-# INLINE listToMaybe' #-} --- 'listToMaybe\'' is a good consumer version of 'listToMaybe'. -listToMaybe' :: [a] -> Maybe a -listToMaybe' = foldr (\ x _ -> Just x) Nothing - --- | /O(i)/ where /i/ is the prefix length. 'takeWhileL', applied --- to a predicate @p@ and a sequence @xs@, returns the longest prefix --- (possibly empty) of @xs@ of elements that satisfy @p@. -takeWhileL :: (a -> Bool) -> Seq a -> Seq a -takeWhileL p = fst . spanl p - --- | /O(i)/ where /i/ is the suffix length. 'takeWhileR', applied --- to a predicate @p@ and a sequence @xs@, returns the longest suffix --- (possibly empty) of @xs@ of elements that satisfy @p@. --- --- @'takeWhileR' p xs@ is equivalent to @'reverse' ('takeWhileL' p ('reverse' xs))@. -takeWhileR :: (a -> Bool) -> Seq a -> Seq a -takeWhileR p = fst . spanr p - --- | /O(i)/ where /i/ is the prefix length. @'dropWhileL' p xs@ returns --- the suffix remaining after @'takeWhileL' p xs@. -dropWhileL :: (a -> Bool) -> Seq a -> Seq a -dropWhileL p = snd . spanl p - --- | /O(i)/ where /i/ is the suffix length. @'dropWhileR' p xs@ returns --- the prefix remaining after @'takeWhileR' p xs@. --- --- @'dropWhileR' p xs@ is equivalent to @'reverse' ('dropWhileL' p ('reverse' xs))@. -dropWhileR :: (a -> Bool) -> Seq a -> Seq a -dropWhileR p = snd . spanr p - --- | /O(i)/ where /i/ is the prefix length. 'spanl', applied to --- a predicate @p@ and a sequence @xs@, returns a pair whose first --- element is the longest prefix (possibly empty) of @xs@ of elements that --- satisfy @p@ and the second element is the remainder of the sequence. -spanl :: (a -> Bool) -> Seq a -> (Seq a, Seq a) -spanl p = breakl (not . p) - --- | /O(i)/ where /i/ is the suffix length. 'spanr', applied to a --- predicate @p@ and a sequence @xs@, returns a pair whose /first/ element --- is the longest /suffix/ (possibly empty) of @xs@ of elements that --- satisfy @p@ and the second element is the remainder of the sequence. -spanr :: (a -> Bool) -> Seq a -> (Seq a, Seq a) -spanr p = breakr (not . p) - -{-# INLINE breakl #-} --- | /O(i)/ where /i/ is the breakpoint index. 'breakl', applied to a --- predicate @p@ and a sequence @xs@, returns a pair whose first element --- is the longest prefix (possibly empty) of @xs@ of elements that --- /do not satisfy/ @p@ and the second element is the remainder of --- the sequence. --- --- @'breakl' p@ is equivalent to @'spanl' (not . p)@. -breakl :: (a -> Bool) -> Seq a -> (Seq a, Seq a) -breakl p xs = foldr (\ i _ -> splitAt i xs) (xs, empty) (findIndicesL p xs) - -{-# INLINE breakr #-} --- | @'breakr' p@ is equivalent to @'spanr' (not . p)@. -breakr :: (a -> Bool) -> Seq a -> (Seq a, Seq a) -breakr p xs = foldr (\ i _ -> flipPair (splitAt (i + 1) xs)) (xs, empty) (findIndicesR p xs) - where flipPair (x, y) = (y, x) - --- | /O(n)/. The 'partition' function takes a predicate @p@ and a --- sequence @xs@ and returns sequences of those elements which do and --- do not satisfy the predicate. -partition :: (a -> Bool) -> Seq a -> (Seq a, Seq a) -partition p = foldl part (empty, empty) - where - part (xs, ys) x - | p x = (xs |> x, ys) - | otherwise = (xs, ys |> x) - --- | /O(n)/. The 'filter' function takes a predicate @p@ and a sequence --- @xs@ and returns a sequence of those elements which satisfy the --- predicate. -filter :: (a -> Bool) -> Seq a -> Seq a -filter p = foldl (\ xs x -> if p x then xs |> x else xs) empty - --- Indexing sequences - --- | 'elemIndexL' finds the leftmost index of the specified element, --- if it is present, and otherwise 'Nothing'. -elemIndexL :: Eq a => a -> Seq a -> Maybe Int -elemIndexL x = findIndexL (x ==) - --- | 'elemIndexR' finds the rightmost index of the specified element, --- if it is present, and otherwise 'Nothing'. -elemIndexR :: Eq a => a -> Seq a -> Maybe Int -elemIndexR x = findIndexR (x ==) - --- | 'elemIndicesL' finds the indices of the specified element, from --- left to right (i.e. in ascending order). -elemIndicesL :: Eq a => a -> Seq a -> [Int] -elemIndicesL x = findIndicesL (x ==) - --- | 'elemIndicesR' finds the indices of the specified element, from --- right to left (i.e. in descending order). -elemIndicesR :: Eq a => a -> Seq a -> [Int] -elemIndicesR x = findIndicesR (x ==) - --- | @'findIndexL' p xs@ finds the index of the leftmost element that --- satisfies @p@, if any exist. -findIndexL :: (a -> Bool) -> Seq a -> Maybe Int -findIndexL p = listToMaybe' . findIndicesL p - --- | @'findIndexR' p xs@ finds the index of the rightmost element that --- satisfies @p@, if any exist. -findIndexR :: (a -> Bool) -> Seq a -> Maybe Int -findIndexR p = listToMaybe' . findIndicesR p - -{-# INLINE findIndicesL #-} --- | @'findIndicesL' p@ finds all indices of elements that satisfy @p@, --- in ascending order. -findIndicesL :: (a -> Bool) -> Seq a -> [Int] -#if __GLASGOW_HASKELL__ -findIndicesL p xs = build (\ c n -> let g i x z = if p x then c i z else z in - foldrWithIndex g n xs) -#else -findIndicesL p xs = foldrWithIndex g [] xs - where g i x is = if p x then i:is else is -#endif - -{-# INLINE findIndicesR #-} --- | @'findIndicesR' p@ finds all indices of elements that satisfy @p@, --- in descending order. -findIndicesR :: (a -> Bool) -> Seq a -> [Int] -#if __GLASGOW_HASKELL__ -findIndicesR p xs = build (\ c n -> - let g z i x = if p x then c i z else z in foldlWithIndex g n xs) -#else -findIndicesR p xs = foldlWithIndex g [] xs - where g is i x = if p x then i:is else is -#endif - ------------------------------------------------------------------------- --- Lists ------------------------------------------------------------------------- - --- The implementation below, by Ross Paterson, avoids the rebuilding --- the previous (|>)-based implementation suffered from. - --- | /O(n)/. Create a sequence from a finite list of elements. --- There is a function 'toList' in the opposite direction for all --- instances of the 'Foldable' class, including 'Seq'. -fromList :: [a] -> Seq a -fromList = Seq . mkTree 1 . map_elem - where - {-# SPECIALIZE mkTree :: Int -> [Elem a] -> FingerTree (Elem a) #-} - {-# SPECIALIZE mkTree :: Int -> [Node a] -> FingerTree (Node a) #-} - mkTree :: (Sized a) => Int -> [a] -> FingerTree a - STRICT_1_OF_2(mkTree) - mkTree _ [] = Empty - mkTree _ [x1] = Single x1 - mkTree s [x1, x2] = Deep (2*s) (One x1) Empty (One x2) - mkTree s [x1, x2, x3] = Deep (3*s) (One x1) Empty (Two x2 x3) - mkTree s (x1:x2:x3:x4:xs) = case getNodes (3*s) x4 xs of - (ns, sf) -> case mkTree (3*s) ns of - m -> m `seq` Deep (3*size x1 + size m + size sf) (Three x1 x2 x3) m sf - - getNodes :: Int -> a -> [a] -> ([Node a], Digit a) - STRICT_1_OF_3(getNodes) - getNodes _ x1 [] = ([], One x1) - getNodes _ x1 [x2] = ([], Two x1 x2) - getNodes _ x1 [x2, x3] = ([], Three x1 x2 x3) - getNodes s x1 (x2:x3:x4:xs) = (Node3 s x1 x2 x3:ns, d) - where (ns, d) = getNodes s x4 xs - - map_elem :: [a] -> [Elem a] -#if __GLASGOW_HASKELL__ >= 708 - map_elem xs = coerce xs -#else - map_elem xs = Data.List.map Elem xs -#endif - {-# INLINE map_elem #-} - -#if __GLASGOW_HASKELL__ >= 708 -instance GHC.Exts.IsList (Seq a) where - type Item (Seq a) = a - fromList = fromList - fromListN = fromList2 - toList = toList -#endif - -#ifdef __GLASGOW_HASKELL__ -instance IsString (Seq Char) where - fromString = fromList -#endif - ------------------------------------------------------------------------- --- Reverse ------------------------------------------------------------------------- - --- | /O(n)/. The reverse of a sequence. -reverse :: Seq a -> Seq a -reverse (Seq xs) = Seq (reverseTree id xs) - -reverseTree :: (a -> a) -> FingerTree a -> FingerTree a -reverseTree _ Empty = Empty -reverseTree f (Single x) = Single (f x) -reverseTree f (Deep s pr m sf) = - Deep s (reverseDigit f sf) - (reverseTree (reverseNode f) m) - (reverseDigit f pr) - -{-# INLINE reverseDigit #-} -reverseDigit :: (a -> a) -> Digit a -> Digit a -reverseDigit f (One a) = One (f a) -reverseDigit f (Two a b) = Two (f b) (f a) -reverseDigit f (Three a b c) = Three (f c) (f b) (f a) -reverseDigit f (Four a b c d) = Four (f d) (f c) (f b) (f a) - -reverseNode :: (a -> a) -> Node a -> Node a -reverseNode f (Node2 s a b) = Node2 s (f b) (f a) -reverseNode f (Node3 s a b c) = Node3 s (f c) (f b) (f a) - ------------------------------------------------------------------------- --- Mapping with a splittable value ------------------------------------------------------------------------- - --- For zipping, it is useful to build a result by --- traversing a sequence while splitting up something else. For zipping, we --- traverse the first sequence while splitting up the second. --- --- What makes all this crazy code a good idea: --- --- Suppose we zip together two sequences of the same length: --- --- zs = zip xs ys --- --- We want to get reasonably fast indexing into zs immediately, rather than --- needing to construct the entire thing first, as the previous implementation --- required. The first aspect is that we build the result "outside-in" or --- "top-down", rather than left to right. That gives us access to both ends --- quickly. But that's not enough, by itself, to give immediate access to the --- center of zs. For that, we need to be able to skip over larger segments of --- zs, delaying their construction until we actually need them. The way we do --- this is to traverse xs, while splitting up ys according to the structure of --- xs. If we have a Deep _ pr m sf, we split ys into three pieces, and hand off --- one piece to the prefix, one to the middle, and one to the suffix of the --- result. The key point is that we don't need to actually do anything further --- with those pieces until we actually need them; the computations to split --- them up further and zip them with their matching pieces can be delayed until --- they're actually needed. We do the same thing for Digits (splitting into --- between one and four pieces) and Nodes (splitting into two or three). The --- ultimate result is that we can index into, or split at, any location in zs --- in polylogarithmic time *immediately*, while still being able to force all --- the thunks in O(n) time. --- --- Benchmark info, and alternatives: --- --- The old zipping code used mapAccumL to traverse the first sequence while --- cutting down the second sequence one piece at a time. --- --- An alternative way to express that basic idea is to convert both sequences --- to lists, zip the lists, and then convert the result back to a sequence. --- I'll call this the "listy" implementation. --- --- I benchmarked two operations: Each started by zipping two sequences --- constructed with replicate and/or fromList. The first would then immediately --- index into the result. The second would apply deepseq to force the entire --- result. The new implementation worked much better than either of the others --- on the immediate indexing test, as expected. It also worked better than the --- old implementation for all the deepseq tests. For short sequences, the listy --- implementation outperformed all the others on the deepseq test. However, the --- splitting implementation caught up and surpassed it once the sequences grew --- long enough. It seems likely that by avoiding rebuilding, it interacts --- better with the cache hierarchy. --- --- David Feuer, with excellent guidance from Carter Schonwald, December 2014 - --- | /O(n)/. Constructs a new sequence with the same structure as an existing --- sequence using a user-supplied mapping function along with a splittable --- value and a way to split it. The value is split up lazily according to the --- structure of the sequence, so one piece of the value is distributed to each --- element of the sequence. The caller should provide a splitter function that --- takes a number, @n@, and a splittable value, breaks off a chunk of size @n@ --- from the value, and returns that chunk and the remainder as a pair. The --- following examples will hopefully make the usage clear: --- --- > zipWith :: (a -> b -> c) -> Seq a -> Seq b -> Seq c --- > zipWith f s1 s2 = splitMap splitAt (\b a -> f a (b `index` 0)) s2' s1' --- > where --- > minLen = min (length s1) (length s2) --- > s1' = take minLen s1 --- > s2' = take minLen s2 --- --- > mapWithIndex :: (Int -> a -> b) -> Seq a -> Seq b --- > mapWithIndex f = splitMap (\n i -> (i, n+i)) f 0 -splitMap :: (Int -> s -> (s,s)) -> (s -> a -> b) -> s -> Seq a -> Seq b -splitMap splt' = go - where - go f s (Seq xs) = Seq $ splitMapTree splt' (\s' (Elem a) -> Elem (f s' a)) s xs - - {-# SPECIALIZE splitMapTree :: (Int -> s -> (s,s)) -> (s -> Elem y -> b) -> s -> FingerTree (Elem y) -> FingerTree b #-} - {-# SPECIALIZE splitMapTree :: (Int -> s -> (s,s)) -> (s -> Node y -> b) -> s -> FingerTree (Node y) -> FingerTree b #-} - splitMapTree :: Sized a => (Int -> s -> (s,s)) -> (s -> a -> b) -> s -> FingerTree a -> FingerTree b - splitMapTree _ _ _ Empty = Empty - splitMapTree _ f s (Single xs) = Single $ f s xs - splitMapTree splt f s (Deep n pr m sf) = Deep n (splitMapDigit splt f prs pr) (splitMapTree splt (splitMapNode splt f) ms m) (splitMapDigit splt f sfs sf) - where - (prs, r) = splt (size pr) s - (ms, sfs) = splt (n - size pr - size sf) r - - {-# SPECIALIZE splitMapDigit :: (Int -> s -> (s,s)) -> (s -> Elem y -> b) -> s -> Digit (Elem y) -> Digit b #-} - {-# SPECIALIZE splitMapDigit :: (Int -> s -> (s,s)) -> (s -> Node y -> b) -> s -> Digit (Node y) -> Digit b #-} - splitMapDigit :: Sized a => (Int -> s -> (s,s)) -> (s -> a -> b) -> s -> Digit a -> Digit b - splitMapDigit _ f s (One a) = One (f s a) - splitMapDigit splt f s (Two a b) = Two (f first a) (f second b) - where - (first, second) = splt (size a) s - splitMapDigit splt f s (Three a b c) = Three (f first a) (f second b) (f third c) - where - (first, r) = splt (size a) s - (second, third) = splt (size b) r - splitMapDigit splt f s (Four a b c d) = Four (f first a) (f second b) (f third c) (f fourth d) - where - (first, s') = splt (size a) s - (middle, fourth) = splt (size b + size c) s' - (second, third) = splt (size b) middle - - {-# SPECIALIZE splitMapNode :: (Int -> s -> (s,s)) -> (s -> Elem y -> b) -> s -> Node (Elem y) -> Node b #-} - {-# SPECIALIZE splitMapNode :: (Int -> s -> (s,s)) -> (s -> Node y -> b) -> s -> Node (Node y) -> Node b #-} - splitMapNode :: Sized a => (Int -> s -> (s,s)) -> (s -> a -> b) -> s -> Node a -> Node b - splitMapNode splt f s (Node2 ns a b) = Node2 ns (f first a) (f second b) - where - (first, second) = splt (size a) s - splitMapNode splt f s (Node3 ns a b c) = Node3 ns (f first a) (f second b) (f third c) - where - (first, r) = splt (size a) s - (second, third) = splt (size b) r - -{-# INLINE splitMap #-} - -getSingleton :: Seq a -> a -getSingleton (Seq (Single (Elem a))) = a -getSingleton (Seq Empty) = error "getSingleton: Empty" -getSingleton _ = error "getSingleton: Not a singleton." - ------------------------------------------------------------------------- --- Zipping ------------------------------------------------------------------------- - --- | /O(min(n1,n2))/. 'zip' takes two sequences and returns a sequence --- of corresponding pairs. If one input is short, excess elements are --- discarded from the right end of the longer sequence. -zip :: Seq a -> Seq b -> Seq (a, b) -zip = zipWith (,) - --- | /O(min(n1,n2))/. 'zipWith' generalizes 'zip' by zipping with the --- function given as the first argument, instead of a tupling function. --- For example, @zipWith (+)@ is applied to two sequences to take the --- sequence of corresponding sums. -zipWith :: (a -> b -> c) -> Seq a -> Seq b -> Seq c -zipWith f s1 s2 = zipWith' f s1' s2' - where - minLen = min (length s1) (length s2) - s1' = take minLen s1 - s2' = take minLen s2 - --- | A version of zipWith that assumes the sequences have the same length. -zipWith' :: (a -> b -> c) -> Seq a -> Seq b -> Seq c -zipWith' f s1 s2 = splitMap splitAt' (\s a -> f a (getSingleton s)) s2 s1 - --- | /O(min(n1,n2,n3))/. 'zip3' takes three sequences and returns a --- sequence of triples, analogous to 'zip'. -zip3 :: Seq a -> Seq b -> Seq c -> Seq (a,b,c) -zip3 = zipWith3 (,,) - --- | /O(min(n1,n2,n3))/. 'zipWith3' takes a function which combines --- three elements, as well as three sequences and returns a sequence of --- their point-wise combinations, analogous to 'zipWith'. -zipWith3 :: (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d -zipWith3 f s1 s2 s3 = zipWith' ($) (zipWith' f s1' s2') s3' - where - minLen = minimum [length s1, length s2, length s3] - s1' = take minLen s1 - s2' = take minLen s2 - s3' = take minLen s3 - -zipWith3' :: (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d -zipWith3' f s1 s2 s3 = zipWith' ($) (zipWith' f s1 s2) s3 - --- | /O(min(n1,n2,n3,n4))/. 'zip4' takes four sequences and returns a --- sequence of quadruples, analogous to 'zip'. -zip4 :: Seq a -> Seq b -> Seq c -> Seq d -> Seq (a,b,c,d) -zip4 = zipWith4 (,,,) - --- | /O(min(n1,n2,n3,n4))/. 'zipWith4' takes a function which combines --- four elements, as well as four sequences and returns a sequence of --- their point-wise combinations, analogous to 'zipWith'. -zipWith4 :: (a -> b -> c -> d -> e) -> Seq a -> Seq b -> Seq c -> Seq d -> Seq e -zipWith4 f s1 s2 s3 s4 = zipWith' ($) (zipWith3' f s1' s2' s3') s4' - where - minLen = minimum [length s1, length s2, length s3, length s4] - s1' = take minLen s1 - s2' = take minLen s2 - s3' = take minLen s3 - s4' = take minLen s4 - ------------------------------------------------------------------------- --- Sorting --- --- sort and sortBy are implemented by simple deforestations of --- \ xs -> fromList2 (length xs) . Data.List.sortBy cmp . toList --- which does not get deforested automatically, it would appear. --- --- Unstable sorting is performed by a heap sort implementation based on --- pairing heaps. Because the internal structure of sequences is quite --- varied, it is difficult to get blocks of elements of roughly the same --- length, which would improve merge sort performance. Pairing heaps, --- on the other hand, are relatively resistant to the effects of merging --- heaps of wildly different sizes, as guaranteed by its amortized --- constant-time merge operation. Moreover, extensive use of SpecConstr --- transformations can be done on pairing heaps, especially when we're --- only constructing them to immediately be unrolled. --- --- On purely random sequences of length 50000, with no RTS options, --- I get the following statistics, in which heapsort is about 42.5% --- faster: (all comparisons done with -O2) --- --- Times (ms) min mean +/-sd median max --- to/from list: 103.802 108.572 7.487 106.436 143.339 --- unstable heapsort: 60.686 62.968 4.275 61.187 79.151 --- --- Heapsort, it would seem, is less of a memory hog than Data.List.sortBy. --- The gap is narrowed when more memory is available, but heapsort still --- wins, 15% faster, with +RTS -H128m: --- --- Times (ms) min mean +/-sd median max --- to/from list: 42.692 45.074 2.596 44.600 56.601 --- unstable heapsort: 37.100 38.344 3.043 37.715 55.526 --- --- In addition, on strictly increasing sequences the gap is even wider --- than normal; heapsort is 68.5% faster with no RTS options: --- Times (ms) min mean +/-sd median max --- to/from list: 52.236 53.574 1.987 53.034 62.098 --- unstable heapsort: 16.433 16.919 0.931 16.681 21.622 --- --- This may be attributed to the elegant nature of the pairing heap. --- --- wasserman.louis@gmail.com, 7/20/09 ------------------------------------------------------------------------- - --- | /O(n log n)/. 'sort' sorts the specified 'Seq' by the natural --- ordering of its elements. The sort is stable. --- If stability is not required, 'unstableSort' can be considerably --- faster, and in particular uses less memory. -sort :: Ord a => Seq a -> Seq a -sort = sortBy compare - --- | /O(n log n)/. 'sortBy' sorts the specified 'Seq' according to the --- specified comparator. The sort is stable. --- If stability is not required, 'unstableSortBy' can be considerably --- faster, and in particular uses less memory. -sortBy :: (a -> a -> Ordering) -> Seq a -> Seq a -sortBy cmp xs = fromList2 (length xs) (Data.List.sortBy cmp (toList xs)) - --- | /O(n log n)/. 'unstableSort' sorts the specified 'Seq' by --- the natural ordering of its elements, but the sort is not stable. --- This algorithm is frequently faster and uses less memory than 'sort', --- and performs extremely well -- frequently twice as fast as 'sort' -- --- when the sequence is already nearly sorted. -unstableSort :: Ord a => Seq a -> Seq a -unstableSort = unstableSortBy compare - --- | /O(n log n)/. A generalization of 'unstableSort', 'unstableSortBy' --- takes an arbitrary comparator and sorts the specified sequence. --- The sort is not stable. This algorithm is frequently faster and --- uses less memory than 'sortBy', and performs extremely well -- --- frequently twice as fast as 'sortBy' -- when the sequence is already --- nearly sorted. -unstableSortBy :: (a -> a -> Ordering) -> Seq a -> Seq a -unstableSortBy cmp (Seq xs) = - fromList2 (size xs) $ maybe [] (unrollPQ cmp) $ - toPQ cmp (\ (Elem x) -> PQueue x Nil) xs - --- | fromList2, given a list and its length, constructs a completely --- balanced Seq whose elements are that list using the replicateA --- generalization. -fromList2 :: Int -> [a] -> Seq a -fromList2 n = execState (replicateA n (State ht)) - where - ht (x:xs) = (xs, x) - ht [] = error "fromList2: short list" - --- | A 'PQueue' is a simple pairing heap. -data PQueue e = PQueue e (PQL e) -data PQL e = Nil | {-# UNPACK #-} !(PQueue e) :& PQL e - -infixr 8 :& - -#if TESTING - -instance Functor PQueue where - fmap f (PQueue x ts) = PQueue (f x) (fmap f ts) - -instance Functor PQL where - fmap f (q :& qs) = fmap f q :& fmap f qs - fmap _ Nil = Nil - -instance Show e => Show (PQueue e) where - show = unlines . draw . fmap show - --- borrowed wholesale from Data.Tree, as Data.Tree actually depends --- on Data.Sequence -draw :: PQueue String -> [String] -draw (PQueue x ts0) = x : drawSubTrees ts0 - where - drawSubTrees Nil = [] - drawSubTrees (t :& Nil) = - "|" : shift "`- " " " (draw t) - drawSubTrees (t :& ts) = - "|" : shift "+- " "| " (draw t) ++ drawSubTrees ts - - shift first other = Data.List.zipWith (++) (first : repeat other) -#endif - --- | 'unrollPQ', given a comparator function, unrolls a 'PQueue' into --- a sorted list. -unrollPQ :: (e -> e -> Ordering) -> PQueue e -> [e] -unrollPQ cmp = unrollPQ' - where - {-# INLINE unrollPQ' #-} - unrollPQ' (PQueue x ts) = x:mergePQs0 ts - (<+>) = mergePQ cmp - mergePQs0 Nil = [] - mergePQs0 (t :& Nil) = unrollPQ' t - mergePQs0 (t1 :& t2 :& ts) = mergePQs (t1 <+> t2) ts - mergePQs t ts = t `seq` case ts of - Nil -> unrollPQ' t - t1 :& Nil -> unrollPQ' (t <+> t1) - t1 :& t2 :& ts' -> mergePQs (t <+> (t1 <+> t2)) ts' - --- | 'toPQ', given an ordering function and a mechanism for queueifying --- elements, converts a 'FingerTree' to a 'PQueue'. -toPQ :: (e -> e -> Ordering) -> (a -> PQueue e) -> FingerTree a -> Maybe (PQueue e) -toPQ _ _ Empty = Nothing -toPQ _ f (Single x) = Just (f x) -toPQ cmp f (Deep _ pr m sf) = Just (maybe (pr' <+> sf') ((pr' <+> sf') <+>) (toPQ cmp fNode m)) - where - fDigit digit = case fmap f digit of - One a -> a - Two a b -> a <+> b - Three a b c -> a <+> b <+> c - Four a b c d -> (a <+> b) <+> (c <+> d) - (<+>) = mergePQ cmp - fNode = fDigit . nodeToDigit - pr' = fDigit pr - sf' = fDigit sf - --- | 'mergePQ' merges two 'PQueue's. -mergePQ :: (a -> a -> Ordering) -> PQueue a -> PQueue a -> PQueue a -mergePQ cmp q1@(PQueue x1 ts1) q2@(PQueue x2 ts2) - | cmp x1 x2 == GT = PQueue x2 (q1 :& ts2) - | otherwise = PQueue x1 (q2 :& ts1) + +#include "containers.h" + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Sequence +-- Copyright : (c) Ross Paterson 2005 +-- (c) Louis Wasserman 2009 +-- (c) Bertram Felgenhauer, David Feuer, Ross Paterson, and +-- Milan Straka 2014 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : experimental +-- Portability : portable +-- +-- General purpose finite sequences. +-- Apart from being finite and having strict operations, sequences +-- also differ from lists in supporting a wider variety of operations +-- efficiently. +-- +-- An amortized running time is given for each operation, with /n/ referring +-- to the length of the sequence and /i/ being the integral index used by +-- some operations. These bounds hold even in a persistent (shared) setting. +-- +-- The implementation uses 2-3 finger trees annotated with sizes, +-- as described in section 4.2 of +-- +-- * Ralf Hinze and Ross Paterson, +-- \"Finger trees: a simple general-purpose data structure\", +-- /Journal of Functional Programming/ 16:2 (2006) pp 197-217. +-- +-- +-- /Note/: Many of these operations have the same names as similar +-- operations on lists in the "Prelude". The ambiguity may be resolved +-- using either qualification or the @hiding@ clause. +-- +-- /Warning/: The size of a 'Seq' must not exceed @maxBound::Int@. Violation +-- of this condition is not detected and if the size limit is exceeded, the +-- behaviour of the sequence is undefined. This is unlikely to occur in most +-- applications, but some care may be required when using '><', '<*>', '*>', or +-- '>>', particularly repeatedly and particularly in combination with +-- 'replicate' or 'fromFunction'. +-- +----------------------------------------------------------------------------- + + +module Data.Sequence ( +#if defined(DEFINE_PATTERN_SYNONYMS) + Seq (Empty, (:<|), (:|>)), +#else + Seq, +#endif + -- * Construction + empty, -- :: Seq a + singleton, -- :: a -> Seq a + (<|), -- :: a -> Seq a -> Seq a + (|>), -- :: Seq a -> a -> Seq a + (><), -- :: Seq a -> Seq a -> Seq a + fromList, -- :: [a] -> Seq a + fromFunction, -- :: Int -> (Int -> a) -> Seq a + fromArray, -- :: Ix i => Array i a -> Seq a + -- ** Repetition + replicate, -- :: Int -> a -> Seq a + replicateA, -- :: Applicative f => Int -> f a -> f (Seq a) + replicateM, -- :: Monad m => Int -> m a -> m (Seq a) + cycleTaking, -- :: Int -> Seq a -> Seq a + -- ** Iterative construction + iterateN, -- :: Int -> (a -> a) -> a -> Seq a + unfoldr, -- :: (b -> Maybe (a, b)) -> b -> Seq a + unfoldl, -- :: (b -> Maybe (b, a)) -> b -> Seq a + -- * Deconstruction + -- | Additional functions for deconstructing sequences are available + -- via the 'Foldable' instance of 'Seq'. + + -- ** Queries + null, -- :: Seq a -> Bool + length, -- :: Seq a -> Int + -- ** Views + ViewL(..), + viewl, -- :: Seq a -> ViewL a + ViewR(..), + viewr, -- :: Seq a -> ViewR a + -- * Scans + scanl, -- :: (a -> b -> a) -> a -> Seq b -> Seq a + scanl1, -- :: (a -> a -> a) -> Seq a -> Seq a + scanr, -- :: (a -> b -> b) -> b -> Seq a -> Seq b + scanr1, -- :: (a -> a -> a) -> Seq a -> Seq a + -- * Sublists + tails, -- :: Seq a -> Seq (Seq a) + inits, -- :: Seq a -> Seq (Seq a) + chunksOf, -- :: Int -> Seq a -> Seq (Seq a) + -- ** Sequential searches + takeWhileL, -- :: (a -> Bool) -> Seq a -> Seq a + takeWhileR, -- :: (a -> Bool) -> Seq a -> Seq a + dropWhileL, -- :: (a -> Bool) -> Seq a -> Seq a + dropWhileR, -- :: (a -> Bool) -> Seq a -> Seq a + spanl, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + spanr, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + breakl, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + breakr, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + partition, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + filter, -- :: (a -> Bool) -> Seq a -> Seq a + -- * Sorting + sort, -- :: Ord a => Seq a -> Seq a + sortBy, -- :: (a -> a -> Ordering) -> Seq a -> Seq a + unstableSort, -- :: Ord a => Seq a -> Seq a + unstableSortBy, -- :: (a -> a -> Ordering) -> Seq a -> Seq a + -- * Indexing + lookup, -- :: Int -> Seq a -> Maybe a + (!?), -- :: Seq a -> Int -> Maybe a + index, -- :: Seq a -> Int -> a + adjust, -- :: (a -> a) -> Int -> Seq a -> Seq a + adjust', -- :: (a -> a) -> Int -> Seq a -> Seq a + update, -- :: Int -> a -> Seq a -> Seq a + take, -- :: Int -> Seq a -> Seq a + drop, -- :: Int -> Seq a -> Seq a + insertAt, -- :: Int -> a -> Seq a -> Seq a + deleteAt, -- :: Int -> Seq a -> Seq a + splitAt, -- :: Int -> Seq a -> (Seq a, Seq a) + -- ** Indexing with predicates + -- | These functions perform sequential searches from the left + -- or right ends of the sequence, returning indices of matching + -- elements. + elemIndexL, -- :: Eq a => a -> Seq a -> Maybe Int + elemIndicesL, -- :: Eq a => a -> Seq a -> [Int] + elemIndexR, -- :: Eq a => a -> Seq a -> Maybe Int + elemIndicesR, -- :: Eq a => a -> Seq a -> [Int] + findIndexL, -- :: (a -> Bool) -> Seq a -> Maybe Int + findIndicesL, -- :: (a -> Bool) -> Seq a -> [Int] + findIndexR, -- :: (a -> Bool) -> Seq a -> Maybe Int + findIndicesR, -- :: (a -> Bool) -> Seq a -> [Int] + -- * Folds + -- | General folds are available via the 'Foldable' instance of 'Seq'. + foldMapWithIndex, -- :: Monoid m => (Int -> a -> m) -> Seq a -> m + foldlWithIndex, -- :: (b -> Int -> a -> b) -> b -> Seq a -> b + foldrWithIndex, -- :: (Int -> a -> b -> b) -> b -> Seq a -> b + -- * Transformations + mapWithIndex, -- :: (Int -> a -> b) -> Seq a -> Seq b + traverseWithIndex, -- :: Applicative f => (Int -> a -> f b) -> Seq a -> f (Seq b) + reverse, -- :: Seq a -> Seq a + intersperse, -- :: a -> Seq a -> Seq a + -- ** Zips + zip, -- :: Seq a -> Seq b -> Seq (a, b) + zipWith, -- :: (a -> b -> c) -> Seq a -> Seq b -> Seq c + zip3, -- :: Seq a -> Seq b -> Seq c -> Seq (a, b, c) + zipWith3, -- :: (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d + zip4, -- :: Seq a -> Seq b -> Seq c -> Seq d -> Seq (a, b, c, d) + zipWith4, -- :: (a -> b -> c -> d -> e) -> Seq a -> Seq b -> Seq c -> Seq d -> Seq e + ) where + +import Data.Sequence.Base +import Prelude () diff --git a/Data/Sequence/Base.hs b/Data/Sequence/Base.hs new file mode 100644 --- /dev/null +++ b/Data/Sequence/Base.hs @@ -0,0 +1,4280 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} +#if __GLASGOW_HASKELL__ >= 800 +#define DEFINE_PATTERN_SYNONYMS 1 +#endif +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +#endif +#if __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +#if __GLASGOW_HASKELL__ >= 702 +{-# LANGUAGE DeriveGeneric #-} +#endif +#if __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE TypeFamilies #-} +#endif +#ifdef DEFINE_PATTERN_SYNONYMS +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ViewPatterns #-} +#endif + +#include "containers.h" + +{-# OPTIONS_HADDOCK hide #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Sequence.Base +-- Copyright : (c) Ross Paterson 2005 +-- (c) Louis Wasserman 2009 +-- (c) Bertram Felgenhauer, David Feuer, Ross Paterson, and +-- Milan Straka 2014 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : experimental +-- Portability : portable +-- +-- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. +-- +-- = Description +-- +-- General purpose finite sequences. +-- Apart from being finite and having strict operations, sequences +-- also differ from lists in supporting a wider variety of operations +-- efficiently. +-- +-- An amortized running time is given for each operation, with /n/ referring +-- to the length of the sequence and /i/ being the integral index used by +-- some operations. These bounds hold even in a persistent (shared) setting. +-- +-- The implementation uses 2-3 finger trees annotated with sizes, +-- as described in section 4.2 of +-- +-- * Ralf Hinze and Ross Paterson, +-- \"Finger trees: a simple general-purpose data structure\", +-- /Journal of Functional Programming/ 16:2 (2006) pp 197-217. +-- +-- +-- /Note/: Many of these operations have the same names as similar +-- operations on lists in the "Prelude". The ambiguity may be resolved +-- using either qualification or the @hiding@ clause. +-- +-- /Warning/: The size of a 'Seq' must not exceed @maxBound::Int@. Violation +-- of this condition is not detected and if the size limit is exceeded, the +-- behaviour of the sequence is undefined. This is unlikely to occur in most +-- applications, but some care may be required when using '><', '<*>', '*>', or +-- '>>', particularly repeatedly and particularly in combination with +-- 'replicate' or 'fromFunction'. +-- +----------------------------------------------------------------------------- + +module Data.Sequence.Base ( + Elem(..), FingerTree(..), Node(..), Digit(..), Sized(..), MaybeForce, +#if defined(DEFINE_PATTERN_SYNONYMS) + Seq (.., Empty, (:<|), (:|>)), +#else + Seq (..), +#endif + + -- * Construction + empty, -- :: Seq a + singleton, -- :: a -> Seq a + (<|), -- :: a -> Seq a -> Seq a + (|>), -- :: Seq a -> a -> Seq a + (><), -- :: Seq a -> Seq a -> Seq a + fromList, -- :: [a] -> Seq a + fromFunction, -- :: Int -> (Int -> a) -> Seq a + fromArray, -- :: Ix i => Array i a -> Seq a + -- ** Repetition + replicate, -- :: Int -> a -> Seq a + replicateA, -- :: Applicative f => Int -> f a -> f (Seq a) + replicateM, -- :: Monad m => Int -> m a -> m (Seq a) + cycleTaking, -- :: Int -> Seq a -> Seq a + -- ** Iterative construction + iterateN, -- :: Int -> (a -> a) -> a -> Seq a + unfoldr, -- :: (b -> Maybe (a, b)) -> b -> Seq a + unfoldl, -- :: (b -> Maybe (b, a)) -> b -> Seq a + -- * Deconstruction + -- | Additional functions for deconstructing sequences are available + -- via the 'Foldable' instance of 'Seq'. + + -- ** Queries + null, -- :: Seq a -> Bool + length, -- :: Seq a -> Int + -- ** Views + ViewL(..), + viewl, -- :: Seq a -> ViewL a + ViewR(..), + viewr, -- :: Seq a -> ViewR a + -- * Scans + scanl, -- :: (a -> b -> a) -> a -> Seq b -> Seq a + scanl1, -- :: (a -> a -> a) -> Seq a -> Seq a + scanr, -- :: (a -> b -> b) -> b -> Seq a -> Seq b + scanr1, -- :: (a -> a -> a) -> Seq a -> Seq a + -- * Sublists + tails, -- :: Seq a -> Seq (Seq a) + inits, -- :: Seq a -> Seq (Seq a) + chunksOf, -- :: Int -> Seq a -> Seq (Seq a) + -- ** Sequential searches + takeWhileL, -- :: (a -> Bool) -> Seq a -> Seq a + takeWhileR, -- :: (a -> Bool) -> Seq a -> Seq a + dropWhileL, -- :: (a -> Bool) -> Seq a -> Seq a + dropWhileR, -- :: (a -> Bool) -> Seq a -> Seq a + spanl, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + spanr, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + breakl, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + breakr, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + partition, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + filter, -- :: (a -> Bool) -> Seq a -> Seq a + -- * Sorting + sort, -- :: Ord a => Seq a -> Seq a + sortBy, -- :: (a -> a -> Ordering) -> Seq a -> Seq a + unstableSort, -- :: Ord a => Seq a -> Seq a + unstableSortBy, -- :: (a -> a -> Ordering) -> Seq a -> Seq a + -- * Indexing + lookup, -- :: Int -> Seq a -> Maybe a + (!?), -- :: Seq a -> Int -> Maybe a + index, -- :: Seq a -> Int -> a + adjust, -- :: (a -> a) -> Int -> Seq a -> Seq a + adjust', -- :: (a -> a) -> Int -> Seq a -> Seq a + update, -- :: Int -> a -> Seq a -> Seq a + take, -- :: Int -> Seq a -> Seq a + drop, -- :: Int -> Seq a -> Seq a + insertAt, -- :: Int -> a -> Seq a -> Seq a + deleteAt, -- :: Int -> Seq a -> Seq a + splitAt, -- :: Int -> Seq a -> (Seq a, Seq a) + -- ** Indexing with predicates + -- | These functions perform sequential searches from the left + -- or right ends of the sequence, returning indices of matching + -- elements. + elemIndexL, -- :: Eq a => a -> Seq a -> Maybe Int + elemIndicesL, -- :: Eq a => a -> Seq a -> [Int] + elemIndexR, -- :: Eq a => a -> Seq a -> Maybe Int + elemIndicesR, -- :: Eq a => a -> Seq a -> [Int] + findIndexL, -- :: (a -> Bool) -> Seq a -> Maybe Int + findIndicesL, -- :: (a -> Bool) -> Seq a -> [Int] + findIndexR, -- :: (a -> Bool) -> Seq a -> Maybe Int + findIndicesR, -- :: (a -> Bool) -> Seq a -> [Int] + -- * Folds + -- | General folds are available via the 'Foldable' instance of 'Seq'. + foldMapWithIndex, -- :: Monoid m => (Int -> a -> m) -> Seq a -> m + foldlWithIndex, -- :: (b -> Int -> a -> b) -> b -> Seq a -> b + foldrWithIndex, -- :: (Int -> a -> b -> b) -> b -> Seq a -> b + -- * Transformations + mapWithIndex, -- :: (Int -> a -> b) -> Seq a -> Seq b + traverseWithIndex, -- :: Applicative f => (Int -> a -> f b) -> Seq a -> f (Seq b) + reverse, -- :: Seq a -> Seq a + intersperse, -- :: a -> Seq a -> Seq a + -- ** Zips + zip, -- :: Seq a -> Seq b -> Seq (a, b) + zipWith, -- :: (a -> b -> c) -> Seq a -> Seq b -> Seq c + zip3, -- :: Seq a -> Seq b -> Seq c -> Seq (a, b, c) + zipWith3, -- :: (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d + zip4, -- :: Seq a -> Seq b -> Seq c -> Seq d -> Seq (a, b, c, d) + zipWith4, -- :: (a -> b -> c -> d -> e) -> Seq a -> Seq b -> Seq c -> Seq d -> Seq e +#if TESTING + deep, + node2, + node3, +#endif + ) where + +import Prelude hiding ( + Functor(..), +#if MIN_VERSION_base(4,8,0) + Applicative, (<$>), foldMap, Monoid, +#endif + null, length, lookup, take, drop, splitAt, foldl, foldl1, foldr, foldr1, + scanl, scanl1, scanr, scanr1, replicate, zip, zipWith, zip3, zipWith3, + takeWhile, dropWhile, iterate, reverse, filter, mapM, sum, all) +import qualified Data.List +import Control.Applicative (Applicative(..), (<$>), (<**>), Alternative, + WrappedMonad(..), liftA, liftA2, liftA3) +import qualified Control.Applicative as Applicative (Alternative(..)) +import Control.DeepSeq (NFData(rnf)) +import Control.Monad (MonadPlus(..), ap) +import Data.Monoid (Monoid(..)) +import Data.Functor (Functor(..)) +#if MIN_VERSION_base(4,6,0) +import Data.Foldable (Foldable(foldl, foldl1, foldr, foldr1, foldMap, foldl', foldr'), toList) +#else +import Data.Foldable (Foldable(foldl, foldl1, foldr, foldr1, foldMap), foldl', toList) +#endif + +#if MIN_VERSION_base(4,9,0) +import qualified Data.Semigroup as Semigroup +#endif +import Data.Traversable +import Data.Typeable + +-- GHC specific stuff +#ifdef __GLASGOW_HASKELL__ +import GHC.Exts (build) +import Text.Read (Lexeme(Ident), lexP, parens, prec, + readPrec, readListPrec, readListPrecDefault) +import Data.Data +import Data.String (IsString(..)) +#endif +#if __GLASGOW_HASKELL__ >= 706 +import GHC.Generics (Generic, Generic1) +#elif __GLASGOW_HASKELL__ >= 702 +import GHC.Generics (Generic) +#endif + +-- Array stuff, with GHC.Arr on GHC +import Data.Array (Ix, Array) +import qualified Data.Array +#ifdef __GLASGOW_HASKELL__ +import qualified GHC.Arr +#endif + +-- Coercion on GHC 7.8+ +#if __GLASGOW_HASKELL__ >= 708 +import Data.Coerce +import qualified GHC.Exts +#else +#endif + +-- Identity functor on base 4.8 (GHC 7.10+) +#if MIN_VERSION_base(4,8,0) +import Data.Functor.Identity (Identity(..)) +#endif + +#if !MIN_VERSION_base(4,8,0) +import Data.Word (Word) +#endif + +import Data.Utils.StrictPair (StrictPair (..), toPair) + +default () + +-- We define our own copy here, for Monoid only, even though this +-- is now a Semigroup operator in base. The essential reason is that +-- we have absolutely no use for semigroups in this module. Everything +-- that needs to sum things up requires a Monoid constraint to deal +-- with empty sequences. I'm not sure if there's a risk of walking +-- through dictionaries to reach <> from Monoid, but I see no reason +-- to risk it. +infixr 6 <> +(<>) :: Monoid m => m -> m -> m +(<>) = mappend +{-# INLINE (<>) #-} + +infixr 5 `consTree` +infixl 5 `snocTree` +infixr 5 `appendTree0` + +infixr 5 >< +infixr 5 <|, :< +infixl 5 |>, :> + +#ifdef DEFINE_PATTERN_SYNONYMS +infixr 5 :<| +infixl 5 :|> + +-- TODO: Once GHC implements some way to prevent non-exhaustive +-- pattern match warnings for pattern synonyms, we should be +-- sure to take advantage of that. + +-- | A pattern synonym matching an empty sequence. +pattern Empty :: Seq a +pattern Empty = Seq EmptyT + +-- | A pattern synonym viewing the front of a non-empty +-- sequence. +pattern (:<|) :: a -> Seq a -> Seq a +pattern x :<| xs <- (viewl -> x :< xs) + where + x :<| xs = x <| xs + +-- | A pattern synonym viewing the rear of a non-empty +-- sequence. +pattern (:|>) :: Seq a -> a -> Seq a +pattern xs :|> x <- (viewr -> xs :> x) + where + xs :|> x = xs |> x +#endif + +class Sized a where + size :: a -> Int + +-- In much the same way that Sized lets us handle the +-- sizes of elements and nodes uniformly, MaybeForce lets +-- us handle their strictness (or lack thereof) uniformly. +-- We can `mseq` something and not have to worry about +-- whether it's an element or a node. +class MaybeForce a where + maybeRwhnf :: a -> () + +mseq :: MaybeForce a => a -> b -> b +mseq a b = case maybeRwhnf a of () -> b +{-# INLINE mseq #-} + +infixr 0 $!? +($!?) :: MaybeForce a => (a -> b) -> a -> b +f $!? a = case maybeRwhnf a of () -> f a +{-# INLINE ($!?) #-} + +instance MaybeForce (Elem a) where + maybeRwhnf _ = () + {-# INLINE maybeRwhnf #-} + +instance MaybeForce (Node a) where + maybeRwhnf !_ = () + {-# INLINE maybeRwhnf #-} + +-- A wrapper making mseq = seq +newtype ForceBox a = ForceBox a +instance MaybeForce (ForceBox a) where + maybeRwhnf !_ = () +instance Sized (ForceBox a) where + size _ = 1 + +-- | General-purpose finite sequences. +newtype Seq a = Seq (FingerTree (Elem a)) + +instance Functor Seq where + fmap = fmapSeq +#ifdef __GLASGOW_HASKELL__ + x <$ s = replicate (length s) x +#endif + +fmapSeq :: (a -> b) -> Seq a -> Seq b +fmapSeq f (Seq xs) = Seq (fmap (fmap f) xs) +#ifdef __GLASGOW_HASKELL__ +{-# NOINLINE [1] fmapSeq #-} +{-# RULES +"fmapSeq/fmapSeq" forall f g xs . fmapSeq f (fmapSeq g xs) = fmapSeq (f . g) xs + #-} +#endif +#if __GLASGOW_HASKELL__ >= 709 +-- Safe coercions were introduced in 7.8, but did not work well with RULES yet. +{-# RULES +"fmapSeq/coerce" fmapSeq coerce = coerce + #-} +#endif + +instance Foldable Seq where + foldMap f (Seq xs) = foldMap (foldMap f) xs +#if __GLASGOW_HASKELL__ >= 708 + foldr f z (Seq xs) = foldr (coerce f) z xs + foldr' f z (Seq xs) = foldr' (coerce f) z xs +#else + foldr f z (Seq xs) = foldr (flip (foldr f)) z xs +#if MIN_VERSION_base(4,6,0) + foldr' f z (Seq xs) = foldr' (flip (foldr' f)) z xs +#endif +#endif + foldl f z (Seq xs) = foldl (foldl f) z xs +#if MIN_VERSION_base(4,6,0) + foldl' f z (Seq xs) = foldl' (foldl' f) z xs +#endif + + foldr1 f (Seq xs) = getElem (foldr1 f' xs) + where f' (Elem x) (Elem y) = Elem (f x y) + + foldl1 f (Seq xs) = getElem (foldl1 f' xs) + where f' (Elem x) (Elem y) = Elem (f x y) + +#if MIN_VERSION_base(4,8,0) + length = length + {-# INLINE length #-} + null = null + {-# INLINE null #-} +#endif + +#if __GLASGOW_HASKELL__ >= 708 +-- The natural definition of traverse, used for implementations that don't +-- support coercions, `fmap`s into each `Elem`, then `fmap`s again over the +-- result to turn it from a `FingerTree` to a `Seq`. None of this mapping is +-- necessary! We could avoid it without coercions, I believe, by writing a +-- bunch of traversal functions to deal with the `Elem` stuff specially (for +-- FingerTrees, Digits, and Nodes), but using coercions we only need to +-- duplicate code at the FingerTree level. We coerce the `Seq a` to a +-- `FingerTree a`, stripping off all the Elem junk, then use a weird FingerTree +-- traversing function that coerces back to Seq within the functor. +instance Traversable Seq where + traverse f xs = traverseFTE f (coerce xs) + +traverseFTE :: Applicative f => (a -> f b) -> FingerTree a -> f (Seq b) +traverseFTE _f EmptyT = pure empty +traverseFTE f (Single x) = Seq . Single . Elem <$> f x +traverseFTE f (Deep s pr m sf) = + (\pr' m' sf' -> coerce $ Deep s pr' m' sf') <$> + traverse f pr <*> traverse (traverse f) m <*> traverse f sf +#else +instance Traversable Seq where + traverse f (Seq xs) = Seq <$> traverse (traverse f) xs +#endif + +instance NFData a => NFData (Seq a) where + rnf (Seq xs) = rnf xs + +instance Monad Seq where + return = pure + xs >>= f = foldl' add empty xs + where add ys x = ys >< f x + (>>) = (*>) + +instance Applicative Seq where + pure = singleton + xs *> ys = cycleNTimes (length xs) ys + + fs <*> xs@(Seq xsFT) = case viewl fs of + EmptyL -> empty + firstf :< fs' -> case viewr fs' of + EmptyR -> fmap firstf xs + Seq fs''FT :> lastf -> case rigidify xsFT of + RigidEmpty -> empty + RigidOne (Elem x) -> fmap ($x) fs + RigidTwo (Elem x1) (Elem x2) -> + Seq $ ap2FT firstf fs''FT lastf (x1, x2) + RigidThree (Elem x1) (Elem x2) (Elem x3) -> + Seq $ ap3FT firstf fs''FT lastf (x1, x2, x3) + RigidFull r@(Rigid s pr _m sf) -> Seq $ + Deep (s * length fs) + (fmap (fmap firstf) (nodeToDigit pr)) + (aptyMiddle (fmap firstf) (fmap lastf) fmap fs''FT r) + (fmap (fmap lastf) (nodeToDigit sf)) + + +ap2FT :: (a -> b) -> FingerTree (Elem (a->b)) -> (a -> b) -> (a,a) -> FingerTree (Elem b) +ap2FT firstf fs lastf (x,y) = + Deep (size fs * 2 + 4) + (Two (Elem $ firstf x) (Elem $ firstf y)) + (mapMulFT 2 (\(Elem f) -> Node2 2 (Elem (f x)) (Elem (f y))) fs) + (Two (Elem $ lastf x) (Elem $ lastf y)) + +ap3FT :: (a -> b) -> FingerTree (Elem (a->b)) -> (a -> b) -> (a,a,a) -> FingerTree (Elem b) +ap3FT firstf fs lastf (x,y,z) = Deep (size fs * 3 + 6) + (Three (Elem $ firstf x) (Elem $ firstf y) (Elem $ firstf z)) + (mapMulFT 3 (\(Elem f) -> Node3 3 (Elem (f x)) (Elem (f y)) (Elem (f z))) fs) + (Three (Elem $ lastf x) (Elem $ lastf y) (Elem $ lastf z)) + + +data Rigidified a = RigidEmpty + | RigidOne a + | RigidTwo a a + | RigidThree a a a + | RigidFull (Rigid a) +#ifdef TESTING + deriving Show +#endif + +-- | A finger tree whose top level has only Two and/or Three digits, and whose +-- other levels have only One and Two digits. A Rigid tree is precisely what one +-- gets by unzipping/inverting a 2-3 tree, so it is precisely what we need to +-- turn a finger tree into in order to transform it into a 2-3 tree. +data Rigid a = Rigid {-# UNPACK #-} !Int !(Digit23 a) (Thin (Node a)) !(Digit23 a) +#ifdef TESTING + deriving Show +#endif + +-- | A finger tree whose digits are all ones and twos +data Thin a = EmptyTh + | SingleTh a + | DeepTh {-# UNPACK #-} !Int !(Digit12 a) (Thin (Node a)) !(Digit12 a) +#ifdef TESTING + deriving Show +#endif + +data Digit12 a = One12 a | Two12 a a +#ifdef TESTING + deriving Show +#endif + +-- | Sometimes, we want to emphasize that we are viewing a node as a top-level +-- digit of a 'Rigid' tree. +type Digit23 a = Node a + +-- | 'aptyMiddle' does most of the hard work of computing @fs<*>xs@. It +-- produces the center part of a finger tree, with a prefix corresponding to +-- the prefix of @xs@ and a suffix corresponding to the suffix of @xs@ omitted; +-- the missing suffix and prefix are added by the caller. For the recursive +-- call, it squashes the prefix and the suffix into the center tree. Once it +-- gets to the bottom, it turns the tree into a 2-3 tree, applies 'mapMulFT' to +-- produce the main body, and glues all the pieces together. +-- +-- 'map23' itself is a bit horrifying because of the nested types involved. Its +-- job is to map over the *elements* of a 2-3 tree, rather than the subtrees. +-- If we used a higher-order nested type with MPTC, we could probably use a +-- class, but as it is we have to build up 'map23' explicitly through the +-- recursion. +aptyMiddle + :: (c -> d) + -> (c -> d) + -> ((a -> b) -> c -> d) + -> FingerTree (Elem (a -> b)) + -> Rigid c + -> FingerTree (Node d) + +-- Not at the bottom yet + +aptyMiddle firstf + lastf + map23 + fs + (Rigid s pr (DeepTh sm prm mm sfm) sf) + = Deep (sm + s * (size fs + 1)) -- note: sm = s - size pr - size sf + (fmap (fmap firstf) (digit12ToDigit prm)) + (aptyMiddle (fmap firstf) + (fmap lastf) + (fmap . map23) + fs + (Rigid s (squashL pr prm) mm (squashR sfm sf))) + (fmap (fmap lastf) (digit12ToDigit sfm)) + +-- At the bottom + +aptyMiddle firstf + lastf + map23 + fs + (Rigid s pr EmptyTh sf) + = deep + (One (fmap firstf sf)) + (mapMulFT s (\(Elem f) -> fmap (fmap (map23 f)) converted) fs) + (One (fmap lastf pr)) + where converted = node2 pr sf + +aptyMiddle firstf + lastf + map23 + fs + (Rigid s pr (SingleTh q) sf) + = deep + (Two (fmap firstf q) (fmap firstf sf)) + (mapMulFT s (\(Elem f) -> fmap (fmap (map23 f)) converted) fs) + (Two (fmap lastf pr) (fmap lastf q)) + where converted = node3 pr q sf + +digit12ToDigit :: Digit12 a -> Digit a +digit12ToDigit (One12 a) = One a +digit12ToDigit (Two12 a b) = Two a b + +-- Squash the first argument down onto the left side of the second. +squashL :: Digit23 a -> Digit12 (Node a) -> Digit23 (Node a) +squashL m (One12 n) = node2 m n +squashL m (Two12 n1 n2) = node3 m n1 n2 + +-- Squash the second argument down onto the right side of the first +squashR :: Digit12 (Node a) -> Digit23 a -> Digit23 (Node a) +squashR (One12 n) m = node2 n m +squashR (Two12 n1 n2) m = node3 n1 n2 m + + +-- | /O(m*n)/ (incremental) Takes an /O(m)/ function and a finger tree of size +-- /n/ and maps the function over the tree leaves. Unlike the usual 'fmap', the +-- function is applied to the "leaves" of the 'FingerTree' (i.e., given a +-- @FingerTree (Elem a)@, it applies the function to elements of type @Elem +-- a@), replacing the leaves with subtrees of at least the same height, e.g., +-- @Node(Node(Elem y))@. The multiplier argument serves to make the annotations +-- match up properly. +mapMulFT :: Int -> (a -> b) -> FingerTree a -> FingerTree b +mapMulFT _ _ EmptyT = EmptyT +mapMulFT _mul f (Single a) = Single (f a) +mapMulFT mul f (Deep s pr m sf) = Deep (mul * s) (fmap f pr) (mapMulFT mul (mapMulNode mul f) m) (fmap f sf) + +mapMulNode :: Int -> (a -> b) -> Node a -> Node b +mapMulNode mul f (Node2 s a b) = Node2 (mul * s) (f a) (f b) +mapMulNode mul f (Node3 s a b c) = Node3 (mul * s) (f a) (f b) (f c) + +-- | /O(log n)/ (incremental) Takes the extra flexibility out of a 'FingerTree' +-- to make it a genuine 2-3 finger tree. The result of 'rigidify' will have +-- only two and three digits at the top level and only one and two +-- digits elsewhere. If the tree has fewer than four elements, 'rigidify' +-- will simply extract them, and will not build a tree. +rigidify :: FingerTree (Elem a) -> Rigidified (Elem a) +-- The patterns below just fix up the top level of the tree; 'rigidify' +-- delegates the hard work to 'thin'. + +rigidify EmptyT = RigidEmpty + +rigidify (Single q) = RigidOne q + +-- The left digit is Two or Three +rigidify (Deep s (Two a b) m sf) = rigidifyRight s (node2 a b) m sf +rigidify (Deep s (Three a b c) m sf) = rigidifyRight s (node3 a b c) m sf + +-- The left digit is Four +rigidify (Deep s (Four a b c d) m sf) = rigidifyRight s (node2 a b) (node2 c d `consTree` m) sf + +-- The left digit is One +rigidify (Deep s (One a) m sf) = case viewLTree m of + ConsLTree (Node2 _ b c) m' -> rigidifyRight s (node3 a b c) m' sf + ConsLTree (Node3 _ b c d) m' -> rigidifyRight s (node2 a b) (node2 c d `consTree` m') sf + EmptyLTree -> case sf of + One b -> RigidTwo a b + Two b c -> RigidThree a b c + Three b c d -> RigidFull $ Rigid s (node2 a b) EmptyTh (node2 c d) + Four b c d e -> RigidFull $ Rigid s (node3 a b c) EmptyTh (node2 d e) + +-- | /O(log n)/ (incremental) Takes a tree whose left side has been rigidified +-- and finishes the job. +rigidifyRight :: Int -> Digit23 (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> Rigidified (Elem a) + +-- The right digit is Two, Three, or Four +rigidifyRight s pr m (Two a b) = RigidFull $ Rigid s pr (thin m) (node2 a b) +rigidifyRight s pr m (Three a b c) = RigidFull $ Rigid s pr (thin m) (node3 a b c) +rigidifyRight s pr m (Four a b c d) = RigidFull $ Rigid s pr (thin $ m `snocTree` node2 a b) (node2 c d) + +-- The right digit is One +rigidifyRight s pr m (One e) = case viewRTree m of + SnocRTree m' (Node2 _ a b) -> RigidFull $ Rigid s pr (thin m') (node3 a b e) + SnocRTree m' (Node3 _ a b c) -> RigidFull $ Rigid s pr (thin $ m' `snocTree` node2 a b) (node2 c e) + EmptyRTree -> case pr of + Node2 _ a b -> RigidThree a b e + Node3 _ a b c -> RigidFull $ Rigid s (node2 a b) EmptyTh (node2 c e) + +-- | /O(log n)/ (incremental) Rejigger a finger tree so the digits are all ones +-- and twos. +thin :: Sized a => FingerTree a -> Thin a +-- Note that 'thin12' will produce a 'DeepTh' constructor immediately before +-- recursively calling 'thin'. +thin EmptyT = EmptyTh +thin (Single a) = SingleTh a +thin (Deep s pr m sf) = + case pr of + One a -> thin12 s (One12 a) m sf + Two a b -> thin12 s (Two12 a b) m sf + Three a b c -> thin12 s (One12 a) (node2 b c `consTree` m) sf + Four a b c d -> thin12 s (Two12 a b) (node2 c d `consTree` m) sf + +thin12 :: Sized a => Int -> Digit12 a -> FingerTree (Node a) -> Digit a -> Thin a +thin12 s pr m (One a) = DeepTh s pr (thin m) (One12 a) +thin12 s pr m (Two a b) = DeepTh s pr (thin m) (Two12 a b) +thin12 s pr m (Three a b c) = DeepTh s pr (thin $ m `snocTree` node2 a b) (One12 c) +thin12 s pr m (Four a b c d) = DeepTh s pr (thin $ m `snocTree` node2 a b) (Two12 c d) + +-- | Intersperse an element between the elements of a sequence. +-- +-- @ +-- intersperse a empty = empty +-- intersperse a (singleton x) = singleton x +-- intersperse a (fromList [x,y]) = fromList [x,a,y] +-- intersperse a (fromList [x,y,z]) = fromList [x,a,y,a,z] +-- @ +-- +-- @since 0.5.8 +intersperse :: a -> Seq a -> Seq a +intersperse y xs = case viewl xs of + EmptyL -> empty + p :< ps -> p <| (ps <**> (const y <| singleton id)) +-- We used to use +-- +-- intersperse y xs = drop 1 $ xs <**> (const y <| singleton id) +-- +-- but if length xs = ((maxBound :: Int) `quot` 2) + 1 then +-- +-- length (xs <**> (const y <| singleton id)) will wrap around to negative +-- and the drop won't work. The new implementation can produce a result +-- right up to maxBound :: Int + +instance MonadPlus Seq where + mzero = empty + mplus = (><) + +instance Alternative Seq where + empty = empty + (<|>) = (><) + +instance Eq a => Eq (Seq a) where + xs == ys = length xs == length ys && toList xs == toList ys + +instance Ord a => Ord (Seq a) where + compare xs ys = compare (toList xs) (toList ys) + +#if TESTING +instance Show a => Show (Seq a) where + showsPrec p (Seq x) = showsPrec p x +#else +instance Show a => Show (Seq a) where + showsPrec p xs = showParen (p > 10) $ + showString "fromList " . shows (toList xs) +#endif + +instance Read a => Read (Seq a) where +#ifdef __GLASGOW_HASKELL__ + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- readPrec + return (fromList xs) + + readListPrec = readListPrecDefault +#else + readsPrec p = readParen (p > 10) $ \ r -> do + ("fromList",s) <- lex r + (xs,t) <- reads s + return (fromList xs,t) +#endif + +instance Monoid (Seq a) where + mempty = empty + mappend = (><) + +#if MIN_VERSION_base(4,9,0) +instance Semigroup.Semigroup (Seq a) where + (<>) = (><) +#endif + +INSTANCE_TYPEABLE1(Seq) + +#if __GLASGOW_HASKELL__ +instance Data a => Data (Seq a) where + gfoldl f z s = case viewl s of + EmptyL -> z empty + x :< xs -> z (<|) `f` x `f` xs + + gunfold k z c = case constrIndex c of + 1 -> z empty + 2 -> k (k (z (<|))) + _ -> error "gunfold" + + toConstr xs + | null xs = emptyConstr + | otherwise = consConstr + + dataTypeOf _ = seqDataType + + dataCast1 f = gcast1 f + +emptyConstr, consConstr :: Constr +emptyConstr = mkConstr seqDataType "empty" [] Prefix +consConstr = mkConstr seqDataType "<|" [] Infix + +seqDataType :: DataType +seqDataType = mkDataType "Data.Sequence.Seq" [emptyConstr, consConstr] +#endif + +-- Finger trees + +data FingerTree a + = EmptyT + | Single a + | Deep {-# UNPACK #-} !Int !(Digit a) (FingerTree (Node a)) !(Digit a) +#if TESTING + deriving Show +#endif + +instance Sized a => Sized (FingerTree a) where + {-# SPECIALIZE instance Sized (FingerTree (Elem a)) #-} + {-# SPECIALIZE instance Sized (FingerTree (Node a)) #-} + size EmptyT = 0 + size (Single x) = size x + size (Deep v _ _ _) = v + +instance Foldable FingerTree where + foldMap _ EmptyT = mempty + foldMap f (Single x) = f x + foldMap f (Deep _ pr m sf) = + foldMap f pr <> foldMap (foldMap f) m <> foldMap f sf + + foldr _ z EmptyT = z + foldr f z (Single x) = x `f` z + foldr f z (Deep _ pr m sf) = + foldr f (foldr (flip (foldr f)) (foldr f z sf) m) pr + + foldl _ z EmptyT = z + foldl f z (Single x) = z `f` x + foldl f z (Deep _ pr m sf) = + foldl f (foldl (foldl f) (foldl f z pr) m) sf + +#if MIN_VERSION_base(4,6,0) + foldr' _ z EmptyT = z + foldr' f z (Single x) = f x z + foldr' f z (Deep _ pr m sf) = foldr' f mres pr + where !sfRes = foldr' f z sf + !mres = foldr' (flip (foldr' f)) sfRes m + + foldl' _ z EmptyT = z + foldl' f z (Single x) = z `f` x + foldl' f z (Deep _ pr m sf) = foldl' f mres sf + where !prRes = foldl' f z pr + !mres = foldl' (foldl' f) prRes m +#endif + + foldr1 _ EmptyT = error "foldr1: empty sequence" + foldr1 _ (Single x) = x + foldr1 f (Deep _ pr m sf) = + foldr f (foldr (flip (foldr f)) (foldr1 f sf) m) pr + + foldl1 _ EmptyT = error "foldl1: empty sequence" + foldl1 _ (Single x) = x + foldl1 f (Deep _ pr m sf) = + foldl f (foldl (foldl f) (foldl1 f pr) m) sf + +instance Functor FingerTree where + fmap _ EmptyT = EmptyT + fmap f (Single x) = Single (f x) + fmap f (Deep v pr m sf) = + Deep v (fmap f pr) (fmap (fmap f) m) (fmap f sf) + +instance Traversable FingerTree where + traverse _ EmptyT = pure EmptyT + traverse f (Single x) = Single <$> f x + traverse f (Deep v pr m sf) = + deep' v <$> traverse f pr <*> traverse (traverse f) m <*> + traverse f sf + +instance NFData a => NFData (FingerTree a) where + rnf EmptyT = () + rnf (Single x) = rnf x + rnf (Deep _ pr m sf) = rnf pr `seq` rnf sf `seq` rnf m + +{-# INLINE deep #-} +deep :: Sized a => Digit a -> FingerTree (Node a) -> Digit a -> FingerTree a +deep pr m sf = Deep (size pr + size m + size sf) pr m sf + +{-# INLINE pullL #-} +pullL :: Int -> FingerTree (Node a) -> Digit a -> FingerTree a +pullL s m sf = case viewLTree m of + EmptyLTree -> digitToTree' s sf + ConsLTree pr m' -> Deep s (nodeToDigit pr) m' sf + +{-# INLINE pullR #-} +pullR :: Int -> Digit a -> FingerTree (Node a) -> FingerTree a +pullR s pr m = case viewRTree m of + EmptyRTree -> digitToTree' s pr + SnocRTree m' sf -> Deep s pr m' (nodeToDigit sf) + +-- Digits + +data Digit a + = One a + | Two a a + | Three a a a + | Four a a a a +#if TESTING + deriving Show +#endif + +instance Foldable Digit where + foldMap f (One a) = f a + foldMap f (Two a b) = f a <> f b + foldMap f (Three a b c) = f a <> f b <> f c + foldMap f (Four a b c d) = f a <> f b <> f c <> f d + + foldr f z (One a) = a `f` z + foldr f z (Two a b) = a `f` (b `f` z) + foldr f z (Three a b c) = a `f` (b `f` (c `f` z)) + foldr f z (Four a b c d) = a `f` (b `f` (c `f` (d `f` z))) + + foldl f z (One a) = z `f` a + foldl f z (Two a b) = (z `f` a) `f` b + foldl f z (Three a b c) = ((z `f` a) `f` b) `f` c + foldl f z (Four a b c d) = (((z `f` a) `f` b) `f` c) `f` d + +#if MIN_VERSION_base(4,6,0) + foldr' f z (One a) = a `f` z + foldr' f z (Two a b) = f a $! f b z + foldr' f z (Three a b c) = f a $! f b $! f c z + foldr' f z (Four a b c d) = f a $! f b $! f c $! f d z + + foldl' f z (One a) = f z a + foldl' f z (Two a b) = (f $! f z a) b + foldl' f z (Three a b c) = (f $! (f $! f z a) b) c + foldl' f z (Four a b c d) = (f $! (f $! (f $! f z a) b) c) d +#endif + + foldr1 _ (One a) = a + foldr1 f (Two a b) = a `f` b + foldr1 f (Three a b c) = a `f` (b `f` c) + foldr1 f (Four a b c d) = a `f` (b `f` (c `f` d)) + + foldl1 _ (One a) = a + foldl1 f (Two a b) = a `f` b + foldl1 f (Three a b c) = (a `f` b) `f` c + foldl1 f (Four a b c d) = ((a `f` b) `f` c) `f` d + +instance Functor Digit where + {-# INLINE fmap #-} + fmap f (One a) = One (f a) + fmap f (Two a b) = Two (f a) (f b) + fmap f (Three a b c) = Three (f a) (f b) (f c) + fmap f (Four a b c d) = Four (f a) (f b) (f c) (f d) + +instance Traversable Digit where + {-# INLINE traverse #-} + traverse f (One a) = One <$> f a + traverse f (Two a b) = Two <$> f a <*> f b + traverse f (Three a b c) = Three <$> f a <*> f b <*> f c + traverse f (Four a b c d) = Four <$> f a <*> f b <*> f c <*> f d + +instance NFData a => NFData (Digit a) where + rnf (One a) = rnf a + rnf (Two a b) = rnf a `seq` rnf b + rnf (Three a b c) = rnf a `seq` rnf b `seq` rnf c + rnf (Four a b c d) = rnf a `seq` rnf b `seq` rnf c `seq` rnf d + +instance Sized a => Sized (Digit a) where + {-# INLINE size #-} + size = foldl1 (+) . fmap size + +{-# SPECIALIZE digitToTree :: Digit (Elem a) -> FingerTree (Elem a) #-} +{-# SPECIALIZE digitToTree :: Digit (Node a) -> FingerTree (Node a) #-} +digitToTree :: Sized a => Digit a -> FingerTree a +digitToTree (One a) = Single a +digitToTree (Two a b) = deep (One a) EmptyT (One b) +digitToTree (Three a b c) = deep (Two a b) EmptyT (One c) +digitToTree (Four a b c d) = deep (Two a b) EmptyT (Two c d) + +-- | Given the size of a digit and the digit itself, efficiently converts +-- it to a FingerTree. +digitToTree' :: Int -> Digit a -> FingerTree a +digitToTree' n (Four a b c d) = Deep n (Two a b) EmptyT (Two c d) +digitToTree' n (Three a b c) = Deep n (Two a b) EmptyT (One c) +digitToTree' n (Two a b) = Deep n (One a) EmptyT (One b) +digitToTree' !_n (One a) = Single a + +-- Nodes + +data Node a + = Node2 {-# UNPACK #-} !Int a a + | Node3 {-# UNPACK #-} !Int a a a +#if TESTING + deriving Show +#endif + +-- Sometimes, we need to apply a Node2, Node3, or Deep constructor +-- to a size and pass the result to a function. If we calculate, +-- say, `Node2 n <$> x <*> y`, then according to -ddump-simpl, +-- GHC boxes up `n`, passes it to the strict constructor for `Node2`, +-- and passes the result to `fmap`. Using `node2'` instead prevents +-- this, forming a closure with the unboxed size. +{-# INLINE node2' #-} +node2' :: Int -> a -> a -> Node a +node2' !s = \a b -> Node2 s a b + +{-# INLINE node3' #-} +node3' :: Int -> a -> a -> a -> Node a +node3' !s = \a b c -> Node3 s a b c + +{-# INLINE deep' #-} +deep' :: Int -> Digit a -> FingerTree (Node a) -> Digit a -> FingerTree a +deep' !s = \pr m sf -> Deep s pr m sf + +instance Foldable Node where + foldMap f (Node2 _ a b) = f a <> f b + foldMap f (Node3 _ a b c) = f a <> f b <> f c + + foldr f z (Node2 _ a b) = a `f` (b `f` z) + foldr f z (Node3 _ a b c) = a `f` (b `f` (c `f` z)) + + foldl f z (Node2 _ a b) = (z `f` a) `f` b + foldl f z (Node3 _ a b c) = ((z `f` a) `f` b) `f` c + +#if MIN_VERSION_base(4,6,0) + foldr' f z (Node2 _ a b) = f a $! f b z + foldr' f z (Node3 _ a b c) = f a $! f b $! f c z + + foldl' f z (Node2 _ a b) = (f $! f z a) b + foldl' f z (Node3 _ a b c) = (f $! (f $! f z a) b) c +#endif + +instance Functor Node where + {-# INLINE fmap #-} + fmap f (Node2 v a b) = Node2 v (f a) (f b) + fmap f (Node3 v a b c) = Node3 v (f a) (f b) (f c) + +instance Traversable Node where + {-# INLINE traverse #-} + traverse f (Node2 v a b) = node2' v <$> f a <*> f b + traverse f (Node3 v a b c) = node3' v <$> f a <*> f b <*> f c + +instance NFData a => NFData (Node a) where + rnf (Node2 _ a b) = rnf a `seq` rnf b + rnf (Node3 _ a b c) = rnf a `seq` rnf b `seq` rnf c + +instance Sized (Node a) where + size (Node2 v _ _) = v + size (Node3 v _ _ _) = v + +{-# INLINE node2 #-} +node2 :: Sized a => a -> a -> Node a +node2 a b = Node2 (size a + size b) a b + +{-# INLINE node3 #-} +node3 :: Sized a => a -> a -> a -> Node a +node3 a b c = Node3 (size a + size b + size c) a b c + +nodeToDigit :: Node a -> Digit a +nodeToDigit (Node2 _ a b) = Two a b +nodeToDigit (Node3 _ a b c) = Three a b c + +-- Elements + +newtype Elem a = Elem { getElem :: a } +#if TESTING + deriving Show +#endif + +instance Sized (Elem a) where + size _ = 1 + +instance Functor Elem where +#if __GLASGOW_HASKELL__ >= 708 +-- This cuts the time for <*> by around a fifth. + fmap = coerce +#else + fmap f (Elem x) = Elem (f x) +#endif + +instance Foldable Elem where + foldr f z (Elem x) = f x z +#if __GLASGOW_HASKELL__ >= 708 + foldMap = coerce + foldl = coerce + foldl' = coerce +#else + foldMap f (Elem x) = f x + foldl f z (Elem x) = f z x +#if MIN_VERSION_base(4,6,0) + foldl' f z (Elem x) = f z x +#endif +#endif + +instance Traversable Elem where + traverse f (Elem x) = Elem <$> f x + +instance NFData a => NFData (Elem a) where + rnf (Elem x) = rnf x + +------------------------------------------------------- +-- Applicative construction +------------------------------------------------------- +#if !MIN_VERSION_base(4,8,0) +newtype Identity a = Identity {runIdentity :: a} + +instance Functor Identity where + fmap f (Identity x) = Identity (f x) + +instance Applicative Identity where + pure = Identity + Identity f <*> Identity x = Identity (f x) +#endif + +-- | This is essentially a clone of Control.Monad.State.Strict. +newtype State s a = State {runState :: s -> (s, a)} + +instance Functor (State s) where + fmap = liftA + +instance Monad (State s) where + {-# INLINE return #-} + {-# INLINE (>>=) #-} + return = pure + m >>= k = State $ \ s -> case runState m s of + (s', x) -> runState (k x) s' + +instance Applicative (State s) where + {-# INLINE pure #-} + pure x = State $ \ s -> (s, x) + (<*>) = ap + +execState :: State s a -> s -> a +execState m x = snd (runState m x) + +-- | 'applicativeTree' takes an Applicative-wrapped construction of a +-- piece of a FingerTree, assumed to always have the same size (which +-- is put in the second argument), and replicates it as many times as +-- specified. This is a generalization of 'replicateA', which itself +-- is a generalization of many Data.Sequence methods. +{-# SPECIALIZE applicativeTree :: Int -> Int -> State s a -> State s (FingerTree a) #-} +{-# SPECIALIZE applicativeTree :: Int -> Int -> Identity a -> Identity (FingerTree a) #-} +-- Special note: the Identity specialization automatically does node sharing, +-- reducing memory usage of the resulting tree to /O(log n)/. +applicativeTree :: Applicative f => Int -> Int -> f a -> f (FingerTree a) +applicativeTree n !mSize m = case n of + 0 -> pure EmptyT + 1 -> fmap Single m + 2 -> deepA one emptyTree one + 3 -> deepA two emptyTree one + 4 -> deepA two emptyTree two + 5 -> deepA three emptyTree two + 6 -> deepA three emptyTree three + _ -> case n `quotRem` 3 of + (q,0) -> deepA three (applicativeTree (q - 2) mSize' n3) three + (q,1) -> deepA two (applicativeTree (q - 1) mSize' n3) two + (q,_) -> deepA three (applicativeTree (q - 1) mSize' n3) two + where !mSize' = 3 * mSize + n3 = liftA3 (node3' mSize') m m m + where + one = fmap One m + two = liftA2 Two m m + three = liftA3 Three m m m + deepA = liftA3 (deep' (n * mSize)) + emptyTree = pure EmptyT + +------------------------------------------------------------------------ +-- Construction +------------------------------------------------------------------------ + +-- | /O(1)/. The empty sequence. +empty :: Seq a +empty = Seq EmptyT + +-- | /O(1)/. A singleton sequence. +singleton :: a -> Seq a +singleton x = Seq (Single (Elem x)) + +-- | /O(log n)/. @replicate n x@ is a sequence consisting of @n@ copies of @x@. +replicate :: Int -> a -> Seq a +replicate n x + | n >= 0 = runIdentity (replicateA n (Identity x)) + | otherwise = error "replicate takes a nonnegative integer argument" + +-- | 'replicateA' is an 'Applicative' version of 'replicate', and makes +-- /O(log n)/ calls to '<*>' and 'pure'. +-- +-- > replicateA n x = sequenceA (replicate n x) +replicateA :: Applicative f => Int -> f a -> f (Seq a) +replicateA n x + | n >= 0 = Seq <$> applicativeTree n 1 (Elem <$> x) + | otherwise = error "replicateA takes a nonnegative integer argument" + +-- | 'replicateM' is a sequence counterpart of 'Control.Monad.replicateM'. +-- +-- > replicateM n x = sequence (replicate n x) +replicateM :: Monad m => Int -> m a -> m (Seq a) +replicateM n x + | n >= 0 = unwrapMonad (replicateA n (WrapMonad x)) + | otherwise = error "replicateM takes a nonnegative integer argument" + +-- | /O(log(k))/. @'cycleTaking' k xs@ forms a sequence of length @k@ by +-- repeatedly concatenating @xs@ with itself. @xs@ may only be empty if +-- @k@ is 0. +-- +-- prop> cycleTaking k = fromList . take k . cycle . toList + +-- If you wish to concatenate a non-empty sequence @xs@ with itself precisely +-- @k@ times, you can use @cycleTaking (k * length xs)@ or just +-- @replicate k () *> xs@. +-- +-- @since 0.5.8 +cycleTaking :: Int -> Seq a -> Seq a +cycleTaking n !_xs | n <= 0 = empty +cycleTaking _n xs | null xs = error "cycleTaking cannot take a positive number of elements from an empty cycle." +cycleTaking n xs = cycleNTimes reps xs >< take final xs + where + (reps, final) = n `quotRem` length xs + +-- | /O(log(kn))/. @'cycleNTimes' k xs@ concatenates @k@ copies of @xs@. This +-- operation uses time and additional space logarithmic in the size of its +-- result. +cycleNTimes :: Int -> Seq a -> Seq a +cycleNTimes n !xs + | n <= 0 = empty + | n == 1 = xs +cycleNTimes n (Seq xsFT) = case rigidify xsFT of + RigidEmpty -> empty + RigidOne (Elem x) -> replicate n x + RigidTwo x1 x2 -> Seq $ + Deep (n*2) pair + (runIdentity $ applicativeTree (n-2) 2 (Identity (node2 x1 x2))) + pair + where pair = Two x1 x2 + RigidThree x1 x2 x3 -> Seq $ + Deep (n*3) triple + (runIdentity $ applicativeTree (n-2) 3 (Identity (node3 x1 x2 x3))) + triple + where triple = Three x1 x2 x3 + RigidFull r@(Rigid s pr _m sf) -> Seq $ + Deep (n*s) + (nodeToDigit pr) + (cycleNMiddle (n-2) r) + (nodeToDigit sf) + +cycleNMiddle + :: Int + -> Rigid c + -> FingerTree (Node c) + +-- Not at the bottom yet + +cycleNMiddle !n + (Rigid s pr (DeepTh sm prm mm sfm) sf) + = Deep (sm + s * (n + 1)) -- note: sm = s - size pr - size sf + (digit12ToDigit prm) + (cycleNMiddle n + (Rigid s (squashL pr prm) mm (squashR sfm sf))) + (digit12ToDigit sfm) + +-- At the bottom + +cycleNMiddle n + (Rigid s pr EmptyTh sf) + = deep + (One sf) + (runIdentity $ applicativeTree n s (Identity converted)) + (One pr) + where converted = node2 pr sf + +cycleNMiddle n + (Rigid s pr (SingleTh q) sf) + = deep + (Two q sf) + (runIdentity $ applicativeTree n s (Identity converted)) + (Two pr q) + where converted = node3 pr q sf + + +-- | /O(1)/. Add an element to the left end of a sequence. +-- Mnemonic: a triangle with the single element at the pointy end. +(<|) :: a -> Seq a -> Seq a +x <| Seq xs = Seq (Elem x `consTree` xs) + +{-# SPECIALIZE consTree :: Elem a -> FingerTree (Elem a) -> FingerTree (Elem a) #-} +{-# SPECIALIZE consTree :: Node a -> FingerTree (Node a) -> FingerTree (Node a) #-} +consTree :: Sized a => a -> FingerTree a -> FingerTree a +consTree a EmptyT = Single a +consTree a (Single b) = deep (One a) EmptyT (One b) +-- As described in the paper, we force the middle of a tree +-- *before* consing onto it; this preserves the amortized +-- bounds but prevents repeated consing from building up +-- gigantic suspensions. +consTree a (Deep s (Four b c d e) m sf) = m `seq` + Deep (size a + s) (Two a b) (node3 c d e `consTree` m) sf +consTree a (Deep s (Three b c d) m sf) = + Deep (size a + s) (Four a b c d) m sf +consTree a (Deep s (Two b c) m sf) = + Deep (size a + s) (Three a b c) m sf +consTree a (Deep s (One b) m sf) = + Deep (size a + s) (Two a b) m sf + +cons' :: a -> Seq a -> Seq a +cons' x (Seq xs) = Seq (Elem x `consTree'` xs) + +snoc' :: Seq a -> a -> Seq a +snoc' (Seq xs) x = Seq (xs `snocTree'` Elem x) + +{-# SPECIALIZE consTree' :: Elem a -> FingerTree (Elem a) -> FingerTree (Elem a) #-} +{-# SPECIALIZE consTree' :: Node a -> FingerTree (Node a) -> FingerTree (Node a) #-} +consTree' :: Sized a => a -> FingerTree a -> FingerTree a +consTree' a EmptyT = Single a +consTree' a (Single b) = deep (One a) EmptyT (One b) +-- As described in the paper, we force the middle of a tree +-- *before* consing onto it; this preserves the amortized +-- bounds but prevents repeated consing from building up +-- gigantic suspensions. +consTree' a (Deep s (Four b c d e) m sf) = + Deep (size a + s) (Two a b) m' sf + where !m' = abc `consTree'` m + !abc = node3 c d e +consTree' a (Deep s (Three b c d) m sf) = + Deep (size a + s) (Four a b c d) m sf +consTree' a (Deep s (Two b c) m sf) = + Deep (size a + s) (Three a b c) m sf +consTree' a (Deep s (One b) m sf) = + Deep (size a + s) (Two a b) m sf + +-- | /O(1)/. Add an element to the right end of a sequence. +-- Mnemonic: a triangle with the single element at the pointy end. +(|>) :: Seq a -> a -> Seq a +Seq xs |> x = Seq (xs `snocTree` Elem x) + +{-# SPECIALIZE snocTree :: FingerTree (Elem a) -> Elem a -> FingerTree (Elem a) #-} +{-# SPECIALIZE snocTree :: FingerTree (Node a) -> Node a -> FingerTree (Node a) #-} +snocTree :: Sized a => FingerTree a -> a -> FingerTree a +snocTree EmptyT a = Single a +snocTree (Single a) b = deep (One a) EmptyT (One b) +-- See note on `seq` in `consTree`. +snocTree (Deep s pr m (Four a b c d)) e = m `seq` + Deep (s + size e) pr (m `snocTree` node3 a b c) (Two d e) +snocTree (Deep s pr m (Three a b c)) d = + Deep (s + size d) pr m (Four a b c d) +snocTree (Deep s pr m (Two a b)) c = + Deep (s + size c) pr m (Three a b c) +snocTree (Deep s pr m (One a)) b = + Deep (s + size b) pr m (Two a b) + +{-# SPECIALIZE snocTree' :: FingerTree (Elem a) -> Elem a -> FingerTree (Elem a) #-} +{-# SPECIALIZE snocTree' :: FingerTree (Node a) -> Node a -> FingerTree (Node a) #-} +snocTree' :: Sized a => FingerTree a -> a -> FingerTree a +snocTree' EmptyT a = Single a +snocTree' (Single a) b = deep (One a) EmptyT (One b) +-- See note on `seq` in `consTree`. +snocTree' (Deep s pr m (Four a b c d)) e = + Deep (s + size e) pr m' (Two d e) + where !m' = m `snocTree'` abc + !abc = node3 a b c +snocTree' (Deep s pr m (Three a b c)) d = + Deep (s + size d) pr m (Four a b c d) +snocTree' (Deep s pr m (Two a b)) c = + Deep (s + size c) pr m (Three a b c) +snocTree' (Deep s pr m (One a)) b = + Deep (s + size b) pr m (Two a b) + +-- | /O(log(min(n1,n2)))/. Concatenate two sequences. +(><) :: Seq a -> Seq a -> Seq a +Seq xs >< Seq ys = Seq (appendTree0 xs ys) + +-- The appendTree/addDigits gunk below is machine generated + +appendTree0 :: FingerTree (Elem a) -> FingerTree (Elem a) -> FingerTree (Elem a) +appendTree0 EmptyT xs = + xs +appendTree0 xs EmptyT = + xs +appendTree0 (Single x) xs = + x `consTree` xs +appendTree0 xs (Single x) = + xs `snocTree` x +appendTree0 (Deep s1 pr1 m1 sf1) (Deep s2 pr2 m2 sf2) = + Deep (s1 + s2) pr1 m sf2 + where !m = addDigits0 m1 sf1 pr2 m2 + +addDigits0 :: FingerTree (Node (Elem a)) -> Digit (Elem a) -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> FingerTree (Node (Elem a)) +addDigits0 m1 (One a) (One b) m2 = + appendTree1 m1 (node2 a b) m2 +addDigits0 m1 (One a) (Two b c) m2 = + appendTree1 m1 (node3 a b c) m2 +addDigits0 m1 (One a) (Three b c d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits0 m1 (One a) (Four b c d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits0 m1 (Two a b) (One c) m2 = + appendTree1 m1 (node3 a b c) m2 +addDigits0 m1 (Two a b) (Two c d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits0 m1 (Two a b) (Three c d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits0 m1 (Two a b) (Four c d e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits0 m1 (Three a b c) (One d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits0 m1 (Three a b c) (Two d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits0 m1 (Three a b c) (Three d e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits0 m1 (Three a b c) (Four d e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits0 m1 (Four a b c d) (One e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits0 m1 (Four a b c d) (Two e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits0 m1 (Four a b c d) (Three e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits0 m1 (Four a b c d) (Four e f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 + +appendTree1 :: FingerTree (Node a) -> Node a -> FingerTree (Node a) -> FingerTree (Node a) +appendTree1 EmptyT !a xs = + a `consTree` xs +appendTree1 xs !a EmptyT = + xs `snocTree` a +appendTree1 (Single x) !a xs = + x `consTree` a `consTree` xs +appendTree1 xs !a (Single x) = + xs `snocTree` a `snocTree` x +appendTree1 (Deep s1 pr1 m1 sf1) a (Deep s2 pr2 m2 sf2) = + Deep (s1 + size a + s2) pr1 m sf2 + where !m = addDigits1 m1 sf1 a pr2 m2 + +addDigits1 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) +addDigits1 m1 (One a) b (One c) m2 = + appendTree1 m1 (node3 a b c) m2 +addDigits1 m1 (One a) b (Two c d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits1 m1 (One a) b (Three c d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits1 m1 (One a) b (Four c d e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits1 m1 (Two a b) c (One d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits1 m1 (Two a b) c (Two d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits1 m1 (Two a b) c (Three d e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits1 m1 (Two a b) c (Four d e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits1 m1 (Three a b c) d (One e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits1 m1 (Three a b c) d (Two e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits1 m1 (Three a b c) d (Three e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits1 m1 (Three a b c) d (Four e f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits1 m1 (Four a b c d) e (One f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits1 m1 (Four a b c d) e (Two f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits1 m1 (Four a b c d) e (Three f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits1 m1 (Four a b c d) e (Four f g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 + +appendTree2 :: FingerTree (Node a) -> Node a -> Node a -> FingerTree (Node a) -> FingerTree (Node a) +appendTree2 EmptyT !a !b xs = + a `consTree` b `consTree` xs +appendTree2 xs !a !b EmptyT = + xs `snocTree` a `snocTree` b +appendTree2 (Single x) a b xs = + x `consTree` a `consTree` b `consTree` xs +appendTree2 xs a b (Single x) = + xs `snocTree` a `snocTree` b `snocTree` x +appendTree2 (Deep s1 pr1 m1 sf1) a b (Deep s2 pr2 m2 sf2) = + Deep (s1 + size a + size b + s2) pr1 m sf2 + where !m = addDigits2 m1 sf1 a b pr2 m2 + +addDigits2 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) +addDigits2 m1 (One a) b c (One d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits2 m1 (One a) b c (Two d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits2 m1 (One a) b c (Three d e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits2 m1 (One a) b c (Four d e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits2 m1 (Two a b) c d (One e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits2 m1 (Two a b) c d (Two e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits2 m1 (Two a b) c d (Three e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits2 m1 (Two a b) c d (Four e f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits2 m1 (Three a b c) d e (One f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits2 m1 (Three a b c) d e (Two f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits2 m1 (Three a b c) d e (Three f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits2 m1 (Three a b c) d e (Four f g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits2 m1 (Four a b c d) e f (One g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits2 m1 (Four a b c d) e f (Two g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits2 m1 (Four a b c d) e f (Three g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits2 m1 (Four a b c d) e f (Four g h i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 + +appendTree3 :: FingerTree (Node a) -> Node a -> Node a -> Node a -> FingerTree (Node a) -> FingerTree (Node a) +appendTree3 EmptyT !a !b !c xs = + a `consTree` b `consTree` c `consTree` xs +appendTree3 xs !a !b !c EmptyT = + xs `snocTree` a `snocTree` b `snocTree` c +appendTree3 (Single x) a b c xs = + x `consTree` a `consTree` b `consTree` c `consTree` xs +appendTree3 xs a b c (Single x) = + xs `snocTree` a `snocTree` b `snocTree` c `snocTree` x +appendTree3 (Deep s1 pr1 m1 sf1) a b c (Deep s2 pr2 m2 sf2) = + Deep (s1 + size a + size b + size c + s2) pr1 m sf2 + where !m = addDigits3 m1 sf1 a b c pr2 m2 + +addDigits3 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Node a -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) +addDigits3 m1 (One a) !b !c !d (One e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits3 m1 (One a) b c d (Two e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits3 m1 (One a) b c d (Three e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits3 m1 (One a) b c d (Four e f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits3 m1 (Two a b) !c !d !e (One f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits3 m1 (Two a b) c d e (Two f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits3 m1 (Two a b) c d e (Three f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits3 m1 (Two a b) c d e (Four f g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits3 m1 (Three a b c) !d !e !f (One g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits3 m1 (Three a b c) d e f (Two g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits3 m1 (Three a b c) d e f (Three g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits3 m1 (Three a b c) d e f (Four g h i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 +addDigits3 m1 (Four a b c d) !e !f !g (One h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits3 m1 (Four a b c d) e f g (Two h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits3 m1 (Four a b c d) e f g (Three h i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 +addDigits3 m1 (Four a b c d) e f g (Four h i j k) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node2 j k) m2 + +appendTree4 :: FingerTree (Node a) -> Node a -> Node a -> Node a -> Node a -> FingerTree (Node a) -> FingerTree (Node a) +appendTree4 EmptyT !a !b !c !d xs = + a `consTree` b `consTree` c `consTree` d `consTree` xs +appendTree4 xs !a !b !c !d EmptyT = + xs `snocTree` a `snocTree` b `snocTree` c `snocTree` d +appendTree4 (Single x) a b c d xs = + x `consTree` a `consTree` b `consTree` c `consTree` d `consTree` xs +appendTree4 xs a b c d (Single x) = + xs `snocTree` a `snocTree` b `snocTree` c `snocTree` d `snocTree` x +appendTree4 (Deep s1 pr1 m1 sf1) a b c d (Deep s2 pr2 m2 sf2) = + Deep (s1 + size a + size b + size c + size d + s2) pr1 m sf2 + where !m = addDigits4 m1 sf1 a b c d pr2 m2 + +addDigits4 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Node a -> Node a -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) +addDigits4 m1 (One a) !b !c !d !e (One f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits4 m1 (One a) b c d e (Two f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits4 m1 (One a) b c d e (Three f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits4 m1 (One a) b c d e (Four f g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits4 m1 (Two a b) !c !d !e !f (One g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits4 m1 (Two a b) c d e f (Two g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits4 m1 (Two a b) c d e f (Three g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits4 m1 (Two a b) c d e f (Four g h i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 +addDigits4 m1 (Three a b c) !d !e !f !g (One h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits4 m1 (Three a b c) d e f g (Two h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits4 m1 (Three a b c) d e f g (Three h i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 +addDigits4 m1 (Three a b c) d e f g (Four h i j k) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node2 j k) m2 +addDigits4 m1 (Four a b c d) !e !f !g !h (One i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits4 m1 (Four a b c d) !e !f !g !h (Two i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 +addDigits4 m1 (Four a b c d) !e !f !g !h (Three i j k) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node2 j k) m2 +addDigits4 m1 (Four a b c d) !e !f !g !h (Four i j k l) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node3 j k l) m2 + +-- | Builds a sequence from a seed value. Takes time linear in the +-- number of generated elements. /WARNING:/ If the number of generated +-- elements is infinite, this method will not terminate. +unfoldr :: (b -> Maybe (a, b)) -> b -> Seq a +unfoldr f = unfoldr' empty + -- uses tail recursion rather than, for instance, the List implementation. + where unfoldr' !as b = maybe as (\ (a, b') -> unfoldr' (as `snoc'` a) b') (f b) + +-- | @'unfoldl' f x@ is equivalent to @'reverse' ('unfoldr' ('fmap' swap . f) x)@. +unfoldl :: (b -> Maybe (b, a)) -> b -> Seq a +unfoldl f = unfoldl' empty + where unfoldl' !as b = maybe as (\ (b', a) -> unfoldl' (a `cons'` as) b') (f b) + +-- | /O(n)/. Constructs a sequence by repeated application of a function +-- to a seed value. +-- +-- > iterateN n f x = fromList (Prelude.take n (Prelude.iterate f x)) +iterateN :: Int -> (a -> a) -> a -> Seq a +iterateN n f x + | n >= 0 = replicateA n (State (\ y -> (f y, y))) `execState` x + | otherwise = error "iterateN takes a nonnegative integer argument" + +------------------------------------------------------------------------ +-- Deconstruction +------------------------------------------------------------------------ + +-- | /O(1)/. Is this the empty sequence? +null :: Seq a -> Bool +null (Seq EmptyT) = True +null _ = False + +-- | /O(1)/. The number of elements in the sequence. +length :: Seq a -> Int +length (Seq xs) = size xs + +-- Views + +data ViewLTree a = ConsLTree a (FingerTree a) | EmptyLTree +data ViewRTree a = SnocRTree (FingerTree a) a | EmptyRTree + +-- | View of the left end of a sequence. +data ViewL a + = EmptyL -- ^ empty sequence + | a :< Seq a -- ^ leftmost element and the rest of the sequence + deriving (Eq, Ord, Show, Read) + +#if __GLASGOW_HASKELL__ +deriving instance Data a => Data (ViewL a) +#endif +#if __GLASGOW_HASKELL__ >= 706 +deriving instance Generic1 ViewL +#endif +#if __GLASGOW_HASKELL__ >= 702 +deriving instance Generic (ViewL a) +#endif + +INSTANCE_TYPEABLE1(ViewL) + +instance Functor ViewL where + {-# INLINE fmap #-} + fmap _ EmptyL = EmptyL + fmap f (x :< xs) = f x :< fmap f xs + +instance Foldable ViewL where + foldr _ z EmptyL = z + foldr f z (x :< xs) = f x (foldr f z xs) + + foldl _ z EmptyL = z + foldl f z (x :< xs) = foldl f (f z x) xs + + foldl1 _ EmptyL = error "foldl1: empty view" + foldl1 f (x :< xs) = foldl f x xs + +#if MIN_VERSION_base(4,8,0) + null EmptyL = True + null (_ :< _) = False + + length EmptyL = 0 + length (_ :< xs) = 1 + length xs +#endif + +instance Traversable ViewL where + traverse _ EmptyL = pure EmptyL + traverse f (x :< xs) = (:<) <$> f x <*> traverse f xs + +-- | /O(1)/. Analyse the left end of a sequence. +viewl :: Seq a -> ViewL a +viewl (Seq xs) = case viewLTree xs of + EmptyLTree -> EmptyL + ConsLTree (Elem x) xs' -> x :< Seq xs' + +{-# SPECIALIZE viewLTree :: FingerTree (Elem a) -> ViewLTree (Elem a) #-} +{-# SPECIALIZE viewLTree :: FingerTree (Node a) -> ViewLTree (Node a) #-} +viewLTree :: Sized a => FingerTree a -> ViewLTree a +viewLTree EmptyT = EmptyLTree +viewLTree (Single a) = ConsLTree a EmptyT +viewLTree (Deep s (One a) m sf) = ConsLTree a (pullL (s - size a) m sf) +viewLTree (Deep s (Two a b) m sf) = + ConsLTree a (Deep (s - size a) (One b) m sf) +viewLTree (Deep s (Three a b c) m sf) = + ConsLTree a (Deep (s - size a) (Two b c) m sf) +viewLTree (Deep s (Four a b c d) m sf) = + ConsLTree a (Deep (s - size a) (Three b c d) m sf) + +-- | View of the right end of a sequence. +data ViewR a + = EmptyR -- ^ empty sequence + | Seq a :> a -- ^ the sequence minus the rightmost element, + -- and the rightmost element + deriving (Eq, Ord, Show, Read) + +#if __GLASGOW_HASKELL__ +deriving instance Data a => Data (ViewR a) +#endif +#if __GLASGOW_HASKELL__ >= 706 +deriving instance Generic1 ViewR +#endif +#if __GLASGOW_HASKELL__ >= 702 +deriving instance Generic (ViewR a) +#endif + +INSTANCE_TYPEABLE1(ViewR) + +instance Functor ViewR where + {-# INLINE fmap #-} + fmap _ EmptyR = EmptyR + fmap f (xs :> x) = fmap f xs :> f x + +instance Foldable ViewR where + foldMap _ EmptyR = mempty + foldMap f (xs :> x) = foldMap f xs <> f x + + foldr _ z EmptyR = z + foldr f z (xs :> x) = foldr f (f x z) xs + + foldl _ z EmptyR = z + foldl f z (xs :> x) = foldl f z xs `f` x + + foldr1 _ EmptyR = error "foldr1: empty view" + foldr1 f (xs :> x) = foldr f x xs +#if MIN_VERSION_base(4,8,0) + null EmptyR = True + null (_ :> _) = False + + length EmptyR = 0 + length (xs :> _) = length xs + 1 +#endif + +instance Traversable ViewR where + traverse _ EmptyR = pure EmptyR + traverse f (xs :> x) = (:>) <$> traverse f xs <*> f x + +-- | /O(1)/. Analyse the right end of a sequence. +viewr :: Seq a -> ViewR a +viewr (Seq xs) = case viewRTree xs of + EmptyRTree -> EmptyR + SnocRTree xs' (Elem x) -> Seq xs' :> x + +{-# SPECIALIZE viewRTree :: FingerTree (Elem a) -> ViewRTree (Elem a) #-} +{-# SPECIALIZE viewRTree :: FingerTree (Node a) -> ViewRTree (Node a) #-} +viewRTree :: Sized a => FingerTree a -> ViewRTree a +viewRTree EmptyT = EmptyRTree +viewRTree (Single z) = SnocRTree EmptyT z +viewRTree (Deep s pr m (One z)) = SnocRTree (pullR (s - size z) pr m) z +viewRTree (Deep s pr m (Two y z)) = + SnocRTree (Deep (s - size z) pr m (One y)) z +viewRTree (Deep s pr m (Three x y z)) = + SnocRTree (Deep (s - size z) pr m (Two x y)) z +viewRTree (Deep s pr m (Four w x y z)) = + SnocRTree (Deep (s - size z) pr m (Three w x y)) z + +------------------------------------------------------------------------ +-- Scans +-- +-- These are not particularly complex applications of the Traversable +-- functor, though making the correspondence with Data.List exact +-- requires the use of (<|) and (|>). +-- +-- Note that save for the single (<|) or (|>), we maintain the original +-- structure of the Seq, not having to do any restructuring of our own. +-- +-- wasserman.louis@gmail.com, 5/23/09 +------------------------------------------------------------------------ + +-- | 'scanl' is similar to 'foldl', but returns a sequence of reduced +-- values from the left: +-- +-- > scanl f z (fromList [x1, x2, ...]) = fromList [z, z `f` x1, (z `f` x1) `f` x2, ...] +scanl :: (a -> b -> a) -> a -> Seq b -> Seq a +scanl f z0 xs = z0 <| snd (mapAccumL (\ x z -> let x' = f x z in (x', x')) z0 xs) + +-- | 'scanl1' is a variant of 'scanl' that has no starting value argument: +-- +-- > scanl1 f (fromList [x1, x2, ...]) = fromList [x1, x1 `f` x2, ...] +scanl1 :: (a -> a -> a) -> Seq a -> Seq a +scanl1 f xs = case viewl xs of + EmptyL -> error "scanl1 takes a nonempty sequence as an argument" + x :< xs' -> scanl f x xs' + +-- | 'scanr' is the right-to-left dual of 'scanl'. +scanr :: (a -> b -> b) -> b -> Seq a -> Seq b +scanr f z0 xs = snd (mapAccumR (\ z x -> let z' = f x z in (z', z')) z0 xs) |> z0 + +-- | 'scanr1' is a variant of 'scanr' that has no starting value argument. +scanr1 :: (a -> a -> a) -> Seq a -> Seq a +scanr1 f xs = case viewr xs of + EmptyR -> error "scanr1 takes a nonempty sequence as an argument" + xs' :> x -> scanr f x xs' + +-- Indexing + +-- | /O(log(min(i,n-i)))/. The element at the specified position, +-- counting from 0. The argument should thus be a non-negative +-- integer less than the size of the sequence. +-- If the position is out of range, 'index' fails with an error. +-- +-- prop> xs `index` i = toList xs !! i +-- +-- Caution: 'index' necessarily delays retrieving the requested +-- element until the result is forced. It can therefore lead to a space +-- leak if the result is stored, unforced, in another structure. To retrieve +-- an element immediately without forcing it, use 'lookup' or '(!?)'. +index :: Seq a -> Int -> a +index (Seq xs) i + -- See note on unsigned arithmetic in splitAt + | fromIntegral i < (fromIntegral (size xs) :: Word) = case lookupTree i xs of + Place _ (Elem x) -> x + | otherwise = error "index out of bounds" + +-- | /O(log(min(i,n-i)))/. The element at the specified position, +-- counting from 0. If the specified position is negative or at +-- least the length of the sequence, 'lookup' returns 'Nothing'. +-- +-- prop> 0 <= i < length xs ==> lookup i xs == Just (toList xs !! i) +-- prop> i < 0 || i >= length xs ==> lookup i xs = Nothing +-- +-- Unlike 'index', this can be used to retrieve an element without +-- forcing it. For example, to insert the fifth element of a sequence +-- @xs@ into a 'Data.Map.Lazy.Map' @m@ at key @k@, you could use +-- +-- @ +-- case lookup 5 xs of +-- Nothing -> m +-- Just x -> 'Data.Map.Lazy.insert' k x m +-- @ +-- +-- @since 0.5.8 +lookup :: Int -> Seq a -> Maybe a +lookup i (Seq xs) + -- Note: we perform the lookup *before* applying the Just constructor + -- to ensure that we don't hold a reference to the whole sequence in + -- a thunk. If we applied the Just constructor around the case, the + -- actual lookup wouldn't be performed unless and until the value was + -- forced. + | fromIntegral i < (fromIntegral (size xs) :: Word) = case lookupTree i xs of + Place _ (Elem x) -> Just x + | otherwise = Nothing + +-- | /O(log(min(i,n-i)))/. A flipped, infix version of `lookup`. +-- +-- @since 0.5.8 +(!?) :: Seq a -> Int -> Maybe a +(!?) = flip lookup + +data Place a = Place {-# UNPACK #-} !Int a +#if TESTING + deriving Show +#endif + +{-# SPECIALIZE lookupTree :: Int -> FingerTree (Elem a) -> Place (Elem a) #-} +{-# SPECIALIZE lookupTree :: Int -> FingerTree (Node a) -> Place (Node a) #-} +lookupTree :: Sized a => Int -> FingerTree a -> Place a +lookupTree !_ EmptyT = error "lookupTree of empty tree" +lookupTree i (Single x) = Place i x +lookupTree i (Deep _ pr m sf) + | i < spr = lookupDigit i pr + | i < spm = case lookupTree (i - spr) m of + Place i' xs -> lookupNode i' xs + | otherwise = lookupDigit (i - spm) sf + where + spr = size pr + spm = spr + size m + +{-# SPECIALIZE lookupNode :: Int -> Node (Elem a) -> Place (Elem a) #-} +{-# SPECIALIZE lookupNode :: Int -> Node (Node a) -> Place (Node a) #-} +lookupNode :: Sized a => Int -> Node a -> Place a +lookupNode i (Node2 _ a b) + | i < sa = Place i a + | otherwise = Place (i - sa) b + where + sa = size a +lookupNode i (Node3 _ a b c) + | i < sa = Place i a + | i < sab = Place (i - sa) b + | otherwise = Place (i - sab) c + where + sa = size a + sab = sa + size b + +{-# SPECIALIZE lookupDigit :: Int -> Digit (Elem a) -> Place (Elem a) #-} +{-# SPECIALIZE lookupDigit :: Int -> Digit (Node a) -> Place (Node a) #-} +lookupDigit :: Sized a => Int -> Digit a -> Place a +lookupDigit i (One a) = Place i a +lookupDigit i (Two a b) + | i < sa = Place i a + | otherwise = Place (i - sa) b + where + sa = size a +lookupDigit i (Three a b c) + | i < sa = Place i a + | i < sab = Place (i - sa) b + | otherwise = Place (i - sab) c + where + sa = size a + sab = sa + size b +lookupDigit i (Four a b c d) + | i < sa = Place i a + | i < sab = Place (i - sa) b + | i < sabc = Place (i - sab) c + | otherwise = Place (i - sabc) d + where + sa = size a + sab = sa + size b + sabc = sab + size c + +-- | /O(log(min(i,n-i)))/. Replace the element at the specified position. +-- If the position is out of range, the original sequence is returned. +update :: Int -> a -> Seq a -> Seq a +update i x (Seq xs) + -- See note on unsigned arithmetic in splitAt + | fromIntegral i < (fromIntegral (size xs) :: Word) = Seq (updateTree (Elem x) i xs) + | otherwise = Seq xs + +-- It seems a shame to copy the implementation of the top layer of +-- `adjust` instead of just using `update i x = adjust (const x) i`. +-- With the latter implementation, updating the same position many +-- times could lead to silly thunks building up around that position. +-- The thunks will each look like @const v a@, where @v@ is the new +-- value and @a@ the old. +updateTree :: Elem a -> Int -> FingerTree (Elem a) -> FingerTree (Elem a) +updateTree _ !_ EmptyT = EmptyT -- Unreachable +updateTree v _i (Single _) = Single v +updateTree v i (Deep s pr m sf) + | i < spr = Deep s (updateDigit v i pr) m sf + | i < spm = let !m' = adjustTree (updateNode v) (i - spr) m + in Deep s pr m' sf + | otherwise = Deep s pr m (updateDigit v (i - spm) sf) + where + spr = size pr + spm = spr + size m + +updateNode :: Elem a -> Int -> Node (Elem a) -> Node (Elem a) +updateNode v i (Node2 s a b) + | i < sa = Node2 s v b + | otherwise = Node2 s a v + where + sa = size a +updateNode v i (Node3 s a b c) + | i < sa = Node3 s v b c + | i < sab = Node3 s a v c + | otherwise = Node3 s a b v + where + sa = size a + sab = sa + size b + +updateDigit :: Elem a -> Int -> Digit (Elem a) -> Digit (Elem a) +updateDigit v !_i (One _) = One v +updateDigit v i (Two a b) + | i < sa = Two v b + | otherwise = Two a v + where + sa = size a +updateDigit v i (Three a b c) + | i < sa = Three v b c + | i < sab = Three a v c + | otherwise = Three a b v + where + sa = size a + sab = sa + size b +updateDigit v i (Four a b c d) + | i < sa = Four v b c d + | i < sab = Four a v c d + | i < sabc = Four a b v d + | otherwise = Four a b c v + where + sa = size a + sab = sa + size b + sabc = sab + size c + +-- | /O(log(min(i,n-i)))/. Update the element at the specified position. If +-- the position is out of range, the original sequence is returned. 'adjust' +-- can lead to poor performance and even memory leaks, because it does not +-- force the new value before installing it in the sequence. 'adjust'' should +-- usually be preferred. +adjust :: (a -> a) -> Int -> Seq a -> Seq a +adjust f i (Seq xs) + -- See note on unsigned arithmetic in splitAt + | fromIntegral i < (fromIntegral (size xs) :: Word) = Seq (adjustTree (`seq` fmap f) i xs) + | otherwise = Seq xs + +-- | /O(log(min(i,n-i)))/. Update the element at the specified position. +-- If the position is out of range, the original sequence is returned. +-- The new value is forced before it is installed in the sequence. +-- +-- @ +-- adjust' f i xs = +-- case xs !? i of +-- Nothing -> xs +-- Just x -> let !x' = f x +-- in update i x' xs +-- @ +-- +-- @since 0.5.8 +adjust' :: forall a . (a -> a) -> Int -> Seq a -> Seq a +#if __GLASGOW_HASKELL__ >= 708 +adjust' f i xs + -- See note on unsigned arithmetic in splitAt + | fromIntegral i < (fromIntegral (length xs) :: Word) = + coerce $ adjustTree (\ !_k (ForceBox a) -> ForceBox (f a)) i (coerce xs) + | otherwise = xs +#else +-- This is inefficient, but fixing it would take a lot of fuss and bother +-- for little immediate gain. We can deal with that when we have another +-- Haskell implementation to worry about. +adjust' f i xs = + case xs !? i of + Nothing -> xs + Just x -> let !x' = f x + in update i x' xs +#endif + +{-# SPECIALIZE adjustTree :: (Int -> ForceBox a -> ForceBox a) -> Int -> FingerTree (ForceBox a) -> FingerTree (ForceBox a) #-} +{-# SPECIALIZE adjustTree :: (Int -> Elem a -> Elem a) -> Int -> FingerTree (Elem a) -> FingerTree (Elem a) #-} +{-# SPECIALIZE adjustTree :: (Int -> Node a -> Node a) -> Int -> FingerTree (Node a) -> FingerTree (Node a) #-} +adjustTree :: (Sized a, MaybeForce a) => (Int -> a -> a) -> + Int -> FingerTree a -> FingerTree a +adjustTree _ !_ EmptyT = EmptyT -- Unreachable +adjustTree f i (Single x) = Single $!? f i x +adjustTree f i (Deep s pr m sf) + | i < spr = Deep s (adjustDigit f i pr) m sf + | i < spm = let !m' = adjustTree (adjustNode f) (i - spr) m + in Deep s pr m' sf + | otherwise = Deep s pr m (adjustDigit f (i - spm) sf) + where + spr = size pr + spm = spr + size m + +{-# SPECIALIZE adjustNode :: (Int -> Elem a -> Elem a) -> Int -> Node (Elem a) -> Node (Elem a) #-} +{-# SPECIALIZE adjustNode :: (Int -> Node a -> Node a) -> Int -> Node (Node a) -> Node (Node a) #-} +adjustNode :: (Sized a, MaybeForce a) => (Int -> a -> a) -> Int -> Node a -> Node a +adjustNode f i (Node2 s a b) + | i < sa = let fia = f i a in fia `mseq` Node2 s fia b + | otherwise = let fisab = f (i - sa) b in fisab `mseq` Node2 s a fisab + where + sa = size a +adjustNode f i (Node3 s a b c) + | i < sa = let fia = f i a in fia `mseq` Node3 s fia b c + | i < sab = let fisab = f (i - sa) b in fisab `mseq` Node3 s a fisab c + | otherwise = let fisabc = f (i - sab) c in fisabc `mseq` Node3 s a b fisabc + where + sa = size a + sab = sa + size b + +{-# SPECIALIZE adjustDigit :: (Int -> Elem a -> Elem a) -> Int -> Digit (Elem a) -> Digit (Elem a) #-} +{-# SPECIALIZE adjustDigit :: (Int -> Node a -> Node a) -> Int -> Digit (Node a) -> Digit (Node a) #-} +adjustDigit :: (Sized a, MaybeForce a) => (Int -> a -> a) -> Int -> Digit a -> Digit a +adjustDigit f !i (One a) = One $!? f i a +adjustDigit f i (Two a b) + | i < sa = let fia = f i a in fia `mseq` Two fia b + | otherwise = let fisab = f (i - sa) b in fisab `mseq` Two a fisab + where + sa = size a +adjustDigit f i (Three a b c) + | i < sa = let fia = f i a in fia `mseq` Three fia b c + | i < sab = let fisab = f (i - sa) b in fisab `mseq` Three a fisab c + | otherwise = let fisabc = f (i - sab) c in fisabc `mseq` Three a b fisabc + where + sa = size a + sab = sa + size b +adjustDigit f i (Four a b c d) + | i < sa = let fia = f i a in fia `mseq` Four fia b c d + | i < sab = let fisab = f (i - sa) b in fisab `mseq` Four a fisab c d + | i < sabc = let fisabc = f (i - sab) c in fisabc `mseq` Four a b fisabc d + | otherwise = let fisabcd = f (i - sabc) d in fisabcd `mseq` Four a b c fisabcd + where + sa = size a + sab = sa + size b + sabc = sab + size c + +-- | /O(log(min(i,n-i)))/. @'insertAt' i x xs@ inserts @x@ into @xs@ +-- at the index @i@, shifting the rest of the sequence over. +-- +-- @ +-- insertAt 2 x (fromList [a,b,c,d]) = fromList [a,b,x,c,d] +-- insertAt 4 x (fromList [a,b,c,d]) = insertAt 10 x (fromList [a,b,c,d]) +-- = fromList [a,b,c,d,x] +-- @ +-- +-- prop> insertAt i x xs = take i xs >< singleton x >< drop i xs +-- +-- @since 0.5.8 +insertAt :: Int -> a -> Seq a -> Seq a +insertAt i a s@(Seq xs) + | fromIntegral i < (fromIntegral (size xs) :: Word) + = Seq (insTree (`seq` InsTwo (Elem a)) i xs) + | i <= 0 = a <| s + | otherwise = s |> a + +data Ins a = InsOne a | InsTwo a a + +{-# SPECIALIZE insTree :: (Int -> Elem a -> Ins (Elem a)) -> Int -> FingerTree (Elem a) -> FingerTree (Elem a) #-} +{-# SPECIALIZE insTree :: (Int -> Node a -> Ins (Node a)) -> Int -> FingerTree (Node a) -> FingerTree (Node a) #-} +insTree :: Sized a => (Int -> a -> Ins a) -> + Int -> FingerTree a -> FingerTree a +insTree _ !_ EmptyT = EmptyT -- Unreachable +insTree f i (Single x) = case f i x of + InsOne x' -> Single x' + InsTwo m n -> deep (One m) EmptyT (One n) +insTree f i (Deep s pr m sf) + | i < spr = case insLeftDigit f i pr of + InsLeftDig pr' -> Deep (s + 1) pr' m sf + InsDigNode pr' n -> m `seq` Deep (s + 1) pr' (n `consTree` m) sf + | i < spm = let !m' = insTree (insNode f) (i - spr) m + in Deep (s + 1) pr m' sf + | otherwise = case insRightDigit f (i - spm) sf of + InsRightDig sf' -> Deep (s + 1) pr m sf' + InsNodeDig n sf' -> m `seq` Deep (s + 1) pr (m `snocTree` n) sf' + where + spr = size pr + spm = spr + size m + +{-# SPECIALIZE insNode :: (Int -> Elem a -> Ins (Elem a)) -> Int -> Node (Elem a) -> Ins (Node (Elem a)) #-} +{-# SPECIALIZE insNode :: (Int -> Node a -> Ins (Node a)) -> Int -> Node (Node a) -> Ins (Node (Node a)) #-} +insNode :: Sized a => (Int -> a -> Ins a) -> Int -> Node a -> Ins (Node a) +insNode f i (Node2 s a b) + | i < sa = case f i a of + InsOne n -> InsOne $ Node2 (s + 1) n b + InsTwo m n -> InsOne $ Node3 (s + 1) m n b + | otherwise = case f (i - sa) b of + InsOne n -> InsOne $ Node2 (s + 1) a n + InsTwo m n -> InsOne $ Node3 (s + 1) a m n + where sa = size a +insNode f i (Node3 s a b c) + | i < sa = case f i a of + InsOne n -> InsOne $ Node3 (s + 1) n b c + InsTwo m n -> InsTwo (Node2 (sa + 1) m n) (Node2 (s - sa) b c) + | i < sab = case f (i - sa) b of + InsOne n -> InsOne $ Node3 (s + 1) a n c + InsTwo m n -> InsTwo am nc + where !am = node2 a m + !nc = node2 n c + | otherwise = case f (i - sab) c of + InsOne n -> InsOne $ Node3 (s + 1) a b n + InsTwo m n -> InsTwo (Node2 sab a b) (Node2 (s - sab + 1) m n) + where sa = size a + sab = sa + size b + +data InsDigNode a = InsLeftDig !(Digit a) | InsDigNode !(Digit a) !(Node a) +{-# SPECIALIZE insLeftDigit :: (Int -> Elem a -> Ins (Elem a)) -> Int -> Digit (Elem a) -> InsDigNode (Elem a) #-} +{-# SPECIALIZE insLeftDigit :: (Int -> Node a -> Ins (Node a)) -> Int -> Digit (Node a) -> InsDigNode (Node a) #-} +insLeftDigit :: Sized a => (Int -> a -> Ins a) -> Int -> Digit a -> InsDigNode a +insLeftDigit f !i (One a) = case f i a of + InsOne a' -> InsLeftDig $ One a' + InsTwo a1 a2 -> InsLeftDig $ Two a1 a2 +insLeftDigit f i (Two a b) + | i < sa = case f i a of + InsOne a' -> InsLeftDig $ Two a' b + InsTwo a1 a2 -> InsLeftDig $ Three a1 a2 b + | otherwise = case f (i - sa) b of + InsOne b' -> InsLeftDig $ Two a b' + InsTwo b1 b2 -> InsLeftDig $ Three a b1 b2 + where sa = size a +insLeftDigit f i (Three a b c) + | i < sa = case f i a of + InsOne a' -> InsLeftDig $ Three a' b c + InsTwo a1 a2 -> InsLeftDig $ Four a1 a2 b c + | i < sab = case f (i - sa) b of + InsOne b' -> InsLeftDig $ Three a b' c + InsTwo b1 b2 -> InsLeftDig $ Four a b1 b2 c + | otherwise = case f (i - sab) c of + InsOne c' -> InsLeftDig $ Three a b c' + InsTwo c1 c2 -> InsLeftDig $ Four a b c1 c2 + where sa = size a + sab = sa + size b +insLeftDigit f i (Four a b c d) + | i < sa = case f i a of + InsOne a' -> InsLeftDig $ Four a' b c d + InsTwo a1 a2 -> InsDigNode (Two a1 a2) (node3 b c d) + | i < sab = case f (i - sa) b of + InsOne b' -> InsLeftDig $ Four a b' c d + InsTwo b1 b2 -> InsDigNode (Two a b1) (node3 b2 c d) + | i < sabc = case f (i - sab) c of + InsOne c' -> InsLeftDig $ Four a b c' d + InsTwo c1 c2 -> InsDigNode (Two a b) (node3 c1 c2 d) + | otherwise = case f (i - sabc) d of + InsOne d' -> InsLeftDig $ Four a b c d' + InsTwo d1 d2 -> InsDigNode (Two a b) (node3 c d1 d2) + where sa = size a + sab = sa + size b + sabc = sab + size c + +data InsNodeDig a = InsRightDig !(Digit a) | InsNodeDig !(Node a) !(Digit a) +{-# SPECIALIZE insRightDigit :: (Int -> Elem a -> Ins (Elem a)) -> Int -> Digit (Elem a) -> InsNodeDig (Elem a) #-} +{-# SPECIALIZE insRightDigit :: (Int -> Node a -> Ins (Node a)) -> Int -> Digit (Node a) -> InsNodeDig (Node a) #-} +insRightDigit :: Sized a => (Int -> a -> Ins a) -> Int -> Digit a -> InsNodeDig a +insRightDigit f !i (One a) = case f i a of + InsOne a' -> InsRightDig $ One a' + InsTwo a1 a2 -> InsRightDig $ Two a1 a2 +insRightDigit f i (Two a b) + | i < sa = case f i a of + InsOne a' -> InsRightDig $ Two a' b + InsTwo a1 a2 -> InsRightDig $ Three a1 a2 b + | otherwise = case f (i - sa) b of + InsOne b' -> InsRightDig $ Two a b' + InsTwo b1 b2 -> InsRightDig $ Three a b1 b2 + where sa = size a +insRightDigit f i (Three a b c) + | i < sa = case f i a of + InsOne a' -> InsRightDig $ Three a' b c + InsTwo a1 a2 -> InsRightDig $ Four a1 a2 b c + | i < sab = case f (i - sa) b of + InsOne b' -> InsRightDig $ Three a b' c + InsTwo b1 b2 -> InsRightDig $ Four a b1 b2 c + | otherwise = case f (i - sab) c of + InsOne c' -> InsRightDig $ Three a b c' + InsTwo c1 c2 -> InsRightDig $ Four a b c1 c2 + where sa = size a + sab = sa + size b +insRightDigit f i (Four a b c d) + | i < sa = case f i a of + InsOne a' -> InsRightDig $ Four a' b c d + InsTwo a1 a2 -> InsNodeDig (node3 a1 a2 b) (Two c d) + | i < sab = case f (i - sa) b of + InsOne b' -> InsRightDig $ Four a b' c d + InsTwo b1 b2 -> InsNodeDig (node3 a b1 b2) (Two c d) + | i < sabc = case f (i - sab) c of + InsOne c' -> InsRightDig $ Four a b c' d + InsTwo c1 c2 -> InsNodeDig (node3 a b c1) (Two c2 d) + | otherwise = case f (i - sabc) d of + InsOne d' -> InsRightDig $ Four a b c d' + InsTwo d1 d2 -> InsNodeDig (node3 a b c) (Two d1 d2) + where sa = size a + sab = sa + size b + sabc = sab + size c + +-- | /O(log(min(i,n-i)))/. Delete the element of a sequence at a given +-- index. Return the original sequence if the index is out of range. +-- +-- @ +-- deleteAt 2 [a,b,c,d] = [a,b,d] +-- deleteAt 4 [a,b,c,d] = deleteAt (-1) [a,b,c,d] = [a,b,c,d] +-- @ +-- +-- @since 0.5.8 +deleteAt :: Int -> Seq a -> Seq a +deleteAt i (Seq xs) + | fromIntegral i < (fromIntegral (size xs) :: Word) = Seq $ delTreeE i xs + | otherwise = Seq xs + +delTreeE :: Int -> FingerTree (Elem a) -> FingerTree (Elem a) +delTreeE !_i EmptyT = EmptyT -- Unreachable +delTreeE _i Single{} = EmptyT +delTreeE i (Deep s pr m sf) + | i < spr = delLeftDigitE i s pr m sf + | i < spm = case delTree delNodeE (i - spr) m of + FullTree m' -> Deep (s - 1) pr m' sf + DefectTree e -> delRebuildMiddle (s - 1) pr e sf + | otherwise = delRightDigitE (i - spm) s pr m sf + where spr = size pr + spm = spr + size m + +delNodeE :: Int -> Node (Elem a) -> Del (Elem a) +delNodeE i (Node3 _ a b c) = case i of + 0 -> Full $ Node2 2 b c + 1 -> Full $ Node2 2 a c + _ -> Full $ Node2 2 a b +delNodeE i (Node2 _ a b) = case i of + 0 -> Defect b + _ -> Defect a + + +delLeftDigitE :: Int -> Int -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> FingerTree (Elem a) +delLeftDigitE !_i s One{} m sf = pullL (s - 1) m sf +delLeftDigitE i s (Two a b) m sf + | i == 0 = Deep (s - 1) (One b) m sf + | otherwise = Deep (s - 1) (One a) m sf +delLeftDigitE i s (Three a b c) m sf + | i == 0 = Deep (s - 1) (Two b c) m sf + | i == 1 = Deep (s - 1) (Two a c) m sf + | otherwise = Deep (s - 1) (Two a b) m sf +delLeftDigitE i s (Four a b c d) m sf + | i == 0 = Deep (s - 1) (Three b c d) m sf + | i == 1 = Deep (s - 1) (Three a c d) m sf + | i == 2 = Deep (s - 1) (Three a b d) m sf + | otherwise = Deep (s - 1) (Three a b c) m sf + +delRightDigitE :: Int -> Int -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> FingerTree (Elem a) +delRightDigitE !_i s pr m One{} = pullR (s - 1) pr m +delRightDigitE i s pr m (Two a b) + | i == 0 = Deep (s - 1) pr m (One b) + | otherwise = Deep (s - 1) pr m (One a) +delRightDigitE i s pr m (Three a b c) + | i == 0 = Deep (s - 1) pr m (Two b c) + | i == 1 = Deep (s - 1) pr m (Two a c) + | otherwise = deep pr m (Two a b) +delRightDigitE i s pr m (Four a b c d) + | i == 0 = Deep (s - 1) pr m (Three b c d) + | i == 1 = Deep (s - 1) pr m (Three a c d) + | i == 2 = Deep (s - 1) pr m (Three a b d) + | otherwise = Deep (s - 1) pr m (Three a b c) + +data DelTree a = FullTree !(FingerTree (Node a)) | DefectTree a + +{-# SPECIALIZE delTree :: (Int -> Node (Elem a) -> Del (Elem a)) -> Int -> FingerTree (Node (Elem a)) -> DelTree (Elem a) #-} +{-# SPECIALIZE delTree :: (Int -> Node (Node a) -> Del (Node a)) -> Int -> FingerTree (Node (Node a)) -> DelTree (Node a) #-} +delTree :: Sized a => (Int -> Node a -> Del a) -> Int -> FingerTree (Node a) -> DelTree a +delTree _f !_i EmptyT = FullTree EmptyT -- Unreachable +delTree f i (Single a) = case f i a of + Full a' -> FullTree (Single a') + Defect e -> DefectTree e +delTree f i (Deep s pr m sf) + | i < spr = case delDigit f i pr of + FullDig pr' -> FullTree $ Deep (s - 1) pr' m sf + DefectDig e -> case viewLTree m of + EmptyLTree -> FullTree $ delRebuildRightDigit (s - 1) e sf + ConsLTree n m' -> FullTree $ delRebuildLeftSide (s - 1) e n m' sf + | i < spm = case delTree (delNode f) (i - spr) m of + FullTree m' -> FullTree (Deep (s - 1) pr m' sf) + DefectTree e -> FullTree $ delRebuildMiddle (s - 1) pr e sf + | otherwise = case delDigit f (i - spm) sf of + FullDig sf' -> FullTree $ Deep (s - 1) pr m sf' + DefectDig e -> case viewRTree m of + EmptyRTree -> FullTree $ delRebuildLeftDigit (s - 1) pr e + SnocRTree m' n -> FullTree $ delRebuildRightSide (s - 1) pr m' n e + where spr = size pr + spm = spr + size m + +data Del a = Full !(Node a) | Defect a + +{-# SPECIALIZE delNode :: (Int -> Node (Elem a) -> Del (Elem a)) -> Int -> Node (Node (Elem a)) -> Del (Node (Elem a)) #-} +{-# SPECIALIZE delNode :: (Int -> Node (Node a) -> Del (Node a)) -> Int -> Node (Node (Node a)) -> Del (Node (Node a)) #-} +delNode :: Sized a => (Int -> Node a -> Del a) -> Int -> Node (Node a) -> Del (Node a) +delNode f i (Node3 s a b c) + | i < sa = case f i a of + Full a' -> Full $ Node3 (s - 1) a' b c + Defect e -> let !se = size e in case b of + Node3 sxyz x y z -> Full $ Node3 (s - 1) (Node2 (se + sx) e x) (Node2 (sxyz - sx) y z) c + where !sx = size x + Node2 sxy x y -> Full $ Node2 (s - 1) (Node3 (sxy + se) e x y) c + | i < sab = case f (i - sa) b of + Full b' -> Full $ Node3 (s - 1) a b' c + Defect e -> let !se = size e in case a of + Node3 sxyz x y z -> Full $ Node3 (s - 1) (Node2 (sxyz - sz) x y) (Node2 (sz + se) z e) c + where !sz = size z + Node2 sxy x y -> Full $ Node2 (s - 1) (Node3 (sxy + se) x y e) c + | otherwise = case f (i - sab) c of + Full c' -> Full $ Node3 (s - 1) a b c' + Defect e -> let !se = size e in case b of + Node3 sxyz x y z -> Full $ Node3 (s - 1) a (Node2 (sxyz - sz) x y) (Node2 (sz + se) z e) + where !sz = size z + Node2 sxy x y -> Full $ Node2 (s - 1) a (Node3 (sxy + se) x y e) + where sa = size a + sab = sa + size b +delNode f i (Node2 s a b) + | i < sa = case f i a of + Full a' -> Full $ Node2 (s - 1) a' b + Defect e -> let !se = size e in case b of + Node3 sxyz x y z -> Full $ Node2 (s - 1) (Node2 (se + sx) e x) (Node2 (sxyz - sx) y z) + where !sx = size x + Node2 _ x y -> Defect $ Node3 (s - 1) e x y + | otherwise = case f (i - sa) b of + Full b' -> Full $ Node2 (s - 1) a b' + Defect e -> let !se = size e in case a of + Node3 sxyz x y z -> Full $ Node2 (s - 1) (Node2 (sxyz - sz) x y) (Node2 (sz + se) z e) + where !sz = size z + Node2 _ x y -> Defect $ Node3 (s - 1) x y e + where sa = size a + +{-# SPECIALIZE delRebuildRightDigit :: Int -> Elem a -> Digit (Node (Elem a)) -> FingerTree (Node (Elem a)) #-} +{-# SPECIALIZE delRebuildRightDigit :: Int -> Node a -> Digit (Node (Node a)) -> FingerTree (Node (Node a)) #-} +delRebuildRightDigit :: Sized a => Int -> a -> Digit (Node a) -> FingerTree (Node a) +delRebuildRightDigit s p (One a) = let !sp = size p in case a of + Node3 sxyz x y z -> Deep s (One (Node2 (sp + sx) p x)) EmptyT (One (Node2 (sxyz - sx) y z)) + where !sx = size x + Node2 sxy x y -> Single (Node3 (sp + sxy) p x y) +delRebuildRightDigit s p (Two a b) = let !sp = size p in case a of + Node3 sxyz x y z -> Deep s (Two (Node2 (sp + sx) p x) (Node2 (sxyz - sx) y z)) EmptyT (One b) + where !sx = size x + Node2 sxy x y -> Deep s (One (Node3 (sp + sxy) p x y)) EmptyT (One b) +delRebuildRightDigit s p (Three a b c) = let !sp = size p in case a of + Node3 sxyz x y z -> Deep s (Two (Node2 (sp + sx) p x) (Node2 (sxyz - sx) y z)) EmptyT (Two b c) + where !sx = size x + Node2 sxy x y -> Deep s (Two (Node3 (sp + sxy) p x y) b) EmptyT (One c) +delRebuildRightDigit s p (Four a b c d) = let !sp = size p in case a of + Node3 sxyz x y z -> Deep s (Three (Node2 (sp + sx) p x) (Node2 (sxyz - sx) y z) b) EmptyT (Two c d) + where !sx = size x + Node2 sxy x y -> Deep s (Two (Node3 (sp + sxy) p x y) b) EmptyT (Two c d) + +{-# SPECIALIZE delRebuildLeftDigit :: Int -> Digit (Node (Elem a)) -> Elem a -> FingerTree (Node (Elem a)) #-} +{-# SPECIALIZE delRebuildLeftDigit :: Int -> Digit (Node (Node a)) -> Node a -> FingerTree (Node (Node a)) #-} +delRebuildLeftDigit :: Sized a => Int -> Digit (Node a) -> a -> FingerTree (Node a) +delRebuildLeftDigit s (One a) p = let !sp = size p in case a of + Node3 sxyz x y z -> Deep s (One (Node2 (sxyz - sz) x y)) EmptyT (One (Node2 (sz + sp) z p)) + where !sz = size z + Node2 sxy x y -> Single (Node3 (sxy + sp) x y p) +delRebuildLeftDigit s (Two a b) p = let !sp = size p in case b of + Node3 sxyz x y z -> Deep s (Two a (Node2 (sxyz - sz) x y)) EmptyT (One (Node2 (sz + sp) z p)) + where !sz = size z + Node2 sxy x y -> Deep s (One a) EmptyT (One (Node3 (sxy + sp) x y p)) +delRebuildLeftDigit s (Three a b c) p = let !sp = size p in case c of + Node3 sxyz x y z -> Deep s (Two a b) EmptyT (Two (Node2 (sxyz - sz) x y) (Node2 (sz + sp) z p)) + where !sz = size z + Node2 sxy x y -> Deep s (Two a b) EmptyT (One (Node3 (sxy + sp) x y p)) +delRebuildLeftDigit s (Four a b c d) p = let !sp = size p in case d of + Node3 sxyz x y z -> Deep s (Three a b c) EmptyT (Two (Node2 (sxyz - sz) x y) (Node2 (sz + sp) z p)) + where !sz = size z + Node2 sxy x y -> Deep s (Two a b) EmptyT (Two c (Node3 (sxy + sp) x y p)) + +delRebuildLeftSide :: Sized a + => Int -> a -> Node (Node a) -> FingerTree (Node (Node a)) -> Digit (Node a) + -> FingerTree (Node a) +delRebuildLeftSide s p (Node2 _ a b) m sf = let !sp = size p in case a of + Node2 sxy x y -> Deep s (Two (Node3 (sp + sxy) p x y) b) m sf + Node3 sxyz x y z -> Deep s (Three (Node2 (sp + sx) p x) (Node2 (sxyz - sx) y z) b) m sf + where !sx = size x +delRebuildLeftSide s p (Node3 _ a b c) m sf = let !sp = size p in case a of + Node2 sxy x y -> Deep s (Three (Node3 (sp + sxy) p x y) b c) m sf + Node3 sxyz x y z -> Deep s (Four (Node2 (sp + sx) p x) (Node2 (sxyz - sx) y z) b c) m sf + where !sx = size x + +delRebuildRightSide :: Sized a + => Int -> Digit (Node a) -> FingerTree (Node (Node a)) -> Node (Node a) -> a + -> FingerTree (Node a) +delRebuildRightSide s pr m (Node2 _ a b) p = let !sp = size p in case b of + Node2 sxy x y -> Deep s pr m (Two a (Node3 (sxy + sp) x y p)) + Node3 sxyz x y z -> Deep s pr m (Three a (Node2 (sxyz - sz) x y) (Node2 (sz + sp) z p)) + where !sz = size z +delRebuildRightSide s pr m (Node3 _ a b c) p = let !sp = size p in case c of + Node2 sxy x y -> Deep s pr m (Three a b (Node3 (sxy + sp) x y p)) + Node3 sxyz x y z -> Deep s pr m (Four a b (Node2 (sxyz - sz) x y) (Node2 (sz + sp) z p)) + where !sz = size z + +delRebuildMiddle :: Sized a + => Int -> Digit a -> a -> Digit a + -> FingerTree a +delRebuildMiddle s (One a) e sf = Deep s (Two a e) EmptyT sf +delRebuildMiddle s (Two a b) e sf = Deep s (Three a b e) EmptyT sf +delRebuildMiddle s (Three a b c) e sf = Deep s (Four a b c e) EmptyT sf +delRebuildMiddle s (Four a b c d) e sf = Deep s (Two a b) (Single (node3 c d e)) sf + +data DelDig a = FullDig !(Digit (Node a)) | DefectDig a + +{-# SPECIALIZE delDigit :: (Int -> Node (Elem a) -> Del (Elem a)) -> Int -> Digit (Node (Elem a)) -> DelDig (Elem a) #-} +{-# SPECIALIZE delDigit :: (Int -> Node (Node a) -> Del (Node a)) -> Int -> Digit (Node (Node a)) -> DelDig (Node a) #-} +delDigit :: Sized a => (Int -> Node a -> Del a) -> Int -> Digit (Node a) -> DelDig a +delDigit f !i (One a) = case f i a of + Full a' -> FullDig $ One a' + Defect e -> DefectDig e +delDigit f i (Two a b) + | i < sa = case f i a of + Full a' -> FullDig $ Two a' b + Defect e -> let !se = size e in case b of + Node3 sxyz x y z -> FullDig $ Two (Node2 (se + sx) e x) (Node2 (sxyz - sx) y z) + where !sx = size x + Node2 sxy x y -> FullDig $ One (Node3 (se + sxy) e x y) + | otherwise = case f (i - sa) b of + Full b' -> FullDig $ Two a b' + Defect e -> let !se = size e in case a of + Node3 sxyz x y z -> FullDig $ Two (Node2 (sxyz - sz) x y) (Node2 (sz + se) z e) + where !sz = size z + Node2 sxy x y -> FullDig $ One (Node3 (sxy + se) x y e) + where sa = size a +delDigit f i (Three a b c) + | i < sa = case f i a of + Full a' -> FullDig $ Three a' b c + Defect e -> let !se = size e in case b of + Node3 sxyz x y z -> FullDig $ Three (Node2 (se + sx) e x) (Node2 (sxyz - sx) y z) c + where !sx = size x + Node2 sxy x y -> FullDig $ Two (Node3 (se + sxy) e x y) c + | i < sab = case f (i - sa) b of + Full b' -> FullDig $ Three a b' c + Defect e -> let !se = size e in case a of + Node3 sxyz x y z -> FullDig $ Three (Node2 (sxyz - sz) x y) (Node2 (sz + se) z e) c + where !sz = size z + Node2 sxy x y -> FullDig $ Two (Node3 (sxy + se) x y e) c + | otherwise = case f (i - sab) c of + Full c' -> FullDig $ Three a b c' + Defect e -> let !se = size e in case b of + Node3 sxyz x y z -> FullDig $ Three a (Node2 (sxyz - sz) x y) (Node2 (sz + se) z e) + where !sz = size z + Node2 sxy x y -> FullDig $ Two a (Node3 (sxy + se) x y e) + where sa = size a + sab = sa + size b +delDigit f i (Four a b c d) + | i < sa = case f i a of + Full a' -> FullDig $ Four a' b c d + Defect e -> let !se = size e in case b of + Node3 sxyz x y z -> FullDig $ Four (Node2 (se + sx) e x) (Node2 (sxyz - sx) y z) c d + where !sx = size x + Node2 sxy x y -> FullDig $ Three (Node3 (se + sxy) e x y) c d + | i < sab = case f (i - sa) b of + Full b' -> FullDig $ Four a b' c d + Defect e -> let !se = size e in case a of + Node3 sxyz x y z -> FullDig $ Four (Node2 (sxyz - sz) x y) (Node2 (sz + se) z e) c d + where !sz = size z + Node2 sxy x y -> FullDig $ Three (Node3 (sxy + se) x y e) c d + | i < sabc = case f (i - sab) c of + Full c' -> FullDig $ Four a b c' d + Defect e -> let !se = size e in case b of + Node3 sxyz x y z -> FullDig $ Four a (Node2 (sxyz - sz) x y) (Node2 (sz + se) z e) d + where !sz = size z + Node2 sxy x y -> FullDig $ Three a (Node3 (sxy + se) x y e) d + | otherwise = case f (i - sabc) d of + Full d' -> FullDig $ Four a b c d' + Defect e -> let !se = size e in case c of + Node3 sxyz x y z -> FullDig $ Four a b (Node2 (sxyz - sz) x y) (Node2 (sz + se) z e) + where !sz = size z + Node2 sxy x y -> FullDig $ Three a b (Node3 (sxy + se) x y e) + where sa = size a + sab = sa + size b + sabc = sab + size c + + +-- | /O(n)/. A generalization of 'fmap', 'mapWithIndex' takes a mapping +-- function that also depends on the element's index, and applies it to every +-- element in the sequence. +mapWithIndex :: (Int -> a -> b) -> Seq a -> Seq b +mapWithIndex f' (Seq xs') = Seq $ mapWithIndexTree (\s (Elem a) -> Elem (f' s a)) 0 xs' + where + {-# SPECIALIZE mapWithIndexTree :: (Int -> Elem y -> b) -> Int -> FingerTree (Elem y) -> FingerTree b #-} + {-# SPECIALIZE mapWithIndexTree :: (Int -> Node y -> b) -> Int -> FingerTree (Node y) -> FingerTree b #-} + mapWithIndexTree :: Sized a => (Int -> a -> b) -> Int -> FingerTree a -> FingerTree b + mapWithIndexTree _ !_s EmptyT = EmptyT + mapWithIndexTree f s (Single xs) = Single $ f s xs + mapWithIndexTree f s (Deep n pr m sf) = + Deep n + (mapWithIndexDigit f s pr) + (mapWithIndexTree (mapWithIndexNode f) sPspr m) + (mapWithIndexDigit f sPsprm sf) + where + !sPspr = s + size pr + !sPsprm = sPspr + size m + + {-# SPECIALIZE mapWithIndexDigit :: (Int -> Elem y -> b) -> Int -> Digit (Elem y) -> Digit b #-} + {-# SPECIALIZE mapWithIndexDigit :: (Int -> Node y -> b) -> Int -> Digit (Node y) -> Digit b #-} + mapWithIndexDigit :: Sized a => (Int -> a -> b) -> Int -> Digit a -> Digit b + mapWithIndexDigit f !s (One a) = One (f s a) + mapWithIndexDigit f s (Two a b) = Two (f s a) (f sPsa b) + where + !sPsa = s + size a + mapWithIndexDigit f s (Three a b c) = + Three (f s a) (f sPsa b) (f sPsab c) + where + !sPsa = s + size a + !sPsab = sPsa + size b + mapWithIndexDigit f s (Four a b c d) = + Four (f s a) (f sPsa b) (f sPsab c) (f sPsabc d) + where + !sPsa = s + size a + !sPsab = sPsa + size b + !sPsabc = sPsab + size c + + {-# SPECIALIZE mapWithIndexNode :: (Int -> Elem y -> b) -> Int -> Node (Elem y) -> Node b #-} + {-# SPECIALIZE mapWithIndexNode :: (Int -> Node y -> b) -> Int -> Node (Node y) -> Node b #-} + mapWithIndexNode :: Sized a => (Int -> a -> b) -> Int -> Node a -> Node b + mapWithIndexNode f s (Node2 ns a b) = Node2 ns (f s a) (f sPsa b) + where + !sPsa = s + size a + mapWithIndexNode f s (Node3 ns a b c) = + Node3 ns (f s a) (f sPsa b) (f sPsab c) + where + !sPsa = s + size a + !sPsab = sPsa + size b + +#ifdef __GLASGOW_HASKELL__ +{-# NOINLINE [1] mapWithIndex #-} +{-# RULES +"mapWithIndex/mapWithIndex" forall f g xs . mapWithIndex f (mapWithIndex g xs) = + mapWithIndex (\k a -> f k (g k a)) xs +"mapWithIndex/fmapSeq" forall f g xs . mapWithIndex f (fmapSeq g xs) = + mapWithIndex (\k a -> f k (g a)) xs +"fmapSeq/mapWithIndex" forall f g xs . fmapSeq f (mapWithIndex g xs) = + mapWithIndex (\k a -> f (g k a)) xs + #-} +#endif + + +-- | /O(n)/. A generalization of 'foldMap', 'foldMapWithIndex' takes a folding +-- function that also depends on the element's index, and applies it to every +-- element in the sequence. +-- +-- @since 0.5.8 +foldMapWithIndex :: Monoid m => (Int -> a -> m) -> Seq a -> m +foldMapWithIndex f' (Seq xs') = foldMapWithIndexTreeE (lift_elem f') 0 xs' + where + lift_elem :: (Int -> a -> m) -> (Int -> Elem a -> m) +#if __GLASGOW_HASKELL__ >= 708 + lift_elem g = coerce g +#else + lift_elem g = \s (Elem a) -> g s a +#endif + {-# INLINE lift_elem #-} +-- We have to specialize these functions by hand, unfortunately, because +-- GHC does not specialize until *all* instances are determined. +-- Although the Sized instance is known at compile time, the Monoid +-- instance generally is not. + foldMapWithIndexTreeE :: Monoid m => (Int -> Elem a -> m) -> Int -> FingerTree (Elem a) -> m + foldMapWithIndexTreeE _ !_s EmptyT = mempty + foldMapWithIndexTreeE f s (Single xs) = f s xs + foldMapWithIndexTreeE f s (Deep _ pr m sf) = + foldMapWithIndexDigitE f s pr <> + foldMapWithIndexTreeN (foldMapWithIndexNodeE f) sPspr m <> + foldMapWithIndexDigitE f sPsprm sf + where + !sPspr = s + size pr + !sPsprm = sPspr + size m + + foldMapWithIndexTreeN :: Monoid m => (Int -> Node a -> m) -> Int -> FingerTree (Node a) -> m + foldMapWithIndexTreeN _ !_s EmptyT = mempty + foldMapWithIndexTreeN f s (Single xs) = f s xs + foldMapWithIndexTreeN f s (Deep _ pr m sf) = + foldMapWithIndexDigitN f s pr <> + foldMapWithIndexTreeN (foldMapWithIndexNodeN f) sPspr m <> + foldMapWithIndexDigitN f sPsprm sf + where + !sPspr = s + size pr + !sPsprm = sPspr + size m + + foldMapWithIndexDigitE :: Monoid m => (Int -> Elem a -> m) -> Int -> Digit (Elem a) -> m + foldMapWithIndexDigitE f i t = foldMapWithIndexDigit f i t + + foldMapWithIndexDigitN :: Monoid m => (Int -> Node a -> m) -> Int -> Digit (Node a) -> m + foldMapWithIndexDigitN f i t = foldMapWithIndexDigit f i t + + {-# INLINE foldMapWithIndexDigit #-} + foldMapWithIndexDigit :: (Monoid m, Sized a) => (Int -> a -> m) -> Int -> Digit a -> m + foldMapWithIndexDigit f !s (One a) = f s a + foldMapWithIndexDigit f s (Two a b) = f s a <> f sPsa b + where + !sPsa = s + size a + foldMapWithIndexDigit f s (Three a b c) = + f s a <> f sPsa b <> f sPsab c + where + !sPsa = s + size a + !sPsab = sPsa + size b + foldMapWithIndexDigit f s (Four a b c d) = + f s a <> f sPsa b <> f sPsab c <> f sPsabc d + where + !sPsa = s + size a + !sPsab = sPsa + size b + !sPsabc = sPsab + size c + + foldMapWithIndexNodeE :: Monoid m => (Int -> Elem a -> m) -> Int -> Node (Elem a) -> m + foldMapWithIndexNodeE f i t = foldMapWithIndexNode f i t + + foldMapWithIndexNodeN :: Monoid m => (Int -> Node a -> m) -> Int -> Node (Node a) -> m + foldMapWithIndexNodeN f i t = foldMapWithIndexNode f i t + + {-# INLINE foldMapWithIndexNode #-} + foldMapWithIndexNode :: (Monoid m, Sized a) => (Int -> a -> m) -> Int -> Node a -> m + foldMapWithIndexNode f !s (Node2 _ a b) = f s a <> f sPsa b + where + !sPsa = s + size a + foldMapWithIndexNode f s (Node3 _ a b c) = + f s a <> f sPsa b <> f sPsab c + where + !sPsa = s + size a + !sPsab = sPsa + size b + +#if __GLASGOW_HASKELL__ +{-# INLINABLE foldMapWithIndex #-} +#endif + +-- | 'traverseWithIndex' is a version of 'traverse' that also offers +-- access to the index of each element. +-- +-- @since 0.5.8 +traverseWithIndex :: Applicative f => (Int -> a -> f b) -> Seq a -> f (Seq b) +traverseWithIndex f' (Seq xs') = Seq <$> traverseWithIndexTreeE (\s (Elem a) -> Elem <$> f' s a) 0 xs' + where +-- We have to specialize these functions by hand, unfortunately, because +-- GHC does not specialize until *all* instances are determined. +-- Although the Sized instance is known at compile time, the Applicative +-- instance generally is not. + traverseWithIndexTreeE :: Applicative f => (Int -> Elem a -> f b) -> Int -> FingerTree (Elem a) -> f (FingerTree b) + traverseWithIndexTreeE _ !_s EmptyT = pure EmptyT + traverseWithIndexTreeE f s (Single xs) = Single <$> f s xs + traverseWithIndexTreeE f s (Deep n pr m sf) = + deep' n <$> + traverseWithIndexDigitE f s pr <*> + traverseWithIndexTreeN (traverseWithIndexNodeE f) sPspr m <*> + traverseWithIndexDigitE f sPsprm sf + where + !sPspr = s + size pr + !sPsprm = sPspr + size m + + traverseWithIndexTreeN :: Applicative f => (Int -> Node a -> f b) -> Int -> FingerTree (Node a) -> f (FingerTree b) + traverseWithIndexTreeN _ !_s EmptyT = pure EmptyT + traverseWithIndexTreeN f s (Single xs) = Single <$> f s xs + traverseWithIndexTreeN f s (Deep n pr m sf) = + deep' n <$> + traverseWithIndexDigitN f s pr <*> + traverseWithIndexTreeN (traverseWithIndexNodeN f) sPspr m <*> + traverseWithIndexDigitN f sPsprm sf + where + !sPspr = s + size pr + !sPsprm = sPspr + size m + + traverseWithIndexDigitE :: Applicative f => (Int -> Elem a -> f b) -> Int -> Digit (Elem a) -> f (Digit b) + traverseWithIndexDigitE f i t = traverseWithIndexDigit f i t + + traverseWithIndexDigitN :: Applicative f => (Int -> Node a -> f b) -> Int -> Digit (Node a) -> f (Digit b) + traverseWithIndexDigitN f i t = traverseWithIndexDigit f i t + + {-# INLINE traverseWithIndexDigit #-} + traverseWithIndexDigit :: (Applicative f, Sized a) => (Int -> a -> f b) -> Int -> Digit a -> f (Digit b) + traverseWithIndexDigit f !s (One a) = One <$> f s a + traverseWithIndexDigit f s (Two a b) = Two <$> f s a <*> f sPsa b + where + !sPsa = s + size a + traverseWithIndexDigit f s (Three a b c) = + Three <$> f s a <*> f sPsa b <*> f sPsab c + where + !sPsa = s + size a + !sPsab = sPsa + size b + traverseWithIndexDigit f s (Four a b c d) = + Four <$> f s a <*> f sPsa b <*> f sPsab c <*> f sPsabc d + where + !sPsa = s + size a + !sPsab = sPsa + size b + !sPsabc = sPsab + size c + + traverseWithIndexNodeE :: Applicative f => (Int -> Elem a -> f b) -> Int -> Node (Elem a) -> f (Node b) + traverseWithIndexNodeE f i t = traverseWithIndexNode f i t + + traverseWithIndexNodeN :: Applicative f => (Int -> Node a -> f b) -> Int -> Node (Node a) -> f (Node b) + traverseWithIndexNodeN f i t = traverseWithIndexNode f i t + + {-# INLINE traverseWithIndexNode #-} + traverseWithIndexNode :: (Applicative f, Sized a) => (Int -> a -> f b) -> Int -> Node a -> f (Node b) + traverseWithIndexNode f !s (Node2 ns a b) = node2' ns <$> f s a <*> f sPsa b + where + !sPsa = s + size a + traverseWithIndexNode f s (Node3 ns a b c) = + node3' ns <$> f s a <*> f sPsa b <*> f sPsab c + where + !sPsa = s + size a + !sPsab = sPsa + size b + + +{-# NOINLINE [1] traverseWithIndex #-} +#ifdef __GLASGOW_HASKELL__ +{-# RULES +"travWithIndex/mapWithIndex" forall f g xs . traverseWithIndex f (mapWithIndex g xs) = + traverseWithIndex (\k a -> f k (g k a)) xs +"travWithIndex/fmapSeq" forall f g xs . traverseWithIndex f (fmapSeq g xs) = + traverseWithIndex (\k a -> f k (g a)) xs + #-} +#endif +{- +It might be nice to be able to rewrite + +traverseWithIndex f (fromFunction i g) +to +replicateAWithIndex i (\k -> f k (g k)) +and +traverse f (fromFunction i g) +to +replicateAWithIndex i (f . g) + +but we don't have replicateAWithIndex as yet. + +We might wish for a rule like +"fmapSeq/travWithIndex" forall f g xs . fmapSeq f <$> traverseWithIndex g xs = + traverseWithIndex (\k a -> f <$> g k a) xs +Unfortunately, this rule could screw up the inliner's treatment of +fmap in general, and it also relies on the arbitrary Functor being +valid. +-} + + +-- | /O(n)/. Convert a given sequence length and a function representing that +-- sequence into a sequence. +fromFunction :: Int -> (Int -> a) -> Seq a +fromFunction len f | len < 0 = error "Data.Sequence.fromFunction called with negative len" + | len == 0 = empty + | otherwise = Seq $ create (lift_elem f) 1 0 len + where + create :: (Int -> a) -> Int -> Int -> Int -> FingerTree a + create b{-tree_builder-} !s{-tree_size-} !i{-start_index-} trees = case trees of + 1 -> Single $ b i + 2 -> Deep (2*s) (One (b i)) EmptyT (One (b (i+s))) + 3 -> Deep (3*s) (createTwo i) EmptyT (One (b (i+2*s))) + 4 -> Deep (4*s) (createTwo i) EmptyT (createTwo (i+2*s)) + 5 -> Deep (5*s) (createThree i) EmptyT (createTwo (i+3*s)) + 6 -> Deep (6*s) (createThree i) EmptyT (createThree (i+3*s)) + _ -> case trees `quotRem` 3 of + (trees', 1) -> Deep (trees*s) (createTwo i) + (create mb (3*s) (i+2*s) (trees'-1)) + (createTwo (i+(2+3*(trees'-1))*s)) + (trees', 2) -> Deep (trees*s) (createThree i) + (create mb (3*s) (i+3*s) (trees'-1)) + (createTwo (i+(3+3*(trees'-1))*s)) + (trees', _) -> Deep (trees*s) (createThree i) + (create mb (3*s) (i+3*s) (trees'-2)) + (createThree (i+(3+3*(trees'-2))*s)) + where + createTwo j = Two (b j) (b (j + s)) + {-# INLINE createTwo #-} + createThree j = Three (b j) (b (j + s)) (b (j + 2*s)) + {-# INLINE createThree #-} + mb j = Node3 (3*s) (b j) (b (j + s)) (b (j + 2*s)) + {-# INLINE mb #-} + + lift_elem :: (Int -> a) -> (Int -> Elem a) +#if __GLASGOW_HASKELL__ >= 708 + lift_elem g = coerce g +#else + lift_elem g = Elem . g +#endif + {-# INLINE lift_elem #-} + +-- | /O(n)/. Create a sequence consisting of the elements of an 'Array'. +-- Note that the resulting sequence elements may be evaluated lazily (as on GHC), +-- so you must force the entire structure to be sure that the original array +-- can be garbage-collected. +fromArray :: Ix i => Array i a -> Seq a +#ifdef __GLASGOW_HASKELL__ +fromArray a = fromFunction (GHC.Arr.numElements a) (GHC.Arr.unsafeAt a) + where + -- The following definition uses (Ix i) constraing, which is needed for the + -- other fromArray definition. + _ = Data.Array.rangeSize (Data.Array.bounds a) +#else +fromArray a = fromList2 (Data.Array.rangeSize (Data.Array.bounds a)) (Data.Array.elems a) +#endif + +-- Splitting + +-- | /O(log(min(i,n-i)))/. The first @i@ elements of a sequence. +-- If @i@ is negative, @'take' i s@ yields the empty sequence. +-- If the sequence contains fewer than @i@ elements, the whole sequence +-- is returned. +take :: Int -> Seq a -> Seq a +take i xs@(Seq t) + -- See note on unsigned arithmetic in splitAt + | fromIntegral i - 1 < (fromIntegral (length xs) - 1 :: Word) = + Seq (takeTreeE i t) + | i <= 0 = empty + | otherwise = xs + +takeTreeE :: Int -> FingerTree (Elem a) -> FingerTree (Elem a) +takeTreeE !_i EmptyT = EmptyT +takeTreeE i t@(Single _) + | i <= 0 = EmptyT + | otherwise = t +takeTreeE i (Deep s pr m sf) + | i < spr = takePrefixE i pr + | i < spm = case takeTreeN im m of + ml :*: xs -> takeMiddleE (im - size ml) spr pr ml xs + | otherwise = takeSuffixE (i - spm) s pr m sf + where + spr = size pr + spm = spr + size m + im = i - spr + +takeTreeN :: Int -> FingerTree (Node a) -> StrictPair (FingerTree (Node a)) (Node a) +takeTreeN !_i EmptyT = error "takeTreeN of empty tree" +takeTreeN _i (Single x) = EmptyT :*: x +takeTreeN i (Deep s pr m sf) + | i < spr = takePrefixN i pr + | i < spm = case takeTreeN im m of + ml :*: xs -> takeMiddleN (im - size ml) spr pr ml xs + | otherwise = takeSuffixN (i - spm) s pr m sf where + spr = size pr + spm = spr + size m + im = i - spr + +takeMiddleN :: Int -> Int + -> Digit (Node a) -> FingerTree (Node (Node a)) -> Node (Node a) + -> StrictPair (FingerTree (Node a)) (Node a) +takeMiddleN i spr pr ml (Node2 _ a b) + | i < sa = pullR sprml pr ml :*: a + | otherwise = Deep sprmla pr ml (One a) :*: b + where + sa = size a + sprml = spr + size ml + sprmla = sa + sprml +takeMiddleN i spr pr ml (Node3 _ a b c) + | i < sa = pullR sprml pr ml :*: a + | i < sab = Deep sprmla pr ml (One a) :*: b + | otherwise = Deep sprmlab pr ml (Two a b) :*: c + where + sa = size a + sab = sa + size b + sprml = spr + size ml + sprmla = sa + sprml + sprmlab = sprmla + size b + +takeMiddleE :: Int -> Int + -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> Node (Elem a) + -> FingerTree (Elem a) +takeMiddleE i spr pr ml (Node2 _ a _) + | i < 1 = pullR sprml pr ml + | otherwise = Deep sprmla pr ml (One a) + where + sprml = spr + size ml + sprmla = 1 + sprml +takeMiddleE i spr pr ml (Node3 _ a b _) + | i < 1 = pullR sprml pr ml + | i < 2 = Deep sprmla pr ml (One a) + | otherwise = Deep sprmlab pr ml (Two a b) + where + sprml = spr + size ml + sprmla = 1 + sprml + sprmlab = sprmla + 1 + +takePrefixE :: Int -> Digit (Elem a) -> FingerTree (Elem a) +takePrefixE !_i (One _) = EmptyT +takePrefixE i (Two a _) + | i < 1 = EmptyT + | otherwise = Single a +takePrefixE i (Three a b _) + | i < 1 = EmptyT + | i < 2 = Single a + | otherwise = Deep 2 (One a) EmptyT (One b) +takePrefixE i (Four a b c _) + | i < 1 = EmptyT + | i < 2 = Single a + | i < 3 = Deep 2 (One a) EmptyT (One b) + | otherwise = Deep 3 (Two a b) EmptyT (One c) + +takePrefixN :: Int -> Digit (Node a) + -> StrictPair (FingerTree (Node a)) (Node a) +takePrefixN !_i (One a) = EmptyT :*: a +takePrefixN i (Two a b) + | i < sa = EmptyT :*: a + | otherwise = Single a :*: b + where + sa = size a +takePrefixN i (Three a b c) + | i < sa = EmptyT :*: a + | i < sab = Single a :*: b + | otherwise = Deep sab (One a) EmptyT (One b) :*: c + where + sa = size a + sab = sa + size b +takePrefixN i (Four a b c d) + | i < sa = EmptyT :*: a + | i < sab = Single a :*: b + | i < sabc = Deep sab (One a) EmptyT (One b) :*: c + | otherwise = Deep sabc (Two a b) EmptyT (One c) :*: d + where + sa = size a + sab = sa + size b + sabc = sab + size c + +takeSuffixE :: Int -> Int -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> + FingerTree (Elem a) +takeSuffixE !_i !s pr m (One _) = pullR (s - 1) pr m +takeSuffixE i s pr m (Two a _) + | i < 1 = pullR (s - 2) pr m + | otherwise = Deep (s - 1) pr m (One a) +takeSuffixE i s pr m (Three a b _) + | i < 1 = pullR (s - 3) pr m + | i < 2 = Deep (s - 2) pr m (One a) + | otherwise = Deep (s - 1) pr m (Two a b) +takeSuffixE i s pr m (Four a b c _) + | i < 1 = pullR (s - 4) pr m + | i < 2 = Deep (s - 3) pr m (One a) + | i < 3 = Deep (s - 2) pr m (Two a b) + | otherwise = Deep (s - 1) pr m (Three a b c) + +takeSuffixN :: Int -> Int -> Digit (Node a) -> FingerTree (Node (Node a)) -> Digit (Node a) -> + StrictPair (FingerTree (Node a)) (Node a) +takeSuffixN !_i !s pr m (One a) = pullR (s - size a) pr m :*: a +takeSuffixN i s pr m (Two a b) + | i < sa = pullR (s - sa - size b) pr m :*: a + | otherwise = Deep (s - size b) pr m (One a) :*: b + where + sa = size a +takeSuffixN i s pr m (Three a b c) + | i < sa = pullR (s - sab - size c) pr m :*: a + | i < sab = Deep (s - size b - size c) pr m (One a) :*: b + | otherwise = Deep (s - size c) pr m (Two a b) :*: c + where + sa = size a + sab = sa + size b +takeSuffixN i s pr m (Four a b c d) + | i < sa = pullR (s - sa - sbcd) pr m :*: a + | i < sab = Deep (s - sbcd) pr m (One a) :*: b + | i < sabc = Deep (s - scd) pr m (Two a b) :*: c + | otherwise = Deep (s - sd) pr m (Three a b c) :*: d + where + sa = size a + sab = sa + size b + sabc = sab + size c + sd = size d + scd = size c + sd + sbcd = size b + scd + +-- | /O(log(min(i,n-i)))/. Elements of a sequence after the first @i@. +-- If @i@ is negative, @'drop' i s@ yields the whole sequence. +-- If the sequence contains fewer than @i@ elements, the empty sequence +-- is returned. +drop :: Int -> Seq a -> Seq a +drop i xs@(Seq t) + -- See note on unsigned arithmetic in splitAt + | fromIntegral i - 1 < (fromIntegral (length xs) - 1 :: Word) = + Seq (takeTreeER (length xs - i) t) + | i <= 0 = xs + | otherwise = empty + +-- We implement `drop` using a "take from the rear" strategy. There's no +-- particular technical reason for this; it just lets us reuse the arithmetic +-- from `take` (which itself reuses the arithmetic from `splitAt`) instead of +-- figuring it out from scratch and ending up with lots of off-by-one errors. +takeTreeER :: Int -> FingerTree (Elem a) -> FingerTree (Elem a) +takeTreeER !_i EmptyT = EmptyT +takeTreeER i t@(Single _) + | i <= 0 = EmptyT + | otherwise = t +takeTreeER i (Deep s pr m sf) + | i < ssf = takeSuffixER i sf + | i < ssm = case takeTreeNR im m of + xs :*: mr -> takeMiddleER (im - size mr) ssf xs mr sf + | otherwise = takePrefixER (i - ssm) s pr m sf + where + ssf = size sf + ssm = ssf + size m + im = i - ssf + +takeTreeNR :: Int -> FingerTree (Node a) -> StrictPair (Node a) (FingerTree (Node a)) +takeTreeNR !_i EmptyT = error "takeTreeNR of empty tree" +takeTreeNR _i (Single x) = x :*: EmptyT +takeTreeNR i (Deep s pr m sf) + | i < ssf = takeSuffixNR i sf + | i < ssm = case takeTreeNR im m of + xs :*: mr -> takeMiddleNR (im - size mr) ssf xs mr sf + | otherwise = takePrefixNR (i - ssm) s pr m sf where + ssf = size sf + ssm = ssf + size m + im = i - ssf + +takeMiddleNR :: Int -> Int + -> Node (Node a) -> FingerTree (Node (Node a)) -> Digit (Node a) + -> StrictPair (Node a) (FingerTree (Node a)) +takeMiddleNR i ssf (Node2 _ a b) mr sf + | i < sb = b :*: pullL ssfmr mr sf + | otherwise = a :*: Deep ssfmrb (One b) mr sf + where + sb = size b + ssfmr = ssf + size mr + ssfmrb = sb + ssfmr +takeMiddleNR i ssf (Node3 _ a b c) mr sf + | i < sc = c :*: pullL ssfmr mr sf + | i < sbc = b :*: Deep ssfmrc (One c) mr sf + | otherwise = a :*: Deep ssfmrbc (Two b c) mr sf + where + sc = size c + sbc = sc + size b + ssfmr = ssf + size mr + ssfmrc = sc + ssfmr + ssfmrbc = ssfmrc + size b + +takeMiddleER :: Int -> Int + -> Node (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) + -> FingerTree (Elem a) +takeMiddleER i ssf (Node2 _ _ b) mr sf + | i < 1 = pullL ssfmr mr sf + | otherwise = Deep ssfmrb (One b) mr sf + where + ssfmr = ssf + size mr + ssfmrb = 1 + ssfmr +takeMiddleER i ssf (Node3 _ _ b c) mr sf + | i < 1 = pullL ssfmr mr sf + | i < 2 = Deep ssfmrc (One c) mr sf + | otherwise = Deep ssfmrbc (Two b c) mr sf + where + ssfmr = ssf + size mr + ssfmrc = 1 + ssfmr + ssfmrbc = ssfmr + 2 + +takeSuffixER :: Int -> Digit (Elem a) -> FingerTree (Elem a) +takeSuffixER !_i (One _) = EmptyT +takeSuffixER i (Two _ b) + | i < 1 = EmptyT + | otherwise = Single b +takeSuffixER i (Three _ b c) + | i < 1 = EmptyT + | i < 2 = Single c + | otherwise = Deep 2 (One b) EmptyT (One c) +takeSuffixER i (Four _ b c d) + | i < 1 = EmptyT + | i < 2 = Single d + | i < 3 = Deep 2 (One c) EmptyT (One d) + | otherwise = Deep 3 (Two b c) EmptyT (One d) + +takeSuffixNR :: Int -> Digit (Node a) + -> StrictPair (Node a) (FingerTree (Node a)) +takeSuffixNR !_i (One a) = a :*: EmptyT +takeSuffixNR i (Two a b) + | i < sb = b :*: EmptyT + | otherwise = a :*: Single b + where + sb = size b +takeSuffixNR i (Three a b c) + | i < sc = c :*: EmptyT + | i < sbc = b :*: Single c + | otherwise = a :*: Deep sbc (One b) EmptyT (One c) + where + sc = size c + sbc = sc + size b +takeSuffixNR i (Four a b c d) + | i < sd = d :*: EmptyT + | i < scd = c :*: Single d + | i < sbcd = b :*: Deep scd (One c) EmptyT (One d) + | otherwise = a :*: Deep sbcd (Two b c) EmptyT (One d) + where + sd = size d + scd = sd + size c + sbcd = scd + size b + +takePrefixER :: Int -> Int -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> + FingerTree (Elem a) +takePrefixER !_i !s (One _) m sf = pullL (s - 1) m sf +takePrefixER i s (Two _ b) m sf + | i < 1 = pullL (s - 2) m sf + | otherwise = Deep (s - 1) (One b) m sf +takePrefixER i s (Three _ b c) m sf + | i < 1 = pullL (s - 3) m sf + | i < 2 = Deep (s - 2) (One c) m sf + | otherwise = Deep (s - 1) (Two b c) m sf +takePrefixER i s (Four _ b c d) m sf + | i < 1 = pullL (s - 4) m sf + | i < 2 = Deep (s - 3) (One d) m sf + | i < 3 = Deep (s - 2) (Two c d) m sf + | otherwise = Deep (s - 1) (Three b c d) m sf + +takePrefixNR :: Int -> Int -> Digit (Node a) -> FingerTree (Node (Node a)) -> Digit (Node a) -> + StrictPair (Node a) (FingerTree (Node a)) +takePrefixNR !_i !s (One a) m sf = a :*: pullL (s - size a) m sf +takePrefixNR i s (Two a b) m sf + | i < sb = b :*: pullL (s - sb - size a) m sf + | otherwise = a :*: Deep (s - size a) (One b) m sf + where + sb = size b +takePrefixNR i s (Three a b c) m sf + | i < sc = c :*: pullL (s - sbc - size a) m sf + | i < sbc = b :*: Deep (s - size b - size a) (One c) m sf + | otherwise = a :*: Deep (s - size a) (Two b c) m sf + where + sc = size c + sbc = sc + size b +takePrefixNR i s (Four a b c d) m sf + | i < sd = d :*: pullL (s - sd - sabc) m sf + | i < scd = c :*: Deep (s - sabc) (One d) m sf + | i < sbcd = b :*: Deep (s - sab) (Two c d) m sf + | otherwise = a :*: Deep (s - sa) (Three b c d) m sf + where + sa = size a + sab = sa + size b + sabc = sab + size c + sd = size d + scd = size c + sd + sbcd = size b + scd + +-- | /O(log(min(i,n-i)))/. Split a sequence at a given position. +-- @'splitAt' i s = ('take' i s, 'drop' i s)@. +splitAt :: Int -> Seq a -> (Seq a, Seq a) +splitAt i xs@(Seq t) + -- We use an unsigned comparison to make the common case + -- faster. This only works because our representation of + -- sizes as (signed) Ints gives us a free high bit to play + -- with. Note also that there's no sharing to lose in the + -- case that the length is 0. + | fromIntegral i - 1 < (fromIntegral (length xs) - 1 :: Word) = + case splitTreeE i t of + l :*: r -> (Seq l, Seq r) + | i <= 0 = (empty, xs) + | otherwise = (xs, empty) + +-- | /O(log(min(i,n-i))) A version of 'splitAt' that does not attempt to +-- enhance sharing when the split point is less than or equal to 0, and that +-- gives completely wrong answers when the split point is at least the length +-- of the sequence, unless the sequence is a singleton. This is used to +-- implement zipWith and chunksOf, which are extremely sensitive to the cost of +-- splitting very short sequences. There is just enough of a speed increase to +-- make this worth the trouble. +uncheckedSplitAt :: Int -> Seq a -> (Seq a, Seq a) +uncheckedSplitAt i (Seq xs) = case splitTreeE i xs of + l :*: r -> (Seq l, Seq r) + +data Split a = Split !(FingerTree (Node a)) !(Node a) !(FingerTree (Node a)) +#if TESTING + deriving Show +#endif + +splitTreeE :: Int -> FingerTree (Elem a) -> StrictPair (FingerTree (Elem a)) (FingerTree (Elem a)) +splitTreeE !_i EmptyT = EmptyT :*: EmptyT +splitTreeE i t@(Single _) + | i <= 0 = EmptyT :*: t + | otherwise = t :*: EmptyT +splitTreeE i (Deep s pr m sf) + | i < spr = splitPrefixE i s pr m sf + | i < spm = case splitTreeN im m of + Split ml xs mr -> splitMiddleE (im - size ml) s spr pr ml xs mr sf + | otherwise = splitSuffixE (i - spm) s pr m sf + where + spr = size pr + spm = spr + size m + im = i - spr + +splitTreeN :: Int -> FingerTree (Node a) -> Split a +splitTreeN !_i EmptyT = error "splitTreeN of empty tree" +splitTreeN _i (Single x) = Split EmptyT x EmptyT +splitTreeN i (Deep s pr m sf) + | i < spr = splitPrefixN i s pr m sf + | i < spm = case splitTreeN im m of + Split ml xs mr -> splitMiddleN (im - size ml) s spr pr ml xs mr sf + | otherwise = splitSuffixN (i - spm) s pr m sf where + spr = size pr + spm = spr + size m + im = i - spr + +splitMiddleN :: Int -> Int -> Int + -> Digit (Node a) -> FingerTree (Node (Node a)) -> Node (Node a) -> FingerTree (Node (Node a)) -> Digit (Node a) + -> Split a +splitMiddleN i s spr pr ml (Node2 _ a b) mr sf + | i < sa = Split (pullR sprml pr ml) a (Deep (s - sprmla) (One b) mr sf) + | otherwise = Split (Deep sprmla pr ml (One a)) b (pullL (s - sprmla - size b) mr sf) + where + sa = size a + sprml = spr + size ml + sprmla = sa + sprml +splitMiddleN i s spr pr ml (Node3 _ a b c) mr sf + | i < sa = Split (pullR sprml pr ml) a (Deep (s - sprmla) (Two b c) mr sf) + | i < sab = Split (Deep sprmla pr ml (One a)) b (Deep (s - sprmlab) (One c) mr sf) + | otherwise = Split (Deep sprmlab pr ml (Two a b)) c (pullL (s - sprmlab - size c) mr sf) + where + sa = size a + sab = sa + size b + sprml = spr + size ml + sprmla = sa + sprml + sprmlab = sprmla + size b + +splitMiddleE :: Int -> Int -> Int + -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> Node (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) + -> StrictPair (FingerTree (Elem a)) (FingerTree (Elem a)) +splitMiddleE i s spr pr ml (Node2 _ a b) mr sf + | i < 1 = pullR sprml pr ml :*: Deep (s - sprml) (Two a b) mr sf + | otherwise = Deep sprmla pr ml (One a) :*: Deep (s - sprmla) (One b) mr sf + where + sprml = spr + size ml + sprmla = 1 + sprml +splitMiddleE i s spr pr ml (Node3 _ a b c) mr sf = case i of + 0 -> pullR sprml pr ml :*: Deep (s - sprml) (Three a b c) mr sf + 1 -> Deep sprmla pr ml (One a) :*: Deep (s - sprmla) (Two b c) mr sf + _ -> Deep sprmlab pr ml (Two a b) :*: Deep (s - sprmlab) (One c) mr sf + where + sprml = spr + size ml + sprmla = 1 + sprml + sprmlab = sprmla + 1 + +splitPrefixE :: Int -> Int -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> + StrictPair (FingerTree (Elem a)) (FingerTree (Elem a)) +splitPrefixE !_i !s (One a) m sf = EmptyT :*: Deep s (One a) m sf +splitPrefixE i s (Two a b) m sf = case i of + 0 -> EmptyT :*: Deep s (Two a b) m sf + _ -> Single a :*: Deep (s - 1) (One b) m sf +splitPrefixE i s (Three a b c) m sf = case i of + 0 -> EmptyT :*: Deep s (Three a b c) m sf + 1 -> Single a :*: Deep (s - 1) (Two b c) m sf + _ -> Deep 2 (One a) EmptyT (One b) :*: Deep (s - 2) (One c) m sf +splitPrefixE i s (Four a b c d) m sf = case i of + 0 -> EmptyT :*: Deep s (Four a b c d) m sf + 1 -> Single a :*: Deep (s - 1) (Three b c d) m sf + 2 -> Deep 2 (One a) EmptyT (One b) :*: Deep (s - 2) (Two c d) m sf + _ -> Deep 3 (Two a b) EmptyT (One c) :*: Deep (s - 3) (One d) m sf + +splitPrefixN :: Int -> Int -> Digit (Node a) -> FingerTree (Node (Node a)) -> Digit (Node a) -> + Split a +splitPrefixN !_i !s (One a) m sf = Split EmptyT a (pullL (s - size a) m sf) +splitPrefixN i s (Two a b) m sf + | i < sa = Split EmptyT a (Deep (s - sa) (One b) m sf) + | otherwise = Split (Single a) b (pullL (s - sa - size b) m sf) + where + sa = size a +splitPrefixN i s (Three a b c) m sf + | i < sa = Split EmptyT a (Deep (s - sa) (Two b c) m sf) + | i < sab = Split (Single a) b (Deep (s - sab) (One c) m sf) + | otherwise = Split (Deep sab (One a) EmptyT (One b)) c (pullL (s - sab - size c) m sf) + where + sa = size a + sab = sa + size b +splitPrefixN i s (Four a b c d) m sf + | i < sa = Split EmptyT a $ Deep (s - sa) (Three b c d) m sf + | i < sab = Split (Single a) b $ Deep (s - sab) (Two c d) m sf + | i < sabc = Split (Deep sab (One a) EmptyT (One b)) c $ Deep (s - sabc) (One d) m sf + | otherwise = Split (Deep sabc (Two a b) EmptyT (One c)) d $ pullL (s - sabc - size d) m sf + where + sa = size a + sab = sa + size b + sabc = sab + size c + +splitSuffixE :: Int -> Int -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> + StrictPair (FingerTree (Elem a)) (FingerTree (Elem a)) +splitSuffixE !_i !s pr m (One a) = pullR (s - 1) pr m :*: Single a +splitSuffixE i s pr m (Two a b) = case i of + 0 -> pullR (s - 2) pr m :*: Deep 2 (One a) EmptyT (One b) + _ -> Deep (s - 1) pr m (One a) :*: Single b +splitSuffixE i s pr m (Three a b c) = case i of + 0 -> pullR (s - 3) pr m :*: Deep 3 (Two a b) EmptyT (One c) + 1 -> Deep (s - 2) pr m (One a) :*: Deep 2 (One b) EmptyT (One c) + _ -> Deep (s - 1) pr m (Two a b) :*: Single c +splitSuffixE i s pr m (Four a b c d) = case i of + 0 -> pullR (s - 4) pr m :*: Deep 4 (Two a b) EmptyT (Two c d) + 1 -> Deep (s - 3) pr m (One a) :*: Deep 3 (Two b c) EmptyT (One d) + 2 -> Deep (s - 2) pr m (Two a b) :*: Deep 2 (One c) EmptyT (One d) + _ -> Deep (s - 1) pr m (Three a b c) :*: Single d + +splitSuffixN :: Int -> Int -> Digit (Node a) -> FingerTree (Node (Node a)) -> Digit (Node a) -> + Split a +splitSuffixN !_i !s pr m (One a) = Split (pullR (s - size a) pr m) a EmptyT +splitSuffixN i s pr m (Two a b) + | i < sa = Split (pullR (s - sa - size b) pr m) a (Single b) + | otherwise = Split (Deep (s - size b) pr m (One a)) b EmptyT + where + sa = size a +splitSuffixN i s pr m (Three a b c) + | i < sa = Split (pullR (s - sab - size c) pr m) a (deep (One b) EmptyT (One c)) + | i < sab = Split (Deep (s - size b - size c) pr m (One a)) b (Single c) + | otherwise = Split (Deep (s - size c) pr m (Two a b)) c EmptyT + where + sa = size a + sab = sa + size b +splitSuffixN i s pr m (Four a b c d) + | i < sa = Split (pullR (s - sa - sbcd) pr m) a (Deep sbcd (Two b c) EmptyT (One d)) + | i < sab = Split (Deep (s - sbcd) pr m (One a)) b (Deep scd (One c) EmptyT (One d)) + | i < sabc = Split (Deep (s - scd) pr m (Two a b)) c (Single d) + | otherwise = Split (Deep (s - sd) pr m (Three a b c)) d EmptyT + where + sa = size a + sab = sa + size b + sabc = sab + size c + sd = size d + scd = size c + sd + sbcd = size b + scd + +-- | /O(n)/. @chunksOf n xs@ splits @xs@ into chunks of size @n>0@. +-- If @n@ does not divide the length of @xs@ evenly, then the last element +-- of the result will be short. +chunksOf :: Int -> Seq a -> Seq (Seq a) +chunksOf n xs | n <= 0 = + if null xs + then empty + else error "chunksOf: A non-empty sequence can only be broken up into positively-sized chunks." +chunksOf 1 s = fmap singleton s +chunksOf n s = splitMap (uncheckedSplitAt . (*n)) const most (replicate numReps ()) + >< if null end then empty else singleton end + where + (numReps, endLength) = length s `quotRem` n + (most, end) = splitAt (length s - endLength) s + +-- | /O(n)/. Returns a sequence of all suffixes of this sequence, +-- longest first. For example, +-- +-- > tails (fromList "abc") = fromList [fromList "abc", fromList "bc", fromList "c", fromList ""] +-- +-- Evaluating the /i/th suffix takes /O(log(min(i, n-i)))/, but evaluating +-- every suffix in the sequence takes /O(n)/ due to sharing. +tails :: Seq a -> Seq (Seq a) +tails (Seq xs) = Seq (tailsTree (Elem . Seq) xs) |> empty + +-- | /O(n)/. Returns a sequence of all prefixes of this sequence, +-- shortest first. For example, +-- +-- > inits (fromList "abc") = fromList [fromList "", fromList "a", fromList "ab", fromList "abc"] +-- +-- Evaluating the /i/th prefix takes /O(log(min(i, n-i)))/, but evaluating +-- every prefix in the sequence takes /O(n)/ due to sharing. +inits :: Seq a -> Seq (Seq a) +inits (Seq xs) = empty <| Seq (initsTree (Elem . Seq) xs) + +-- This implementation of tails (and, analogously, inits) has the +-- following algorithmic advantages: +-- Evaluating each tail in the sequence takes linear total time, +-- which is better than we could say for +-- @fromList [drop n xs | n <- [0..length xs]]@. +-- Evaluating any individual tail takes logarithmic time, which is +-- better than we can say for either +-- @scanr (<|) empty xs@ or @iterateN (length xs + 1) (\ xs -> let _ :< xs' = viewl xs in xs') xs@. +-- +-- Moreover, if we actually look at every tail in the sequence, the +-- following benchmarks demonstrate that this implementation is modestly +-- faster than any of the above: +-- +-- Times (ms) +-- min mean +/-sd median max +-- Seq.tails: 21.986 24.961 10.169 22.417 86.485 +-- scanr: 85.392 87.942 2.488 87.425 100.217 +-- iterateN: 29.952 31.245 1.574 30.412 37.268 +-- +-- The algorithm for tails (and, analogously, inits) is as follows: +-- +-- A Node in the FingerTree of tails is constructed by evaluating the +-- corresponding tail of the FingerTree of Nodes, considering the first +-- Node in this tail, and constructing a Node in which each tail of this +-- Node is made to be the prefix of the remaining tree. This ends up +-- working quite elegantly, as the remainder of the tail of the FingerTree +-- of Nodes becomes the middle of a new tail, the suffix of the Node is +-- the prefix, and the suffix of the original tree is retained. +-- +-- In particular, evaluating the /i/th tail involves making as +-- many partial evaluations as the Node depth of the /i/th element. +-- In addition, when we evaluate the /i/th tail, and we also evaluate +-- the /j/th tail, and /m/ Nodes are on the path to both /i/ and /j/, +-- each of those /m/ evaluations are shared between the computation of +-- the /i/th and /j/th tails. +-- +-- wasserman.louis@gmail.com, 7/16/09 + +tailsDigit :: Digit a -> Digit (Digit a) +tailsDigit (One a) = One (One a) +tailsDigit (Two a b) = Two (Two a b) (One b) +tailsDigit (Three a b c) = Three (Three a b c) (Two b c) (One c) +tailsDigit (Four a b c d) = Four (Four a b c d) (Three b c d) (Two c d) (One d) + +initsDigit :: Digit a -> Digit (Digit a) +initsDigit (One a) = One (One a) +initsDigit (Two a b) = Two (One a) (Two a b) +initsDigit (Three a b c) = Three (One a) (Two a b) (Three a b c) +initsDigit (Four a b c d) = Four (One a) (Two a b) (Three a b c) (Four a b c d) + +tailsNode :: Node a -> Node (Digit a) +tailsNode (Node2 s a b) = Node2 s (Two a b) (One b) +tailsNode (Node3 s a b c) = Node3 s (Three a b c) (Two b c) (One c) + +initsNode :: Node a -> Node (Digit a) +initsNode (Node2 s a b) = Node2 s (One a) (Two a b) +initsNode (Node3 s a b c) = Node3 s (One a) (Two a b) (Three a b c) + +{-# SPECIALIZE tailsTree :: (FingerTree (Elem a) -> Elem b) -> FingerTree (Elem a) -> FingerTree (Elem b) #-} +{-# SPECIALIZE tailsTree :: (FingerTree (Node a) -> Node b) -> FingerTree (Node a) -> FingerTree (Node b) #-} +-- | Given a function to apply to tails of a tree, applies that function +-- to every tail of the specified tree. +tailsTree :: Sized a => (FingerTree a -> b) -> FingerTree a -> FingerTree b +tailsTree _ EmptyT = EmptyT +tailsTree f (Single x) = Single (f (Single x)) +tailsTree f (Deep n pr m sf) = + Deep n (fmap (\ pr' -> f (deep pr' m sf)) (tailsDigit pr)) + (tailsTree f' m) + (fmap (f . digitToTree) (tailsDigit sf)) + where + f' ms = let ConsLTree node m' = viewLTree ms in + fmap (\ pr' -> f (deep pr' m' sf)) (tailsNode node) + +{-# SPECIALIZE initsTree :: (FingerTree (Elem a) -> Elem b) -> FingerTree (Elem a) -> FingerTree (Elem b) #-} +{-# SPECIALIZE initsTree :: (FingerTree (Node a) -> Node b) -> FingerTree (Node a) -> FingerTree (Node b) #-} +-- | Given a function to apply to inits of a tree, applies that function +-- to every init of the specified tree. +initsTree :: Sized a => (FingerTree a -> b) -> FingerTree a -> FingerTree b +initsTree _ EmptyT = EmptyT +initsTree f (Single x) = Single (f (Single x)) +initsTree f (Deep n pr m sf) = + Deep n (fmap (f . digitToTree) (initsDigit pr)) + (initsTree f' m) + (fmap (f . deep pr m) (initsDigit sf)) + where + f' ms = let SnocRTree m' node = viewRTree ms in + fmap (\ sf' -> f (deep pr m' sf')) (initsNode node) + +{-# INLINE foldlWithIndex #-} +-- | 'foldlWithIndex' is a version of 'foldl' that also provides access +-- to the index of each element. +foldlWithIndex :: (b -> Int -> a -> b) -> b -> Seq a -> b +foldlWithIndex f z xs = foldl (\ g x !i -> f (g (i - 1)) i x) (const z) xs (length xs - 1) + +{-# INLINE foldrWithIndex #-} +-- | 'foldrWithIndex' is a version of 'foldr' that also provides access +-- to the index of each element. +foldrWithIndex :: (Int -> a -> b -> b) -> b -> Seq a -> b +foldrWithIndex f z xs = foldr (\ x g !i -> f i x (g (i+1))) (const z) xs 0 + +{-# INLINE listToMaybe' #-} +-- 'listToMaybe\'' is a good consumer version of 'listToMaybe'. +listToMaybe' :: [a] -> Maybe a +listToMaybe' = foldr (\ x _ -> Just x) Nothing + +-- | /O(i)/ where /i/ is the prefix length. 'takeWhileL', applied +-- to a predicate @p@ and a sequence @xs@, returns the longest prefix +-- (possibly empty) of @xs@ of elements that satisfy @p@. +takeWhileL :: (a -> Bool) -> Seq a -> Seq a +takeWhileL p = fst . spanl p + +-- | /O(i)/ where /i/ is the suffix length. 'takeWhileR', applied +-- to a predicate @p@ and a sequence @xs@, returns the longest suffix +-- (possibly empty) of @xs@ of elements that satisfy @p@. +-- +-- @'takeWhileR' p xs@ is equivalent to @'reverse' ('takeWhileL' p ('reverse' xs))@. +takeWhileR :: (a -> Bool) -> Seq a -> Seq a +takeWhileR p = fst . spanr p + +-- | /O(i)/ where /i/ is the prefix length. @'dropWhileL' p xs@ returns +-- the suffix remaining after @'takeWhileL' p xs@. +dropWhileL :: (a -> Bool) -> Seq a -> Seq a +dropWhileL p = snd . spanl p + +-- | /O(i)/ where /i/ is the suffix length. @'dropWhileR' p xs@ returns +-- the prefix remaining after @'takeWhileR' p xs@. +-- +-- @'dropWhileR' p xs@ is equivalent to @'reverse' ('dropWhileL' p ('reverse' xs))@. +dropWhileR :: (a -> Bool) -> Seq a -> Seq a +dropWhileR p = snd . spanr p + +-- | /O(i)/ where /i/ is the prefix length. 'spanl', applied to +-- a predicate @p@ and a sequence @xs@, returns a pair whose first +-- element is the longest prefix (possibly empty) of @xs@ of elements that +-- satisfy @p@ and the second element is the remainder of the sequence. +spanl :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +spanl p = breakl (not . p) + +-- | /O(i)/ where /i/ is the suffix length. 'spanr', applied to a +-- predicate @p@ and a sequence @xs@, returns a pair whose /first/ element +-- is the longest /suffix/ (possibly empty) of @xs@ of elements that +-- satisfy @p@ and the second element is the remainder of the sequence. +spanr :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +spanr p = breakr (not . p) + +{-# INLINE breakl #-} +-- | /O(i)/ where /i/ is the breakpoint index. 'breakl', applied to a +-- predicate @p@ and a sequence @xs@, returns a pair whose first element +-- is the longest prefix (possibly empty) of @xs@ of elements that +-- /do not satisfy/ @p@ and the second element is the remainder of +-- the sequence. +-- +-- @'breakl' p@ is equivalent to @'spanl' (not . p)@. +breakl :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +breakl p xs = foldr (\ i _ -> splitAt i xs) (xs, empty) (findIndicesL p xs) + +{-# INLINE breakr #-} +-- | @'breakr' p@ is equivalent to @'spanr' (not . p)@. +breakr :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +breakr p xs = foldr (\ i _ -> flipPair (splitAt (i + 1) xs)) (xs, empty) (findIndicesR p xs) + where flipPair (x, y) = (y, x) + +-- | /O(n)/. The 'partition' function takes a predicate @p@ and a +-- sequence @xs@ and returns sequences of those elements which do and +-- do not satisfy the predicate. +partition :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +partition p = toPair . foldl' part (empty :*: empty) + where + part (xs :*: ys) x + | p x = (xs `snoc'` x) :*: ys + | otherwise = xs :*: (ys `snoc'` x) + +-- | /O(n)/. The 'filter' function takes a predicate @p@ and a sequence +-- @xs@ and returns a sequence of those elements which satisfy the +-- predicate. +filter :: (a -> Bool) -> Seq a -> Seq a +filter p = foldl' (\ xs x -> if p x then xs `snoc'` x else xs) empty + +-- Indexing sequences + +-- | 'elemIndexL' finds the leftmost index of the specified element, +-- if it is present, and otherwise 'Nothing'. +elemIndexL :: Eq a => a -> Seq a -> Maybe Int +elemIndexL x = findIndexL (x ==) + +-- | 'elemIndexR' finds the rightmost index of the specified element, +-- if it is present, and otherwise 'Nothing'. +elemIndexR :: Eq a => a -> Seq a -> Maybe Int +elemIndexR x = findIndexR (x ==) + +-- | 'elemIndicesL' finds the indices of the specified element, from +-- left to right (i.e. in ascending order). +elemIndicesL :: Eq a => a -> Seq a -> [Int] +elemIndicesL x = findIndicesL (x ==) + +-- | 'elemIndicesR' finds the indices of the specified element, from +-- right to left (i.e. in descending order). +elemIndicesR :: Eq a => a -> Seq a -> [Int] +elemIndicesR x = findIndicesR (x ==) + +-- | @'findIndexL' p xs@ finds the index of the leftmost element that +-- satisfies @p@, if any exist. +findIndexL :: (a -> Bool) -> Seq a -> Maybe Int +findIndexL p = listToMaybe' . findIndicesL p + +-- | @'findIndexR' p xs@ finds the index of the rightmost element that +-- satisfies @p@, if any exist. +findIndexR :: (a -> Bool) -> Seq a -> Maybe Int +findIndexR p = listToMaybe' . findIndicesR p + +{-# INLINE findIndicesL #-} +-- | @'findIndicesL' p@ finds all indices of elements that satisfy @p@, +-- in ascending order. +findIndicesL :: (a -> Bool) -> Seq a -> [Int] +#if __GLASGOW_HASKELL__ +findIndicesL p xs = build (\ c n -> let g i x z = if p x then c i z else z in + foldrWithIndex g n xs) +#else +findIndicesL p xs = foldrWithIndex g [] xs + where g i x is = if p x then i:is else is +#endif + +{-# INLINE findIndicesR #-} +-- | @'findIndicesR' p@ finds all indices of elements that satisfy @p@, +-- in descending order. +findIndicesR :: (a -> Bool) -> Seq a -> [Int] +#if __GLASGOW_HASKELL__ +findIndicesR p xs = build (\ c n -> + let g z i x = if p x then c i z else z in foldlWithIndex g n xs) +#else +findIndicesR p xs = foldlWithIndex g [] xs + where g is i x = if p x then i:is else is +#endif + +------------------------------------------------------------------------ +-- Lists +------------------------------------------------------------------------ + +-- The implementation below is based on an idea by Ross Paterson and +-- implemented by Lennart Spitzner. It avoids the rebuilding the original +-- (|>)-based implementation suffered from. It also avoids the excessive pair +-- allocations Paterson's implementation suffered from. +-- +-- David Feuer suggested building in nine-element chunks, which reduces +-- intermediate conses from around (1/2)*n to around (1/8)*n with a concomitant +-- improvement in benchmark constant factors. In fact, it should be even +-- better to work in chunks of 27 `Elem`s and chunks of three `Node`s, rather +-- than nine of each, but it seems hard to avoid a code explosion with +-- such large chunks. +-- +-- Paterson's code can be seen, for example, in +-- https://github.com/haskell/containers/blob/74034b3244fa4817c7bef1202e639b887a975d9e/Data/Sequence.hs#L3532 +-- +-- Given a list +-- +-- [1..302] +-- +-- the original code forms Three 1 2 3 | [node3 4 5 6, node3 7 8 9, node3 10 11 +-- 12, ...] | Two 301 302 +-- +-- Then it recurses on the middle list. The middle lists become successively +-- shorter as their elements become successively deeper nodes. +-- +-- The original implementation of the list shortener, getNodes, included the +-- recursive step + +-- getNodes s x1 (x2:x3:x4:xs) = (Node3 s x1 x2 x3:ns, d) +-- where (ns, d) = getNodes s x4 xs + +-- This allocates a cons and a lazy pair at each 3-element step. It relies on +-- the Haskell implementation using Wadler's technique, described in "Fixing +-- some space leaks with a garbage collector" +-- http://homepages.inf.ed.ac.uk/wadler/papers/leak/leak.ps.gz, to repeatedly +-- simplify the `d` thunk. Although GHC uses this GC trick, heap profiling at +-- least appears to indicate that the pair constructors and conses build up +-- with this implementation. +-- +-- Spitzner's implementation uses a similar approach, but replaces the middle +-- list, in each level, with a customized stream type that finishes off with +-- the final digit in that level and (since it works in nines) in the one +-- above. To work around the nested tree structure, the overall computation is +-- structured using continuation-passing style, with a function that, at the +-- bottom of the tree, deals with a stream that terminates in a nested-pair +-- representation of the entire right side of the tree. Perhaps someone will +-- eventually find a less mind-bending way to accomplish this. + +-- | /O(n)/. Create a sequence from a finite list of elements. +-- There is a function 'toList' in the opposite direction for all +-- instances of the 'Foldable' class, including 'Seq'. +fromList :: [a] -> Seq a +-- Note: we can avoid map_elem if we wish by scattering +-- Elem applications throughout mkTreeE and getNodesE, but +-- it gets a bit hard to read. +fromList = Seq . mkTree . map_elem + where +#ifdef __GLASGOW_HASKELL__ + mkTree :: forall a' . [Elem a'] -> FingerTree (Elem a') +#else + mkTree :: [Elem a] -> FingerTree (Elem a) +#endif + mkTree [] = EmptyT + mkTree [x1] = Single x1 + mkTree [x1, x2] = Deep 2 (One x1) EmptyT (One x2) + mkTree [x1, x2, x3] = Deep 3 (Two x1 x2) EmptyT (One x3) + mkTree [x1, x2, x3, x4] = Deep 4 (Two x1 x2) EmptyT (Two x3 x4) + mkTree [x1, x2, x3, x4, x5] = Deep 5 (Three x1 x2 x3) EmptyT (Two x4 x5) + mkTree [x1, x2, x3, x4, x5, x6] = + Deep 6 (Three x1 x2 x3) EmptyT (Three x4 x5 x6) + mkTree [x1, x2, x3, x4, x5, x6, x7] = + Deep 7 (Two x1 x2) (Single (Node3 3 x3 x4 x5)) (Two x6 x7) + mkTree [x1, x2, x3, x4, x5, x6, x7, x8] = + Deep 8 (Three x1 x2 x3) (Single (Node3 3 x4 x5 x6)) (Two x7 x8) + mkTree [x1, x2, x3, x4, x5, x6, x7, x8, x9] = + Deep 9 (Three x1 x2 x3) (Single (Node3 3 x4 x5 x6)) (Three x7 x8 x9) + mkTree [x1, x2, x3, x4, x5, x6, x7, x8, y0, y1] = + Deep 10 (Two x1 x2) + (Deep 6 (One (Node3 3 x3 x4 x5)) EmptyT (One (Node3 3 x6 x7 x8))) + (Two y0 y1) + mkTree [x1, x2, x3, x4, x5, x6, x7, x8, x9, y0, y1] = + Deep 11 (Three x1 x2 x3) + (Deep 6 (One (Node3 3 x4 x5 x6)) EmptyT (One (Node3 3 x7 x8 x9))) + (Two y0 y1) + mkTree [x1, x2, x3, x4, x5, x6, x7, x8, x9, y0, y1, y2] = + Deep 12 (Three x1 x2 x3) + (Deep 6 (One (Node3 3 x4 x5 x6)) EmptyT (One (Node3 3 x7 x8 x9))) + (Three y0 y1 y2) + mkTree [x1, x2, x3, x4, x5, x6, x7, x8, y0, y1, y2, y3, y4] = + Deep 13 (Two x1 x2) + (Deep 9 (Two (Node3 3 x3 x4 x5) (Node3 3 x6 x7 x8)) EmptyT (One (Node3 3 y0 y1 y2))) + (Two y3 y4) + mkTree [x1, x2, x3, x4, x5, x6, x7, x8, x9, y0, y1, y2, y3, y4] = + Deep 14 (Three x1 x2 x3) + (Deep 9 (Two (Node3 3 x4 x5 x6) (Node3 3 x7 x8 x9)) EmptyT (One (Node3 3 y0 y1 y2))) + (Two y3 y4) + mkTree [x1, x2, x3, x4, x5, x6, x7, x8, x9, y0, y1, y2, y3, y4, y5] = + Deep 15 (Three x1 x2 x3) + (Deep 9 (Two (Node3 3 x4 x5 x6) (Node3 3 x7 x8 x9)) EmptyT (One (Node3 3 y0 y1 y2))) + (Three y3 y4 y5) + mkTree (x1:x2:x3:x4:x5:x6:x7:x8:x9:y0:y1:y2:y3:y4:y5:y6:xs) = + mkTreeC cont 9 (getNodes 3 (Node3 3 y3 y4 y5) y6 xs) + where + d2 = Three x1 x2 x3 + d1 = Three (Node3 3 x4 x5 x6) (Node3 3 x7 x8 x9) (Node3 3 y0 y1 y2) +#ifdef __GLASGOW_HASKELL__ + cont :: (Digit (Node (Elem a')), Digit (Elem a')) -> FingerTree (Node (Node (Elem a'))) -> FingerTree (Elem a') +#endif + cont (!r1, !r2) !sub = + let !sub1 = Deep (9 + size r1 + size sub) d1 sub r1 + in Deep (3 + size r2 + size sub1) d2 sub1 r2 + + getNodes :: forall a . Int + -> Node a + -> a + -> [a] + -> ListFinal (Node (Node a)) (Digit (Node a), Digit a) + getNodes !_ n1 x1 [] = LFinal (One n1, One x1) + getNodes _ n1 x1 [x2] = LFinal (One n1, Two x1 x2) + getNodes _ n1 x1 [x2, x3] = LFinal (One n1, Three x1 x2 x3) + getNodes s n1 x1 [x2, x3, x4] = LFinal (Two n1 (Node3 s x1 x2 x3), One x4) + getNodes s n1 x1 [x2, x3, x4, x5] = LFinal (Two n1 (Node3 s x1 x2 x3), Two x4 x5) + getNodes s n1 x1 [x2, x3, x4, x5, x6] = LFinal (Two n1 (Node3 s x1 x2 x3), Three x4 x5 x6) + getNodes s n1 x1 [x2, x3, x4, x5, x6, x7] = LFinal (Three n1 (Node3 s x1 x2 x3) (Node3 s x4 x5 x6), One x7) + getNodes s n1 x1 [x2, x3, x4, x5, x6, x7, x8] = LFinal (Three n1 (Node3 s x1 x2 x3) (Node3 s x4 x5 x6), Two x7 x8) + getNodes s n1 x1 [x2, x3, x4, x5, x6, x7, x8, x9] = LFinal (Three n1 (Node3 s x1 x2 x3) (Node3 s x4 x5 x6), Three x7 x8 x9) + getNodes s n1 x1 (x2:x3:x4:x5:x6:x7:x8:x9:x10:xs) = LCons n10 (getNodes s (Node3 s x7 x8 x9) x10 xs) + where !n2 = Node3 s x1 x2 x3 + !n3 = Node3 s x4 x5 x6 + !n10 = Node3 (3*s) n1 n2 n3 + + mkTreeC :: +#ifdef __GLASGOW_HASKELL__ + forall a b c . +#endif + (b -> FingerTree (Node a) -> c) + -> Int + -> ListFinal (Node a) b + -> c + mkTreeC cont !_ (LFinal b) = + cont b EmptyT + mkTreeC cont _ (LCons x1 (LFinal b)) = + cont b (Single x1) + mkTreeC cont s (LCons x1 (LCons x2 (LFinal b))) = + cont b (Deep (2*s) (One x1) EmptyT (One x2)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LFinal b)))) = + cont b (Deep (3*s) (Two x1 x2) EmptyT (One x3)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LFinal b))))) = + cont b (Deep (4*s) (Two x1 x2) EmptyT (Two x3 x4)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LFinal b)))))) = + cont b (Deep (5*s) (Three x1 x2 x3) EmptyT (Two x4 x5)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LFinal b))))))) = + cont b (Deep (6*s) (Three x1 x2 x3) EmptyT (Three x4 x5 x6)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LFinal b)))))))) = + cont b (Deep (7*s) (Two x1 x2) (Single (Node3 (3*s) x3 x4 x5)) (Two x6 x7)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LFinal b))))))))) = + cont b (Deep (8*s) (Three x1 x2 x3) (Single (Node3 (3*s) x4 x5 x6)) (Two x7 x8)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LCons x9 (LFinal b)))))))))) = + cont b (Deep (9*s) (Three x1 x2 x3) (Single (Node3 (3*s) x4 x5 x6)) (Three x7 x8 x9)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LCons y0 (LCons y1 (LFinal b))))))))))) = + cont b (Deep (10*s) (Two x1 x2) (Deep (6*s) (One (Node3 (3*s) x3 x4 x5)) EmptyT (One (Node3 (3*s) x6 x7 x8))) (Two y0 y1)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LCons x9 (LCons y0 (LCons y1 (LFinal b)))))))))))) = + cont b (Deep (11*s) (Three x1 x2 x3) (Deep (6*s) (One (Node3 (3*s) x4 x5 x6)) EmptyT (One (Node3 (3*s) x7 x8 x9))) (Two y0 y1)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LCons x9 (LCons y0 (LCons y1 (LCons y2 (LFinal b))))))))))))) = + cont b (Deep (12*s) (Three x1 x2 x3) (Deep (6*s) (One (Node3 (3*s) x4 x5 x6)) EmptyT (One (Node3 (3*s) x7 x8 x9))) (Three y0 y1 y2)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LCons y0 (LCons y1 (LCons y2 (LCons y3 (LCons y4 (LFinal b)))))))))))))) = + cont b (Deep (13*s) (Two x1 x2) (Deep (9*s) (Two (Node3 (3*s) x3 x4 x5) (Node3 (3*s) x6 x7 x8)) EmptyT (One (Node3 (3*s) y0 y1 y2))) (Two y3 y4)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LCons x9 (LCons y0 (LCons y1 (LCons y2 (LCons y3 (LCons y4 (LFinal b))))))))))))))) = + cont b (Deep (14*s) (Three x1 x2 x3) (Deep (9*s) (Two (Node3 (3*s) x4 x5 x6) (Node3 (3*s) x7 x8 x9)) EmptyT (One (Node3 (3*s) y0 y1 y2))) (Two y3 y4)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LCons x9 (LCons y0 (LCons y1 (LCons y2 (LCons y3 (LCons y4 (LCons y5 (LFinal b)))))))))))))))) = + cont b (Deep (15*s) (Three x1 x2 x3) (Deep (9*s) (Two (Node3 (3*s) x4 x5 x6) (Node3 (3*s) x7 x8 x9)) EmptyT (One (Node3 (3*s) y0 y1 y2))) (Three y3 y4 y5)) + mkTreeC cont s (LCons x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LCons x9 (LCons y0 (LCons y1 (LCons y2 (LCons y3 (LCons y4 (LCons y5 (LCons y6 xs)))))))))))))))) = + mkTreeC cont2 (9*s) (getNodesC (3*s) (Node3 (3*s) y3 y4 y5) y6 xs) + where +#ifdef __GLASGOW_HASKELL__ + cont2 :: (b, Digit (Node (Node a)), Digit (Node a)) -> FingerTree (Node (Node (Node a))) -> c +#endif + cont2 (b, r1, r2) !sub = + let d2 = Three x1 x2 x3 + d1 = Three (Node3 (3*s) x4 x5 x6) (Node3 (3*s) x7 x8 x9) (Node3 (3*s) y0 y1 y2) + !sub1 = Deep (9*s + size r1 + size sub) d1 sub r1 + in cont b $! Deep (3*s + size r2 + size sub1) d2 sub1 r2 + + getNodesC :: Int + -> Node a + -> a + -> ListFinal a b + -> ListFinal (Node (Node a)) (b, Digit (Node a), Digit a) + getNodesC !_ n1 x1 (LFinal b) = LFinal $ (b, One n1, One x1) + getNodesC _ n1 x1 (LCons x2 (LFinal b)) = LFinal $ (b, One n1, Two x1 x2) + getNodesC _ n1 x1 (LCons x2 (LCons x3 (LFinal b))) = LFinal $ (b, One n1, Three x1 x2 x3) + getNodesC s n1 x1 (LCons x2 (LCons x3 (LCons x4 (LFinal b)))) = + let !n2 = Node3 s x1 x2 x3 + in LFinal $ (b, Two n1 n2, One x4) + getNodesC s n1 x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LFinal b))))) = + let !n2 = Node3 s x1 x2 x3 + in LFinal $ (b, Two n1 n2, Two x4 x5) + getNodesC s n1 x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LFinal b)))))) = + let !n2 = Node3 s x1 x2 x3 + in LFinal $ (b, Two n1 n2, Three x4 x5 x6) + getNodesC s n1 x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LFinal b))))))) = + let !n2 = Node3 s x1 x2 x3 + !n3 = Node3 s x4 x5 x6 + in LFinal $ (b, Three n1 n2 n3, One x7) + getNodesC s n1 x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LFinal b)))))))) = + let !n2 = Node3 s x1 x2 x3 + !n3 = Node3 s x4 x5 x6 + in LFinal $ (b, Three n1 n2 n3, Two x7 x8) + getNodesC s n1 x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LCons x9 (LFinal b))))))))) = + let !n2 = Node3 s x1 x2 x3 + !n3 = Node3 s x4 x5 x6 + in LFinal $ (b, Three n1 n2 n3, Three x7 x8 x9) + getNodesC s n1 x1 (LCons x2 (LCons x3 (LCons x4 (LCons x5 (LCons x6 (LCons x7 (LCons x8 (LCons x9 (LCons x10 xs))))))))) = + LCons n10 $ getNodesC s (Node3 s x7 x8 x9) x10 xs + where !n2 = Node3 s x1 x2 x3 + !n3 = Node3 s x4 x5 x6 + !n10 = Node3 (3*s) n1 n2 n3 + + map_elem :: [a] -> [Elem a] +#if __GLASGOW_HASKELL__ >= 708 + map_elem xs = coerce xs +#else + map_elem xs = Data.List.map Elem xs +#endif + {-# INLINE map_elem #-} + +-- essentially: Free ((,) a) b. +data ListFinal a cont = LFinal !cont | LCons !a (ListFinal a cont) + +#if __GLASGOW_HASKELL__ >= 708 +instance GHC.Exts.IsList (Seq a) where + type Item (Seq a) = a + fromList = fromList + fromListN = fromList2 + toList = toList +#endif + +#ifdef __GLASGOW_HASKELL__ +instance IsString (Seq Char) where + fromString = fromList +#endif + +------------------------------------------------------------------------ +-- Reverse +------------------------------------------------------------------------ + +-- | /O(n)/. The reverse of a sequence. +reverse :: Seq a -> Seq a +reverse (Seq xs) = Seq (fmapReverseTree id xs) + +#ifdef __GLASGOW_HASKELL__ +{-# NOINLINE [1] reverse #-} + +-- | /O(n)/. Reverse a sequence while mapping over it. This is not +-- currently exported, but is used in rewrite rules. +fmapReverse :: (a -> b) -> Seq a -> Seq b +fmapReverse f (Seq xs) = Seq (fmapReverseTree (lift_elem f) xs) + where + lift_elem :: (a -> b) -> (Elem a -> Elem b) +#if __GLASGOW_HASKELL__ >= 708 + lift_elem = coerce +#else + lift_elem g (Elem a) = Elem (g a) +#endif + +-- If we're mapping over a sequence, we can reverse it at the same time +-- at no extra charge. +{-# RULES +"fmapSeq/reverse" forall f xs . fmapSeq f (reverse xs) = fmapReverse f xs +"reverse/fmapSeq" forall f xs . reverse (fmapSeq f xs) = fmapReverse f xs + #-} +#endif + +fmapReverseTree :: (a -> b) -> FingerTree a -> FingerTree b +fmapReverseTree _ EmptyT = EmptyT +fmapReverseTree f (Single x) = Single (f x) +fmapReverseTree f (Deep s pr m sf) = + Deep s (reverseDigit f sf) + (fmapReverseTree (reverseNode f) m) + (reverseDigit f pr) + +{-# INLINE reverseDigit #-} +reverseDigit :: (a -> b) -> Digit a -> Digit b +reverseDigit f (One a) = One (f a) +reverseDigit f (Two a b) = Two (f b) (f a) +reverseDigit f (Three a b c) = Three (f c) (f b) (f a) +reverseDigit f (Four a b c d) = Four (f d) (f c) (f b) (f a) + +reverseNode :: (a -> b) -> Node a -> Node b +reverseNode f (Node2 s a b) = Node2 s (f b) (f a) +reverseNode f (Node3 s a b c) = Node3 s (f c) (f b) (f a) + +------------------------------------------------------------------------ +-- Mapping with a splittable value +------------------------------------------------------------------------ + +-- For zipping, it is useful to build a result by +-- traversing a sequence while splitting up something else. For zipping, we +-- traverse the first sequence while splitting up the second. +-- +-- What makes all this crazy code a good idea: +-- +-- Suppose we zip together two sequences of the same length: +-- +-- zs = zip xs ys +-- +-- We want to get reasonably fast indexing into zs immediately, rather than +-- needing to construct the entire thing first, as the previous implementation +-- required. The first aspect is that we build the result "outside-in" or +-- "top-down", rather than left to right. That gives us access to both ends +-- quickly. But that's not enough, by itself, to give immediate access to the +-- center of zs. For that, we need to be able to skip over larger segments of +-- zs, delaying their construction until we actually need them. The way we do +-- this is to traverse xs, while splitting up ys according to the structure of +-- xs. If we have a Deep _ pr m sf, we split ys into three pieces, and hand off +-- one piece to the prefix, one to the middle, and one to the suffix of the +-- result. The key point is that we don't need to actually do anything further +-- with those pieces until we actually need them; the computations to split +-- them up further and zip them with their matching pieces can be delayed until +-- they're actually needed. We do the same thing for Digits (splitting into +-- between one and four pieces) and Nodes (splitting into two or three). The +-- ultimate result is that we can index into, or split at, any location in zs +-- in polylogarithmic time *immediately*, while still being able to force all +-- the thunks in O(n) time. +-- +-- Benchmark info, and alternatives: +-- +-- The old zipping code used mapAccumL to traverse the first sequence while +-- cutting down the second sequence one piece at a time. +-- +-- An alternative way to express that basic idea is to convert both sequences +-- to lists, zip the lists, and then convert the result back to a sequence. +-- I'll call this the "listy" implementation. +-- +-- I benchmarked two operations: Each started by zipping two sequences +-- constructed with replicate and/or fromList. The first would then immediately +-- index into the result. The second would apply deepseq to force the entire +-- result. The new implementation worked much better than either of the others +-- on the immediate indexing test, as expected. It also worked better than the +-- old implementation for all the deepseq tests. For short sequences, the listy +-- implementation outperformed all the others on the deepseq test. However, the +-- splitting implementation caught up and surpassed it once the sequences grew +-- long enough. It seems likely that by avoiding rebuilding, it interacts +-- better with the cache hierarchy. +-- +-- David Feuer, with some guidance from Carter Schonwald, December 2014 + +-- | /O(n)/. Constructs a new sequence with the same structure as an existing +-- sequence using a user-supplied mapping function along with a splittable +-- value and a way to split it. The value is split up lazily according to the +-- structure of the sequence, so one piece of the value is distributed to each +-- element of the sequence. The caller should provide a splitter function that +-- takes a number, @n@, and a splittable value, breaks off a chunk of size @n@ +-- from the value, and returns that chunk and the remainder as a pair. The +-- following examples will hopefully make the usage clear: +-- +-- > zipWith :: (a -> b -> c) -> Seq a -> Seq b -> Seq c +-- > zipWith f s1 s2 = splitMap splitAt (\b a -> f a (b `index` 0)) s2' s1' +-- > where +-- > minLen = min (length s1) (length s2) +-- > s1' = take minLen s1 +-- > s2' = take minLen s2 +-- +-- > mapWithIndex :: (Int -> a -> b) -> Seq a -> Seq b +-- > mapWithIndex f = splitMap (\n i -> (i, n+i)) f 0 +#ifdef __GLASGOW_HASKELL__ +-- We use ScopedTypeVariables to improve performance and make +-- performance less sensitive to minor changes. + +-- We INLINE this so GHC can see that the function passed in is +-- strict in its Int argument. +{-# INLINE splitMap #-} +splitMap :: forall s a' b' . (Int -> s -> (s,s)) -> (s -> a' -> b') -> s -> Seq a' -> Seq b' +splitMap splt f0 s0 (Seq xs0) = Seq $ splitMapTreeE (\s' (Elem a) -> Elem (f0 s' a)) s0 xs0 + where + {-# INLINE splitMapTreeE #-} + splitMapTreeE :: (s -> Elem y -> b) -> s -> FingerTree (Elem y) -> FingerTree b + splitMapTreeE _ _ EmptyT = EmptyT + splitMapTreeE f s (Single xs) = Single $ f s xs + splitMapTreeE f s (Deep n pr m sf) = Deep n (splitMapDigit f prs pr) (splitMapTreeN (\eta1 eta2 -> splitMapNode f eta1 eta2) ms m) (splitMapDigit f sfs sf) + where + !spr = size pr + !sm = n - spr - size sf + (prs, r) = splt spr s + (ms, sfs) = splt sm r + + splitMapTreeN :: (s -> Node a -> b) -> s -> FingerTree (Node a) -> FingerTree b + splitMapTreeN _ _ EmptyT = EmptyT + splitMapTreeN f s (Single xs) = Single $ f s xs + splitMapTreeN f s (Deep n pr m sf) = Deep n (splitMapDigit f prs pr) (splitMapTreeN (\eta1 eta2 -> splitMapNode f eta1 eta2) ms m) (splitMapDigit f sfs sf) + where + (prs, r) = splt (size pr) s + (ms, sfs) = splt (size m) r + + {-# INLINE splitMapDigit #-} + splitMapDigit :: Sized a => (s -> a -> b) -> s -> Digit a -> Digit b + splitMapDigit f s (One a) = One (f s a) + splitMapDigit f s (Two a b) = Two (f first a) (f second b) + where + (first, second) = splt (size a) s + splitMapDigit f s (Three a b c) = Three (f first a) (f second b) (f third c) + where + (first, r) = splt (size a) s + (second, third) = splt (size b) r + splitMapDigit f s (Four a b c d) = Four (f first a) (f second b) (f third c) (f fourth d) + where + (first, s') = splt (size a) s + (middle, fourth) = splt (size b + size c) s' + (second, third) = splt (size b) middle + + {-# INLINE splitMapNode #-} + splitMapNode :: Sized a => (s -> a -> b) -> s -> Node a -> Node b + splitMapNode f s (Node2 ns a b) = Node2 ns (f first a) (f second b) + where + (first, second) = splt (size a) s + splitMapNode f s (Node3 ns a b c) = Node3 ns (f first a) (f second b) (f third c) + where + (first, r) = splt (size a) s + (second, third) = splt (size b) r + +#else +-- Implementation without ScopedTypeVariables--somewhat slower, +-- and much more sensitive to minor changes in various places. + +{-# INLINE splitMap #-} +splitMap :: (Int -> s -> (s,s)) -> (s -> a -> b) -> s -> Seq a -> Seq b +splitMap splt' f0 s0 (Seq xs0) = Seq $ splitMapTreeE splt' (\s' (Elem a) -> Elem (f0 s' a)) s0 xs0 + +{-# INLINE splitMapTreeE #-} +splitMapTreeE :: (Int -> s -> (s,s)) -> (s -> Elem y -> b) -> s -> FingerTree (Elem y) -> FingerTree b +splitMapTreeE _ _ _ EmptyT = EmptyT +splitMapTreeE _ f s (Single xs) = Single $ f s xs +splitMapTreeE splt f s (Deep n pr m sf) = Deep n (splitMapDigit splt f prs pr) (splitMapTreeN splt (\eta1 eta2 -> splitMapNode splt f eta1 eta2) ms m) (splitMapDigit splt f sfs sf) + where + !spr = size pr + sm = n - spr - size sf + (prs, r) = splt spr s + (ms, sfs) = splt sm r + +splitMapTreeN :: (Int -> s -> (s,s)) -> (s -> Node a -> b) -> s -> FingerTree (Node a) -> FingerTree b +splitMapTreeN _ _ _ EmptyT = EmptyT +splitMapTreeN _ f s (Single xs) = Single $ f s xs +splitMapTreeN splt f s (Deep n pr m sf) = Deep n (splitMapDigit splt f prs pr) (splitMapTreeN splt (\eta1 eta2 -> splitMapNode splt f eta1 eta2) ms m) (splitMapDigit splt f sfs sf) + where + (prs, r) = splt (size pr) s + (ms, sfs) = splt (size m) r + +{-# INLINE splitMapDigit #-} +splitMapDigit :: Sized a => (Int -> s -> (s,s)) -> (s -> a -> b) -> s -> Digit a -> Digit b +splitMapDigit _ f s (One a) = One (f s a) +splitMapDigit splt f s (Two a b) = Two (f first a) (f second b) + where + (first, second) = splt (size a) s +splitMapDigit splt f s (Three a b c) = Three (f first a) (f second b) (f third c) + where + (first, r) = splt (size a) s + (second, third) = splt (size b) r +splitMapDigit splt f s (Four a b c d) = Four (f first a) (f second b) (f third c) (f fourth d) + where + (first, s') = splt (size a) s + (middle, fourth) = splt (size b + size c) s' + (second, third) = splt (size b) middle + +{-# INLINE splitMapNode #-} +splitMapNode :: Sized a => (Int -> s -> (s,s)) -> (s -> a -> b) -> s -> Node a -> Node b +splitMapNode splt f s (Node2 ns a b) = Node2 ns (f first a) (f second b) + where + (first, second) = splt (size a) s +splitMapNode splt f s (Node3 ns a b c) = Node3 ns (f first a) (f second b) (f third c) + where + (first, r) = splt (size a) s + (second, third) = splt (size b) r +#endif + +getSingleton :: Seq a -> a +getSingleton (Seq (Single (Elem a))) = a +getSingleton _ = error "getSingleton: Not a singleton." + +------------------------------------------------------------------------ +-- Zipping +------------------------------------------------------------------------ + +-- | /O(min(n1,n2))/. 'zip' takes two sequences and returns a sequence +-- of corresponding pairs. If one input is short, excess elements are +-- discarded from the right end of the longer sequence. +zip :: Seq a -> Seq b -> Seq (a, b) +zip = zipWith (,) + +-- | /O(min(n1,n2))/. 'zipWith' generalizes 'zip' by zipping with the +-- function given as the first argument, instead of a tupling function. +-- For example, @zipWith (+)@ is applied to two sequences to take the +-- sequence of corresponding sums. +zipWith :: (a -> b -> c) -> Seq a -> Seq b -> Seq c +zipWith f s1 s2 = zipWith' f s1' s2' + where + minLen = min (length s1) (length s2) + s1' = take minLen s1 + s2' = take minLen s2 + +-- | A version of zipWith that assumes the sequences have the same length. +zipWith' :: (a -> b -> c) -> Seq a -> Seq b -> Seq c +zipWith' f s1 s2 = splitMap uncheckedSplitAt (\s a -> f a (getSingleton s)) s2 s1 + +-- | /O(min(n1,n2,n3))/. 'zip3' takes three sequences and returns a +-- sequence of triples, analogous to 'zip'. +zip3 :: Seq a -> Seq b -> Seq c -> Seq (a,b,c) +zip3 = zipWith3 (,,) + +-- | /O(min(n1,n2,n3))/. 'zipWith3' takes a function which combines +-- three elements, as well as three sequences and returns a sequence of +-- their point-wise combinations, analogous to 'zipWith'. +zipWith3 :: (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d +zipWith3 f s1 s2 s3 = zipWith' ($) (zipWith' f s1' s2') s3' + where + minLen = minimum [length s1, length s2, length s3] + s1' = take minLen s1 + s2' = take minLen s2 + s3' = take minLen s3 + +zipWith3' :: (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d +zipWith3' f s1 s2 s3 = zipWith' ($) (zipWith' f s1 s2) s3 + +-- | /O(min(n1,n2,n3,n4))/. 'zip4' takes four sequences and returns a +-- sequence of quadruples, analogous to 'zip'. +zip4 :: Seq a -> Seq b -> Seq c -> Seq d -> Seq (a,b,c,d) +zip4 = zipWith4 (,,,) + +-- | /O(min(n1,n2,n3,n4))/. 'zipWith4' takes a function which combines +-- four elements, as well as four sequences and returns a sequence of +-- their point-wise combinations, analogous to 'zipWith'. +zipWith4 :: (a -> b -> c -> d -> e) -> Seq a -> Seq b -> Seq c -> Seq d -> Seq e +zipWith4 f s1 s2 s3 s4 = zipWith' ($) (zipWith3' f s1' s2' s3') s4' + where + minLen = minimum [length s1, length s2, length s3, length s4] + s1' = take minLen s1 + s2' = take minLen s2 + s3' = take minLen s3 + s4' = take minLen s4 + +------------------------------------------------------------------------ +-- Sorting +-- +-- sort and sortBy are implemented by simple deforestations of +-- \ xs -> fromList2 (length xs) . Data.List.sortBy cmp . toList +-- which does not get deforested automatically, it would appear. +-- +-- Unstable sorting is performed by a heap sort implementation based on +-- pairing heaps. Because the internal structure of sequences is quite +-- varied, it is difficult to get blocks of elements of roughly the same +-- length, which would improve merge sort performance. Pairing heaps, +-- on the other hand, are relatively resistant to the effects of merging +-- heaps of wildly different sizes, as guaranteed by its amortized +-- constant-time merge operation. Moreover, extensive use of SpecConstr +-- transformations can be done on pairing heaps, especially when we're +-- only constructing them to immediately be unrolled. +-- +-- On purely random sequences of length 50000, with no RTS options, +-- I get the following statistics, in which heapsort is about 42.5% +-- faster: (all comparisons done with -O2) +-- +-- Times (ms) min mean +/-sd median max +-- to/from list: 103.802 108.572 7.487 106.436 143.339 +-- unstable heapsort: 60.686 62.968 4.275 61.187 79.151 +-- +-- Heapsort, it would seem, is less of a memory hog than Data.List.sortBy. +-- The gap is narrowed when more memory is available, but heapsort still +-- wins, 15% faster, with +RTS -H128m: +-- +-- Times (ms) min mean +/-sd median max +-- to/from list: 42.692 45.074 2.596 44.600 56.601 +-- unstable heapsort: 37.100 38.344 3.043 37.715 55.526 +-- +-- In addition, on strictly increasing sequences the gap is even wider +-- than normal; heapsort is 68.5% faster with no RTS options: +-- Times (ms) min mean +/-sd median max +-- to/from list: 52.236 53.574 1.987 53.034 62.098 +-- unstable heapsort: 16.433 16.919 0.931 16.681 21.622 +-- +-- This may be attributed to the elegant nature of the pairing heap. +-- +-- wasserman.louis@gmail.com, 7/20/09 +------------------------------------------------------------------------ + +-- | /O(n log n)/. 'sort' sorts the specified 'Seq' by the natural +-- ordering of its elements. The sort is stable. +-- If stability is not required, 'unstableSort' can be considerably +-- faster, and in particular uses less memory. +sort :: Ord a => Seq a -> Seq a +sort = sortBy compare + +-- | /O(n log n)/. 'sortBy' sorts the specified 'Seq' according to the +-- specified comparator. The sort is stable. +-- If stability is not required, 'unstableSortBy' can be considerably +-- faster, and in particular uses less memory. +sortBy :: (a -> a -> Ordering) -> Seq a -> Seq a +sortBy cmp xs = fromList2 (length xs) (Data.List.sortBy cmp (toList xs)) + +-- | /O(n log n)/. 'unstableSort' sorts the specified 'Seq' by +-- the natural ordering of its elements, but the sort is not stable. +-- This algorithm is frequently faster and uses less memory than 'sort', +-- and performs extremely well -- frequently twice as fast as 'sort' -- +-- when the sequence is already nearly sorted. +unstableSort :: Ord a => Seq a -> Seq a +unstableSort = unstableSortBy compare + +-- | /O(n log n)/. A generalization of 'unstableSort', 'unstableSortBy' +-- takes an arbitrary comparator and sorts the specified sequence. +-- The sort is not stable. This algorithm is frequently faster and +-- uses less memory than 'sortBy', and performs extremely well -- +-- frequently twice as fast as 'sortBy' -- when the sequence is already +-- nearly sorted. +unstableSortBy :: (a -> a -> Ordering) -> Seq a -> Seq a +unstableSortBy cmp (Seq xs) = + fromList2 (size xs) $ maybe [] (unrollPQ cmp) $ + toPQ cmp (\ (Elem x) -> PQueue x Nil) xs + +-- | fromList2, given a list and its length, constructs a completely +-- balanced Seq whose elements are that list using the replicateA +-- generalization. +fromList2 :: Int -> [a] -> Seq a +fromList2 n = execState (replicateA n (State ht)) + where + ht (x:xs) = (xs, x) + ht [] = error "fromList2: short list" + +-- | A 'PQueue' is a simple pairing heap. +data PQueue e = PQueue e (PQL e) +data PQL e = Nil | {-# UNPACK #-} !(PQueue e) :& PQL e + +infixr 8 :& + +#if TESTING + +instance Functor PQueue where + fmap f (PQueue x ts) = PQueue (f x) (fmap f ts) + +instance Functor PQL where + fmap f (q :& qs) = fmap f q :& fmap f qs + fmap _ Nil = Nil + +instance Show e => Show (PQueue e) where + show = unlines . draw . fmap show + +-- borrowed wholesale from Data.Tree, as Data.Tree actually depends +-- on Data.Sequence +draw :: PQueue String -> [String] +draw (PQueue x ts0) = x : drawSubTrees ts0 + where + drawSubTrees Nil = [] + drawSubTrees (t :& Nil) = + "|" : shift "`- " " " (draw t) + drawSubTrees (t :& ts) = + "|" : shift "+- " "| " (draw t) ++ drawSubTrees ts + + shift first other = Data.List.zipWith (++) (first : repeat other) +#endif + +-- | 'unrollPQ', given a comparator function, unrolls a 'PQueue' into +-- a sorted list. +unrollPQ :: (e -> e -> Ordering) -> PQueue e -> [e] +unrollPQ cmp = unrollPQ' + where + {-# INLINE unrollPQ' #-} + unrollPQ' (PQueue x ts) = x:mergePQs0 ts + (<+>) = mergePQ cmp + mergePQs0 Nil = [] + mergePQs0 (t :& Nil) = unrollPQ' t + mergePQs0 (t1 :& t2 :& ts) = mergePQs (t1 <+> t2) ts + mergePQs !t ts = case ts of + Nil -> unrollPQ' t + t1 :& Nil -> unrollPQ' (t <+> t1) + t1 :& t2 :& ts' -> mergePQs (t <+> (t1 <+> t2)) ts' + +-- | 'toPQ', given an ordering function and a mechanism for queueifying +-- elements, converts a 'FingerTree' to a 'PQueue'. +toPQ :: (e -> e -> Ordering) -> (a -> PQueue e) -> FingerTree a -> Maybe (PQueue e) +toPQ _ _ EmptyT = Nothing +toPQ _ f (Single x) = Just (f x) +toPQ cmp f (Deep _ pr m sf) = Just (maybe (pr' <+> sf') ((pr' <+> sf') <+>) (toPQ cmp fNode m)) + where + fDigit digit = case fmap f digit of + One a -> a + Two a b -> a <+> b + Three a b c -> a <+> b <+> c + Four a b c d -> (a <+> b) <+> (c <+> d) + (<+>) = mergePQ cmp + fNode = fDigit . nodeToDigit + pr' = fDigit pr + sf' = fDigit sf + +-- | 'mergePQ' merges two 'PQueue's. +mergePQ :: (a -> a -> Ordering) -> PQueue a -> PQueue a -> PQueue a +mergePQ cmp q1@(PQueue x1 ts1) q2@(PQueue x2 ts2) + | cmp x1 x2 == GT = PQueue x2 (q1 :& ts2) + | otherwise = PQueue x1 (q2 :& ts1) diff --git a/Data/Set.hs b/Data/Set.hs --- a/Data/Set.hs +++ b/Data/Set.hs @@ -28,11 +28,17 @@ -- * Stephen Adams, \"/Efficient sets: a balancing act/\", -- Journal of Functional Programming 3(4):553-562, October 1993, -- . --- -- * J. Nievergelt and E.M. Reingold, -- \"/Binary search trees of bounded balance/\", -- SIAM journal of computing 2(1), March 1973. -- +-- Bounds for 'union', 'intersection', and 'difference' are as given +-- by +-- +-- * Guy Blelloch, Daniel Ferizovic, and Yihan Sun, +-- \"/Just Join for Parallel Ordered Sets/\", +-- . +-- -- Note that the implementation is /left-biased/ -- the elements of a -- first argument are always preferred to the second, for example in -- 'union' or 'insert'. Of course, left-biasing can only be observed @@ -84,6 +90,9 @@ -- * Filter , S.filter + , takeWhileAntitone + , dropWhileAntitone + , spanAntitone , partition , split , splitMember @@ -94,6 +103,9 @@ , findIndex , elemAt , deleteAt + , S.take + , S.drop + , S.splitAt -- * Map , S.map @@ -129,7 +141,9 @@ , toAscList , toDescList , fromAscList + , fromDescList , fromDistinctAscList + , fromDistinctDescList -- * Debugging , showTree diff --git a/Data/Set/Base.hs b/Data/Set/Base.hs --- a/Data/Set/Base.hs +++ b/Data/Set/Base.hs @@ -1,4 +1,5 @@ {-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} #if __GLASGOW_HASKELL__ {-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} #endif @@ -12,6 +13,8 @@ #include "containers.h" +{-# OPTIONS_HADDOCK hide #-} + ----------------------------------------------------------------------------- -- | -- Module : Data.Set.Base @@ -35,11 +38,17 @@ -- * Stephen Adams, \"/Efficient sets: a balancing act/\", -- Journal of Functional Programming 3(4):553-562, October 1993, -- . --- -- * J. Nievergelt and E.M. Reingold, -- \"/Binary search trees of bounded balance/\", -- SIAM journal of computing 2(1), March 1973. -- +-- Bounds for 'union', 'intersection', and 'difference' are as given +-- by +-- +-- * Guy Blelloch, Daniel Ferizovic, and Yihan Sun, +-- \"/Just Join for Parallel Ordered Sets/\", +-- . +-- -- Note that the implementation is /left-biased/ -- the elements of a -- first argument are always preferred to the second, for example in -- 'union' or 'insert'. Of course, left-biasing can only be observed @@ -47,8 +56,8 @@ -- equality. -- -- /Warning/: The size of the set must not exceed @maxBound::Int@. Violation of --- this condition is not detected and if the size limit is exceeded, its --- behaviour is undefined. +-- this condition is not detected and if the size limit is exceeded, the +-- behavior of the set is completely undefined. ----------------------------------------------------------------------------- -- [Note: Using INLINABLE] @@ -87,11 +96,7 @@ -- floats out of its enclosing function and then it heap-allocates the -- dictionary and the argument. Maybe it floats out too late and strictness -- analyzer cannot see that these could be passed on stack. --- --- For example, change 'member' so that its local 'go' function is not passing --- argument x and then look at the resulting code for hedgeInt. - -- [Note: Order of constructors] -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -- The order of constructors of Set matters when considering performance. @@ -134,6 +139,9 @@ -- * Filter , filter + , takeWhileAntitone + , dropWhileAntitone + , spanAntitone , partition , split , splitMember @@ -144,6 +152,9 @@ , findIndex , elemAt , deleteAt + , take + , drop + , splitAt -- * Map , map @@ -180,6 +191,8 @@ , toDescList , fromAscList , fromDistinctAscList + , fromDescList + , fromDistinctDescList -- * Debugging , showTree @@ -193,7 +206,7 @@ , merge ) where -import Prelude hiding (filter,foldl,foldr,null,map) +import Prelude hiding (filter,foldl,foldr,null,map,take,drop,splitAt) import qualified Data.List as List import Data.Bits (shiftL, shiftR) #if !MIN_VERSION_base(4,8,0) @@ -208,6 +221,7 @@ import Data.Utils.StrictFold import Data.Utils.StrictPair +import Data.Utils.PtrEquality #if __GLASGOW_HASKELL__ import GHC.Exts ( build ) @@ -224,10 +238,10 @@ --------------------------------------------------------------------} infixl 9 \\ -- --- | /O(n+m)/. See 'difference'. +-- | /O(m*log(n\/m+1)), m <= n/. See 'difference'. (\\) :: Ord a => Set a -> Set a -> Set a m1 \\ m2 = difference m1 m2 -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE (\\) #-} #endif @@ -290,8 +304,7 @@ toList = toList {-# INLINE toList #-} elem = go - where STRICT_1_OF_2(go) - go _ Tip = False + where go !_ Tip = False go x (Bin _ y l r) = x == y || go x l || go x r {-# INLINABLE elem #-} minimum = findMin @@ -350,13 +363,12 @@ member :: Ord a => a -> Set a -> Bool member = go where - STRICT_1_OF_2(go) - go _ Tip = False + go !_ Tip = False go x (Bin _ y l r) = case compare x y of LT -> go x l GT -> go x r EQ -> True -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE member #-} #else {-# INLINE member #-} @@ -365,7 +377,7 @@ -- | /O(log n)/. Is the element not in the set? notMember :: Ord a => a -> Set a -> Bool notMember a t = not $ member a t -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE notMember #-} #else {-# INLINE notMember #-} @@ -378,16 +390,14 @@ lookupLT :: Ord a => a -> Set a -> Maybe a lookupLT = goNothing where - STRICT_1_OF_2(goNothing) - goNothing _ Tip = Nothing + goNothing !_ Tip = Nothing goNothing x (Bin _ y l r) | x <= y = goNothing x l | otherwise = goJust x y r - STRICT_1_OF_3(goJust) - goJust _ best Tip = Just best + goJust !_ best Tip = Just best goJust x best (Bin _ y l r) | x <= y = goJust x best l | otherwise = goJust x y r -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE lookupLT #-} #else {-# INLINE lookupLT #-} @@ -400,16 +410,14 @@ lookupGT :: Ord a => a -> Set a -> Maybe a lookupGT = goNothing where - STRICT_1_OF_2(goNothing) - goNothing _ Tip = Nothing + goNothing !_ Tip = Nothing goNothing x (Bin _ y l r) | x < y = goJust x y l | otherwise = goNothing x r - STRICT_1_OF_3(goJust) - goJust _ best Tip = Just best + goJust !_ best Tip = Just best goJust x best (Bin _ y l r) | x < y = goJust x y l | otherwise = goJust x best r -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE lookupGT #-} #else {-# INLINE lookupGT #-} @@ -423,18 +431,16 @@ lookupLE :: Ord a => a -> Set a -> Maybe a lookupLE = goNothing where - STRICT_1_OF_2(goNothing) - goNothing _ Tip = Nothing + goNothing !_ Tip = Nothing goNothing x (Bin _ y l r) = case compare x y of LT -> goNothing x l EQ -> Just y GT -> goJust x y r - STRICT_1_OF_3(goJust) - goJust _ best Tip = Just best + goJust !_ best Tip = Just best goJust x best (Bin _ y l r) = case compare x y of LT -> goJust x best l EQ -> Just y GT -> goJust x y r -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE lookupLE #-} #else {-# INLINE lookupLE #-} @@ -448,18 +454,16 @@ lookupGE :: Ord a => a -> Set a -> Maybe a lookupGE = goNothing where - STRICT_1_OF_2(goNothing) - goNothing _ Tip = Nothing + goNothing !_ Tip = Nothing goNothing x (Bin _ y l r) = case compare x y of LT -> goJust x y l EQ -> Just y GT -> goNothing x r - STRICT_1_OF_3(goJust) - goJust _ best Tip = Just best + goJust !_ best Tip = Just best goJust x best (Bin _ y l r) = case compare x y of LT -> goJust x y l EQ -> Just y GT -> goJust x best r -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE lookupGE #-} #else {-# INLINE lookupGE #-} @@ -490,13 +494,17 @@ insert = go where go :: Ord a => a -> Set a -> Set a - STRICT_1_OF_2(go) - go x Tip = singleton x - go x (Bin sz y l r) = case compare x y of - LT -> balanceL y (go x l) r - GT -> balanceR y l (go x r) - EQ -> Bin sz x l r -#if __GLASGOW_HASKELL__ >= 700 + go !x Tip = singleton x + go !x t@(Bin sz y l r) = case compare x y of + LT | l' `ptrEq` l -> t + | otherwise -> balanceL y l' r + where !l' = go x l + GT | r' `ptrEq` r -> t + | otherwise -> balanceR y l r' + where !r' = go x r + EQ | x `ptrEq` y -> t + | otherwise -> Bin sz x l r +#if __GLASGOW_HASKELL__ {-# INLINABLE insert #-} #else {-# INLINE insert #-} @@ -510,13 +518,16 @@ insertR = go where go :: Ord a => a -> Set a -> Set a - STRICT_1_OF_2(go) - go x Tip = singleton x - go x t@(Bin _ y l r) = case compare x y of - LT -> balanceL y (go x l) r - GT -> balanceR y l (go x r) + go !x Tip = singleton x + go !x t@(Bin _ y l r) = case compare x y of + LT | l' `ptrEq` l -> t + | otherwise -> balanceL y l' r + where !l' = go x l + GT | r' `ptrEq` r -> t + | otherwise -> balanceR y l r' + where !r' = go x r EQ -> t -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE insertR #-} #else {-# INLINE insertR #-} @@ -529,13 +540,16 @@ delete = go where go :: Ord a => a -> Set a -> Set a - STRICT_1_OF_2(go) - go _ Tip = Tip - go x (Bin _ y l r) = case compare x y of - LT -> balanceR y (go x l) r - GT -> balanceL y l (go x r) + go !_ Tip = Tip + go x t@(Bin _ y l r) = case compare x y of + LT | l' `ptrEq` l -> t + | otherwise -> balanceR y l' r + where !l' = go x l + GT | r' `ptrEq` r -> t + | otherwise -> balanceL y l r' + where !r' = go x r EQ -> glue l r -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE delete #-} #else {-# INLINE delete #-} @@ -548,7 +562,7 @@ isProperSubsetOf :: Ord a => Set a -> Set a -> Bool isProperSubsetOf s1 s2 = (size s1 < size s2) && (isSubsetOf s1 s2) -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE isProperSubsetOf #-} #endif @@ -558,7 +572,7 @@ isSubsetOf :: Ord a => Set a -> Set a -> Bool isSubsetOf t1 t2 = (size t1 <= size t2) && (isSubsetOfX t1 t2) -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE isSubsetOf #-} #endif @@ -569,7 +583,7 @@ = found && isSubsetOfX l lt && isSubsetOfX r gt where (lt,found,gt) = splitMember x t -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE isSubsetOfX #-} #endif @@ -607,62 +621,49 @@ -- | The union of a list of sets: (@'unions' == 'foldl' 'union' 'empty'@). unions :: Ord a => [Set a] -> Set a unions = foldlStrict union empty -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE unions #-} #endif --- | /O(n+m)/. The union of two sets, preferring the first set when +-- | /O(m*log(n/m + 1)), m <= n/. The union of two sets, preferring the first set when -- equal elements are encountered. --- The implementation uses the efficient /hedge-union/ algorithm. union :: Ord a => Set a -> Set a -> Set a -union Tip t2 = t2 union t1 Tip = t1 -union t1 t2 = hedgeUnion NothingS NothingS t1 t2 -#if __GLASGOW_HASKELL__ >= 700 +union t1 (Bin _ x Tip Tip) = insertR x t1 +union (Bin _ x Tip Tip) t2 = insert x t2 +union Tip t2 = t2 +union t1@(Bin _ x l1 r1) t2 = case splitS x t2 of + (l2 :*: r2) + | l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 -> t1 + | otherwise -> link x l1l2 r1r2 + where !l1l2 = union l1 l2 + !r1r2 = union r1 r2 +#if __GLASGOW_HASKELL__ {-# INLINABLE union #-} #endif -hedgeUnion :: Ord a => MaybeS a -> MaybeS a -> Set a -> Set a -> Set a -hedgeUnion _ _ t1 Tip = t1 -hedgeUnion blo bhi Tip (Bin _ x l r) = link x (filterGt blo l) (filterLt bhi r) -hedgeUnion _ _ t1 (Bin _ x Tip Tip) = insertR x t1 -- According to benchmarks, this special case increases - -- performance up to 30%. It does not help in difference or intersection. -hedgeUnion blo bhi (Bin _ x l r) t2 = link x (hedgeUnion blo bmi l (trim blo bmi t2)) - (hedgeUnion bmi bhi r (trim bmi bhi t2)) - where bmi = JustS x -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE hedgeUnion #-} -#endif - {-------------------------------------------------------------------- Difference --------------------------------------------------------------------} --- | /O(n+m)/. Difference of two sets. --- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/. +-- | /O(m*log(n/m + 1)), m <= n/. Difference of two sets. difference :: Ord a => Set a -> Set a -> Set a difference Tip _ = Tip difference t1 Tip = t1 -difference t1 t2 = hedgeDiff NothingS NothingS t1 t2 -#if __GLASGOW_HASKELL__ >= 700 +difference t1 (Bin _ x l2 r2) = case split x t1 of + (l1, r1) + | size l1l2 + size r1r2 == size t1 -> t1 + | otherwise -> merge l1l2 r1r2 + where !l1l2 = difference l1 l2 + !r1r2 = difference r1 r2 +#if __GLASGOW_HASKELL__ {-# INLINABLE difference #-} #endif -hedgeDiff :: Ord a => MaybeS a -> MaybeS a -> Set a -> Set a -> Set a -hedgeDiff _ _ Tip _ = Tip -hedgeDiff blo bhi (Bin _ x l r) Tip = link x (filterGt blo l) (filterLt bhi r) -hedgeDiff blo bhi t (Bin _ x l r) = merge (hedgeDiff blo bmi (trim blo bmi t) l) - (hedgeDiff bmi bhi (trim bmi bhi t) r) - where bmi = JustS x -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE hedgeDiff #-} -#endif - {-------------------------------------------------------------------- Intersection --------------------------------------------------------------------} --- | /O(n+m)/. The intersection of two sets. The implementation uses an --- efficient /hedge/ algorithm comparable with /hedge-union/. Elements of the --- result come from the first set, so for example +-- | /O(m*log(n/m + 1)), m <= n/. The intersection of two sets. +-- Elements of the result come from the first set, so for example -- -- > import qualified Data.Set as S -- > data AB = A | B deriving Show @@ -675,31 +676,33 @@ intersection :: Ord a => Set a -> Set a -> Set a intersection Tip _ = Tip intersection _ Tip = Tip -intersection t1 t2 = hedgeInt NothingS NothingS t1 t2 -#if __GLASGOW_HASKELL__ >= 700 +intersection t1@(Bin _ x l1 r1) t2 + | b = if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 + then t1 + else link x l1l2 r1r2 + | otherwise = merge l1l2 r1r2 + where + !(l2, b, r2) = splitMember x t2 + !l1l2 = intersection l1 l2 + !r1r2 = intersection r1 r2 +#if __GLASGOW_HASKELL__ {-# INLINABLE intersection #-} #endif -hedgeInt :: Ord a => MaybeS a -> MaybeS a -> Set a -> Set a -> Set a -hedgeInt _ _ _ Tip = Tip -hedgeInt _ _ Tip _ = Tip -hedgeInt blo bhi (Bin _ x l r) t2 = let l' = hedgeInt blo bmi l (trim blo bmi t2) - r' = hedgeInt bmi bhi r (trim bmi bhi t2) - in if x `member` t2 then link x l' r' else merge l' r' - where bmi = JustS x -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE hedgeInt #-} -#endif - {-------------------------------------------------------------------- Filter and partition --------------------------------------------------------------------} -- | /O(n)/. Filter all elements that satisfy the predicate. filter :: (a -> Bool) -> Set a -> Set a filter _ Tip = Tip -filter p (Bin _ x l r) - | p x = link x (filter p l) (filter p r) - | otherwise = merge (filter p l) (filter p r) +filter p t@(Bin _ x l r) + | p x = if l `ptrEq` l' && r `ptrEq` r' + then t + else link x l' r' + | otherwise = merge l' r' + where + !l' = filter p l + !r' = filter p r -- | /O(n)/. Partition the set into two sets, one with all elements that satisfy -- the predicate and one with all elements that don't satisfy the predicate. @@ -708,10 +711,15 @@ partition p0 t0 = toPair $ go p0 t0 where go _ Tip = (Tip :*: Tip) - go p (Bin _ x l r) = case (go p l, go p r) of + go p t@(Bin _ x l r) = case (go p l, go p r) of ((l1 :*: l2), (r1 :*: r2)) - | p x -> link x l1 r1 :*: merge l2 r2 - | otherwise -> merge l1 r1 :*: link x l2 r2 + | p x -> (if l1 `ptrEq` l && r1 `ptrEq` r + then t + else link x l1 r1) :*: merge l2 r2 + | otherwise -> merge l1 r1 :*: + (if l2 `ptrEq` l && r2 `ptrEq` r + then t + else link x l2 r2) {---------------------------------------------------------------------- Map @@ -725,13 +733,13 @@ map :: Ord b => (a->b) -> Set a -> Set b map f = fromList . List.map f . toList -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE map #-} #endif -- | /O(n)/. The -- --- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is monotonic. +-- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is strictly increasing. -- /The precondition is not checked./ -- Semi-formally, we have: -- @@ -774,8 +782,7 @@ foldr' :: (a -> b -> b) -> b -> Set a -> b foldr' f z = go z where - STRICT_1_OF_2(go) - go z' Tip = z' + go !z' Tip = z' go z' (Bin _ x l r) = go (f x (go z' r)) l {-# INLINE foldr' #-} @@ -798,8 +805,7 @@ foldl' :: (a -> b -> a) -> a -> Set b -> a foldl' f z = go z where - STRICT_1_OF_2(go) - go z' Tip = z' + go !z' Tip = z' go z' (Bin _ x l r) = go (f (go z' l) x) r {-# INLINE foldl' #-} @@ -883,8 +889,7 @@ fromList' t0 xs = foldlStrict ins t0 xs where ins t x = insert x t - STRICT_1_OF_3(go) - go _ t [] = t + go !_ t [] = t go _ t [x] = insertMax x t go s l xs@(x : xss) | not_ordered x xss = fromList' l xs | otherwise = case create s xss of @@ -896,8 +901,7 @@ -- If ys is nonempty, the keys in ys are not ordered with respect to tree -- and must be inserted using fromList'. Otherwise the keys have been -- ordered so far. - STRICT_1_OF_2(create) - create _ [] = (Tip, [], []) + create !_ [] = (Tip, [], []) create s xs@(x : xss) | s == 1 = if not_ordered x xss then (Bin 1 x Tip Tip, [], xss) else (Bin 1 x Tip Tip, xss, []) @@ -907,7 +911,7 @@ (l, ys@(y:yss), _) | not_ordered y yss -> (l, [], ys) | otherwise -> case create (s `shiftR` 1) yss of (r, zs, ws) -> (link y l r, zs, ws) -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE fromList #-} #endif @@ -920,25 +924,33 @@ -- | /O(n)/. Build a set from an ascending list in linear time. -- /The precondition (input list is ascending) is not checked./ fromAscList :: Eq a => [a] -> Set a -fromAscList xs - = fromDistinctAscList (combineEq xs) - where - -- [combineEq xs] combines equal elements with [const] in an ordered list [xs] - combineEq xs' - = case xs' of - [] -> [] - [x] -> [x] - (x:xx) -> combineEq' x xx - - combineEq' z [] = [z] - combineEq' z (x:xs') - | z==x = combineEq' z xs' - | otherwise = z:combineEq' x xs' -#if __GLASGOW_HASKELL__ >= 700 +fromAscList xs = fromDistinctAscList (combineEq xs) +#if __GLASGOW_HASKELL__ {-# INLINABLE fromAscList #-} #endif +-- | /O(n)/. Build a set from a descending list in linear time. +-- /The precondition (input list is descending) is not checked./ +fromDescList :: Eq a => [a] -> Set a +fromDescList xs = fromDistinctDescList (combineEq xs) +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromDescList #-} +#endif +-- [combineEq xs] combines equal elements with [const] in an ordered list [xs] +-- +-- TODO: combineEq allocates an intermediate list. It *should* be better to +-- make fromAscListBy and fromDescListBy the fundamental operations, and to +-- implement the rest using those. +combineEq :: Eq a => [a] -> [a] +combineEq [] = [] +combineEq (x : xs) = combineEq' x xs + where + combineEq' z [] = [z] + combineEq' z (y:ys) + | z == y = combineEq' z ys + | otherwise = z : combineEq' y ys + -- | /O(n)/. Build a set from an ascending list of distinct elements in linear time. -- /The precondition (input list is strictly ascending) is not checked./ @@ -948,20 +960,39 @@ fromDistinctAscList [] = Tip fromDistinctAscList (x0 : xs0) = go (1::Int) (Bin 1 x0 Tip Tip) xs0 where - STRICT_1_OF_3(go) - go _ t [] = t + go !_ t [] = t go s l (x : xs) = case create s xs of - (r, ys) -> go (s `shiftL` 1) (link x l r) ys + (r :*: ys) -> go (s `shiftL` 1) (link x l r) ys - STRICT_1_OF_2(create) - create _ [] = (Tip, []) + create !_ [] = (Tip :*: []) create s xs@(x : xs') - | s == 1 = (Bin 1 x Tip Tip, xs') + | s == 1 = (Bin 1 x Tip Tip :*: xs') | otherwise = case create (s `shiftR` 1) xs of - res@(_, []) -> res - (l, y:ys) -> case create (s `shiftR` 1) ys of - (r, zs) -> (link y l r, zs) + res@(_ :*: []) -> res + (l :*: (y:ys)) -> case create (s `shiftR` 1) ys of + (r :*: zs) -> (link y l r :*: zs) +-- | /O(n)/. Build a set from a descending list of distinct elements in linear time. +-- /The precondition (input list is strictly descending) is not checked./ + +-- For some reason, when 'singleton' is used in fromDistinctDescList or in +-- create, it is not inlined, so we inline it manually. +fromDistinctDescList :: [a] -> Set a +fromDistinctDescList [] = Tip +fromDistinctDescList (x0 : xs0) = go (1::Int) (Bin 1 x0 Tip Tip) xs0 + where + go !_ t [] = t + go s r (x : xs) = case create s xs of + (l :*: ys) -> go (s `shiftL` 1) (link x l r) ys + + create !_ [] = (Tip :*: []) + create s xs@(x : xs') + | s == 1 = (Bin 1 x Tip Tip :*: xs') + | otherwise = case create (s `shiftR` 1) xs of + res@(_ :*: []) -> res + (r :*: (y:ys)) -> case create (s `shiftR` 1) ys of + (l :*: zs) -> (link y l r :*: zs) + {-------------------------------------------------------------------- Eq converts the set to a list. In a lazy setting, this actually seems one of the faster methods to compare two trees @@ -1006,7 +1037,7 @@ Typeable/Data --------------------------------------------------------------------} -INSTANCE_TYPEABLE1(Set,setTc,"Set") +INSTANCE_TYPEABLE1(Set) {-------------------------------------------------------------------- NFData @@ -1017,89 +1048,23 @@ rnf (Bin _ y l r) = rnf y `seq` rnf l `seq` rnf r {-------------------------------------------------------------------- - Utility functions that return sub-ranges of the original - tree. Some functions take a `Maybe value` as an argument to - allow comparisons against infinite values. These are called `blow` - (Nothing is -\infty) and `bhigh` (here Nothing is +\infty). - We use MaybeS value, which is a Maybe strict in the Just case. - - [trim blow bhigh t] A tree that is either empty or where [x > blow] - and [x < bhigh] for the value [x] of the root. - [filterGt blow t] A tree where for all values [k]. [k > blow] - [filterLt bhigh t] A tree where for all values [k]. [k < bhigh] - - [split k t] Returns two trees [l] and [r] where all values - in [l] are <[k] and all keys in [r] are >[k]. - [splitMember k t] Just like [split] but also returns whether [k] - was found in the tree. ---------------------------------------------------------------------} - -data MaybeS a = NothingS | JustS !a - -{-------------------------------------------------------------------- - [trim blo bhi t] trims away all subtrees that surely contain no - values between the range [blo] to [bhi]. The returned tree is either - empty or the key of the root is between @blo@ and @bhi@. ---------------------------------------------------------------------} -trim :: Ord a => MaybeS a -> MaybeS a -> Set a -> Set a -trim NothingS NothingS t = t -trim (JustS lx) NothingS t = greater lx t where greater lo (Bin _ x _ r) | x <= lo = greater lo r - greater _ t' = t' -trim NothingS (JustS hx) t = lesser hx t where lesser hi (Bin _ x l _) | x >= hi = lesser hi l - lesser _ t' = t' -trim (JustS lx) (JustS hx) t = middle lx hx t where middle lo hi (Bin _ x _ r) | x <= lo = middle lo hi r - middle lo hi (Bin _ x l _) | x >= hi = middle lo hi l - middle _ _ t' = t' -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE trim #-} -#endif - -{-------------------------------------------------------------------- - [filterGt b t] filter all values >[b] from tree [t] - [filterLt b t] filter all values <[b] from tree [t] ---------------------------------------------------------------------} -filterGt :: Ord a => MaybeS a -> Set a -> Set a -filterGt NothingS t = t -filterGt (JustS b) t = filter' b t - where filter' _ Tip = Tip - filter' b' (Bin _ x l r) = - case compare b' x of LT -> link x (filter' b' l) r - EQ -> r - GT -> filter' b' r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE filterGt #-} -#endif - -filterLt :: Ord a => MaybeS a -> Set a -> Set a -filterLt NothingS t = t -filterLt (JustS b) t = filter' b t - where filter' _ Tip = Tip - filter' b' (Bin _ x l r) = - case compare x b' of LT -> link x l (filter' b' r) - EQ -> l - GT -> filter' b' l -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE filterLt #-} -#endif - -{-------------------------------------------------------------------- Split --------------------------------------------------------------------} -- | /O(log n)/. The expression (@'split' x set@) is a pair @(set1,set2)@ -- where @set1@ comprises the elements of @set@ less than @x@ and @set2@ -- comprises the elements of @set@ greater than @x@. split :: Ord a => a -> Set a -> (Set a,Set a) -split x0 t0 = toPair $ go x0 t0 - where - go _ Tip = (Tip :*: Tip) - go x (Bin _ y l r) +split x t = toPair $ splitS x t +{-# INLINABLE split #-} + +splitS :: Ord a => a -> Set a -> StrictPair (Set a) (Set a) +splitS _ Tip = (Tip :*: Tip) +splitS x (Bin _ y l r) = case compare x y of - LT -> let (lt :*: gt) = go x l in (lt :*: link y gt r) - GT -> let (lt :*: gt) = go x r in (link y l lt :*: gt) + LT -> let (lt :*: gt) = splitS x l in (lt :*: link y gt r) + GT -> let (lt :*: gt) = splitS x r in (link y l lt :*: gt) EQ -> (l :*: r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE split #-} -#endif +{-# INLINABLE splitS #-} -- | /O(log n)/. Performs a 'split' but also returns whether the pivot -- element was found in the original set. @@ -1108,13 +1073,13 @@ splitMember x (Bin _ y l r) = case compare x y of LT -> let (lt, found, gt) = splitMember x l - gt' = link y gt r - in gt' `seq` (lt, found, gt') + !gt' = link y gt r + in (lt, found, gt') GT -> let (lt, found, gt) = splitMember x r - lt' = link y l lt - in lt' `seq` (lt', found, gt) + !lt' = link y l lt + in (lt', found, gt) EQ -> (l, True, r) -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE splitMember #-} #endif @@ -1137,14 +1102,12 @@ findIndex = go 0 where go :: Ord a => Int -> a -> Set a -> Int - STRICT_1_OF_3(go) - STRICT_2_OF_3(go) - go _ _ Tip = error "Set.findIndex: element is not in the set" + go !_ !_ Tip = error "Set.findIndex: element is not in the set" go idx x (Bin _ kx l r) = case compare x kx of LT -> go idx x l GT -> go (idx + size l + 1) x r EQ -> idx + size l -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE findIndex #-} #endif @@ -1162,14 +1125,12 @@ lookupIndex = go 0 where go :: Ord a => Int -> a -> Set a -> Maybe Int - STRICT_1_OF_3(go) - STRICT_2_OF_3(go) - go _ _ Tip = Nothing + go !_ !_ Tip = Nothing go idx x (Bin _ kx l r) = case compare x kx of LT -> go idx x l GT -> go (idx + size l + 1) x r EQ -> Just $! idx + size l -#if __GLASGOW_HASKELL__ >= 700 +#if __GLASGOW_HASKELL__ {-# INLINABLE lookupIndex #-} #endif @@ -1182,8 +1143,7 @@ -- > elemAt 2 (fromList [5,3]) Error: index out of range elemAt :: Int -> Set a -> a -STRICT_1_OF_2(elemAt) -elemAt _ Tip = error "Set.elemAt: index out of range" +elemAt !_ Tip = error "Set.elemAt: index out of range" elemAt i (Bin _ x l r) = case compare i sizeL of LT -> elemAt i l @@ -1202,7 +1162,7 @@ -- > deleteAt (-1) (fromList [5,3]) Error: index out of range deleteAt :: Int -> Set a -> Set a -deleteAt i t = i `seq` +deleteAt !i t = case t of Tip -> error "Set.deleteAt: index out of range" Bin _ x l r -> case compare i sizeL of @@ -1212,7 +1172,118 @@ where sizeL = size l +-- | Take a given number of elements in order, beginning +-- with the smallest ones. +-- +-- @ +-- take n = 'fromDistinctAscList' . 'Prelude.take' n . 'toAscList' +-- @ +take :: Int -> Set a -> Set a +take i m | i >= size m = m +take i0 m0 = go i0 m0 + where + go i !_ | i <= 0 = Tip + go !_ Tip = Tip + go i (Bin _ x l r) = + case compare i sizeL of + LT -> go i l + GT -> link x l (go (i - sizeL - 1) r) + EQ -> l + where sizeL = size l +-- | Drop a given number of elements in order, beginning +-- with the smallest ones. +-- +-- @ +-- drop n = 'fromDistinctAscList' . 'Prelude.drop' n . 'toAscList' +-- @ +drop :: Int -> Set a -> Set a +drop i m | i >= size m = Tip +drop i0 m0 = go i0 m0 + where + go i m | i <= 0 = m + go !_ Tip = Tip + go i (Bin _ x l r) = + case compare i sizeL of + LT -> link x (go i l) r + GT -> go (i - sizeL - 1) r + EQ -> insertMin x r + where sizeL = size l + +-- | /O(log n)/. Split a set at a particular index. +-- +-- @ +-- splitAt !n !xs = ('take' n xs, 'drop' n xs) +-- @ +splitAt :: Int -> Set a -> (Set a, Set a) +splitAt i0 m0 + | i0 >= size m0 = (m0, Tip) + | otherwise = toPair $ go i0 m0 + where + go i m | i <= 0 = Tip :*: m + go !_ Tip = Tip :*: Tip + go i (Bin _ x l r) + = case compare i sizeL of + LT -> case go i l of + ll :*: lr -> ll :*: link x lr r + GT -> case go (i - sizeL - 1) r of + rl :*: rr -> link x l rl :*: rr + EQ -> l :*: insertMin x r + where sizeL = size l + +-- | /O(log n)/. Take while a predicate on the elements holds. +-- The user is responsible for ensuring that for all elements @j@ and @k@ in the set, +-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'. +-- +-- @ +-- takeWhileAntitone p = 'fromDistinctAscList' . 'Data.List.takeWhile' p . 'toList' +-- takeWhileAntitone p = 'filter' p +-- @ + +takeWhileAntitone :: (a -> Bool) -> Set a -> Set a +takeWhileAntitone _ Tip = Tip +takeWhileAntitone p (Bin _ x l r) + | p x = link x l (takeWhileAntitone p r) + | otherwise = takeWhileAntitone p l + +-- | /O(log n)/. Drop while a predicate on the elements holds. +-- The user is responsible for ensuring that for all elements @j@ and @k@ in the set, +-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'. +-- +-- @ +-- dropWhileAntitone p = 'fromDistinctAscList' . 'Data.List.dropWhile' p . 'toList' +-- dropWhileAntitone p = 'filter' (not . p) +-- @ + +dropWhileAntitone :: (a -> Bool) -> Set a -> Set a +dropWhileAntitone _ Tip = Tip +dropWhileAntitone p (Bin _ x l r) + | p x = dropWhileAntitone p r + | otherwise = link x (dropWhileAntitone p l) r + +-- | /O(log n)/. Divide a set at the point where a predicate on the elements stops holding. +-- The user is responsible for ensuring that for all elements @j@ and @k@ in the set, +-- @j \< k ==\> p j \>= p k@. +-- +-- @ +-- spanAntitone p xs = ('takeWhileAntitone' p xs, 'dropWhileAntitone' p xs) +-- spanAntitone p xs = partition p xs +-- @ +-- +-- Note: if @p@ is not actually antitone, then @spanAntitone@ will split the set +-- at some /unspecified/ point where the predicate switches from holding to not +-- holding (where the predicate is seen to hold before the first element and to fail +-- after the last element). + +spanAntitone :: (a -> Bool) -> Set a -> (Set a, Set a) +spanAntitone p0 m = toPair (go p0 m) + where + go _ Tip = Tip :*: Tip + go p (Bin _ x l r) + | p x = let u :*: v = go p r in link x l u :*: v + | otherwise = let u :*: v = go p l in u :*: link x v r + + {-------------------------------------------------------------------- Utility functions that maintain the balance properties of the tree. All constructors assume that all values in [l] < [x] and all values @@ -1233,13 +1304,6 @@ [glue l r] Glues [l] and [r] together. Assumes that [l] and [r] are already balanced with respect to each other. [merge l r] Merges two trees and restores balance. - - Note: in contrast to Adam's paper, we use (<=) comparisons instead - of (<) comparisons in [link], [merge] and [balance]. - Quickcheck (on [difference]) showed that this was necessary in order - to maintain the invariants. It is quite unsatisfactory that I haven't - been able to find out why this is actually the case! Fortunately, it - doesn't hurt to be a bit more conservative. --------------------------------------------------------------------} {-------------------------------------------------------------------- diff --git a/Data/Tree.hs b/Data/Tree.hs --- a/Data/Tree.hs +++ b/Data/Tree.hs @@ -2,6 +2,9 @@ #if __GLASGOW_HASKELL__ {-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} #endif +#if __GLASGOW_HASKELL__ >= 702 +{-# LANGUAGE DeriveGeneric #-} +#endif #if __GLASGOW_HASKELL__ >= 703 {-# LANGUAGE Trustworthy #-} #endif @@ -27,7 +30,7 @@ -- * Two-dimensional drawing drawTree, drawForest, -- * Extraction - flatten, levels, + flatten, levels, foldTree, -- * Building trees unfoldTree, unfoldForest, unfoldTreeM, unfoldForestM, @@ -52,6 +55,11 @@ #ifdef __GLASGOW_HASKELL__ import Data.Data (Data) #endif +#if __GLASGOW_HASKELL__ >= 706 +import GHC.Generics (Generic, Generic1) +#elif __GLASGOW_HASKELL__ >= 702 +import GHC.Generics (Generic) +#endif #if MIN_VERSION_base(4,8,0) import Data.Coerce @@ -63,13 +71,19 @@ subForest :: Forest a -- ^ zero or more child trees } #ifdef __GLASGOW_HASKELL__ +#if __GLASGOW_HASKELL__ >= 706 + deriving (Eq, Read, Show, Data, Generic, Generic1) +#elif __GLASGOW_HASKELL__ >= 702 + deriving (Eq, Read, Show, Data, Generic) +#else deriving (Eq, Read, Show, Data) +#endif #else deriving (Eq, Read, Show) #endif type Forest a = [Tree a] -INSTANCE_TYPEABLE1(Tree,treeTc,"Tree") +INSTANCE_TYPEABLE1(Tree) instance Functor Tree where fmap = fmapTree @@ -120,7 +134,7 @@ drawForest = unlines . map drawTree draw :: Tree String -> [String] -draw (Node x ts0) = x : drawSubTrees ts0 +draw (Node x ts0) = lines x ++ drawSubTrees ts0 where drawSubTrees [] = [] drawSubTrees [t] = @@ -142,6 +156,11 @@ takeWhile (not . null) $ iterate (concatMap subForest) [t] +-- | Catamorphism on trees. +foldTree :: (a -> [b] -> b) -> Tree a -> b +foldTree f = go where + go (Node x ts) = f x (map go ts) + -- | Build a tree from a seed value unfoldTree :: (b -> (a, [b])) -> b -> Tree a unfoldTree f b = let (a, bs) = f b in Node a (unfoldForest f bs) @@ -158,9 +177,7 @@ return (Node a ts) -- | Monadic forest builder, in depth-first order -#ifndef __NHC__ unfoldForestM :: Monad m => (b -> m (a, [b])) -> [b] -> m (Forest a) -#endif unfoldForestM f = Prelude.mapM (unfoldTreeM f) -- | Monadic tree builder, in breadth-first order, diff --git a/Data/Utils/BitQueue.hs b/Data/Utils/BitQueue.hs new file mode 100644 --- /dev/null +++ b/Data/Utils/BitQueue.hs @@ -0,0 +1,148 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} + +#include "containers.h" + +{-# OPTIONS_HADDOCK hide #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Utils.BitQueue +-- Copyright : (c) David Feuer 2016 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. +-- +-- = Description +-- +-- An extremely light-weight, fast, and limited representation of a string of +-- up to (2*WORDSIZE - 2) bits. In fact, there are two representations, +-- misleadingly named bit queue builder and bit queue. The builder supports +-- only `emptyQB`, creating an empty builder, and `snocQB`, enqueueing a bit. +-- The bit queue builder is then turned into a bit queue using `buildQ`, after +-- which bits can be removed one by one using `unconsQ`. If the size limit is +-- exceeded, further operations will silently produce nonsense. +----------------------------------------------------------------------------- + +module Data.Utils.BitQueue + ( BitQueue + , BitQueueB + , emptyQB + , snocQB + , buildQ + , unconsQ + , toListQ + ) where + +#if !MIN_VERSION_base(4,8,0) +import Data.Word (Word) +#endif +import Data.Utils.BitUtil (shiftLL, shiftRL, wordSize) +import Data.Bits ((.|.), (.&.), testBit) +#if MIN_VERSION_base(4,8,0) +import Data.Bits (countTrailingZeros) +#elif MIN_VERSION_base(4,5,0) +import Data.Bits (popCount) +#endif + +#if !MIN_VERSION_base(4,5,0) +-- We could almost certainly improve this fall-back (copied straight from the +-- default definition in Data.Bits), but it hardly seems worth the trouble +-- to speed things up on GHC 7.4 and below. +countTrailingZeros :: Word -> Int +countTrailingZeros x = go 0 + where + go i | i >= wordSize = i + | testBit x i = i + | otherwise = go (i+1) + +#elif !MIN_VERSION_base(4,8,0) +countTrailingZeros :: Word -> Int +countTrailingZeros x = popCount ((x .&. (-x)) - 1) +{-# INLINE countTrailingZeros #-} +#endif + +-- A bit queue builder. We represent a double word using two words +-- because we don't currently have access to proper double words. +data BitQueueB = BQB {-# UNPACK #-} !Word + {-# UNPACK #-} !Word + +newtype BitQueue = BQ BitQueueB deriving Show + +-- Intended for debugging. +instance Show BitQueueB where + show (BQB hi lo) = "BQ"++ + show (map (testBit hi) [(wordSize - 1),(wordSize - 2)..0] + ++ map (testBit lo) [(wordSize - 1),(wordSize - 2)..0]) + +-- | Create an empty bit queue builder. This is represented as a single guard +-- bit in the most significant position. +emptyQB :: BitQueueB +emptyQB = BQB (1 `shiftLL` (wordSize - 1)) 0 +{-# INLINE emptyQB #-} + +-- Shift the double word to the right by one bit. +shiftQBR1 :: BitQueueB -> BitQueueB +shiftQBR1 (BQB hi lo) = BQB hi' lo' where + lo' = (lo `shiftRL` 1) .|. (hi `shiftLL` (wordSize - 1)) + hi' = hi `shiftRL` 1 +{-# INLINE shiftQBR1 #-} + +-- | Enqueue a bit. This works by shifting the queue right one bit, +-- then setting the most significant bit as requested. +{-# INLINE snocQB #-} +snocQB :: BitQueueB -> Bool -> BitQueueB +snocQB bq b = case shiftQBR1 bq of + BQB hi lo -> BQB (hi .|. (fromIntegral (fromEnum b) `shiftLL` (wordSize - 1))) lo + +-- | Convert a bit queue builder to a bit queue. This shifts in a new +-- guard bit on the left, and shifts right until the old guard bit falls +-- off. +{-# INLINE buildQ #-} +buildQ :: BitQueueB -> BitQueue +buildQ (BQB hi 0) = BQ (BQB 0 lo') where + zeros = countTrailingZeros hi + lo' = ((hi `shiftRL` 1) .|. (1 `shiftLL` (wordSize - 1))) `shiftRL` zeros +buildQ (BQB hi lo) = BQ (BQB hi' lo') where + zeros = countTrailingZeros lo + lo1 = (lo `shiftRL` 1) .|. (hi `shiftLL` (wordSize - 1)) + hi1 = (hi `shiftRL` 1) .|. (1 `shiftLL` (wordSize - 1)) + lo' = (lo1 `shiftRL` zeros) .|. (hi1 `shiftLL` (wordSize - zeros)) + hi' = hi1 `shiftRL` zeros + +-- Test if the queue is empty, which occurs when theres +-- nothing left but a guard bit in the least significant +-- place. +nullQ :: BitQueue -> Bool +nullQ (BQ (BQB 0 1)) = True +nullQ _ = False +{-# INLINE nullQ #-} + +-- | Dequeue an element, or discover the queue is empty. +unconsQ :: BitQueue -> Maybe (Bool, BitQueue) +unconsQ q | nullQ q = Nothing +unconsQ (BQ bq@(BQB _ lo)) = Just (hd, BQ tl) + where + !hd = (lo .&. 1) /= 0 + !tl = shiftQBR1 bq +{-# INLINE unconsQ #-} + +-- | Convert a bit queue to a list of bits by unconsing. +-- This is used to test that the queue functions properly. +toListQ :: BitQueue -> [Bool] +toListQ bq = case unconsQ bq of + Nothing -> [] + Just (hd, tl) -> hd : toListQ tl diff --git a/Data/Utils/BitUtil.hs b/Data/Utils/BitUtil.hs --- a/Data/Utils/BitUtil.hs +++ b/Data/Utils/BitUtil.hs @@ -8,6 +8,7 @@ #include "containers.h" +{-# OPTIONS_HADDOCK hide #-} ----------------------------------------------------------------------------- -- | -- Module : Data.Utils.BitUtil @@ -18,15 +19,34 @@ -- Stability : provisional -- Portability : portable ----------------------------------------------------------------------------- +-- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. module Data.Utils.BitUtil ( highestBitMask , shiftLL , shiftRL + , wordSize ) where import Data.Bits ((.|.), xor) +#if MIN_VERSION_base(4,7,0) +import Data.Bits (finiteBitSize) +#else +import Data.Bits (bitSize) +#endif + #if __GLASGOW_HASKELL__ import GHC.Exts (Word(..), Int(..)) import GHC.Prim (uncheckedShiftL#, uncheckedShiftRL#) @@ -61,9 +81,19 @@ --------------------------------------------------------------------} shiftRL (W# x) (I# i) = W# (uncheckedShiftRL# x i) shiftLL (W# x) (I# i) = W# (uncheckedShiftL# x i) +{-# INLINE CONLIKE shiftRL #-} +{-# INLINE CONLIKE shiftLL #-} #else shiftRL x i = shiftR x i shiftLL x i = shiftL x i -#endif {-# INLINE shiftRL #-} {-# INLINE shiftLL #-} +#endif + +{-# INLINE wordSize #-} +wordSize :: Int +#if MIN_VERSION_base(4,7,0) +wordSize = finiteBitSize (0 :: Word) +#else +wordSize = bitSize (0 :: Word) +#endif diff --git a/Data/Utils/PtrEquality.hs b/Data/Utils/PtrEquality.hs new file mode 100644 --- /dev/null +++ b/Data/Utils/PtrEquality.hs @@ -0,0 +1,64 @@ +{-# LANGUAGE CPP #-} +#ifdef __GLASGOW_HASKELL__ +{-# LANGUAGE MagicHash #-} +#endif + +{-# OPTIONS_HADDOCK hide #-} + +-- | Really unsafe pointer equality +-- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. + +module Data.Utils.PtrEquality (ptrEq, hetPtrEq) where + +#ifdef __GLASGOW_HASKELL__ +import GHC.Exts ( reallyUnsafePtrEquality# ) +import Unsafe.Coerce ( unsafeCoerce ) +#if __GLASGOW_HASKELL__ < 707 +import GHC.Exts ( (==#) ) +#else +import GHC.Exts ( isTrue# ) +#endif +#endif + +-- | Checks if two pointers are equal. Yes means yes; +-- no means maybe. The values should be forced to at least +-- WHNF before comparison to get moderately reliable results. +ptrEq :: a -> a -> Bool + +-- | Checks if two pointers are equal, without requiring +-- them to have the same type. The values should be forced +-- to at least WHNF before comparison to get moderately +-- reliable results. +hetPtrEq :: a -> b -> Bool + +#ifdef __GLASGOW_HASKELL__ +#if __GLASGOW_HASKELL__ < 707 +ptrEq x y = reallyUnsafePtrEquality# x y ==# 1# +hetPtrEq x y = unsafeCoerce reallyUnsafePtrEquality# x y ==# 1# +#else +ptrEq x y = isTrue# (reallyUnsafePtrEquality# x y) +hetPtrEq x y = isTrue# (unsafeCoerce reallyUnsafePtrEquality# x y) +#endif + +#else +-- Not GHC +ptrEq _ _ = False +hetPtrEq _ _ = False +#endif + +{-# INLINE ptrEq #-} +{-# INLINE hetPtrEq #-} + +infix 4 `ptrEq` +infix 4 `hetPtrEq` diff --git a/Data/Utils/StrictFold.hs b/Data/Utils/StrictFold.hs --- a/Data/Utils/StrictFold.hs +++ b/Data/Utils/StrictFold.hs @@ -4,6 +4,19 @@ #endif #include "containers.h" +{-# OPTIONS_HADDOCK hide #-} + +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. module Data.Utils.StrictFold (foldlStrict) where diff --git a/Data/Utils/StrictMaybe.hs b/Data/Utils/StrictMaybe.hs new file mode 100644 --- /dev/null +++ b/Data/Utils/StrictMaybe.hs @@ -0,0 +1,43 @@ +{-# LANGUAGE CPP #-} + +#include "containers.h" + +{-# OPTIONS_HADDOCK hide #-} +-- | Strict 'Maybe' +-- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. + +module Data.Utils.StrictMaybe (MaybeS (..), maybeS, toMaybe, toMaybeS) where + +#if !MIN_VERSION_base(4,8,0) +import Data.Foldable (Foldable (..)) +import Data.Monoid (Monoid (..)) +#endif + +data MaybeS a = NothingS | JustS !a + +instance Foldable MaybeS where + foldMap _ NothingS = mempty + foldMap f (JustS a) = f a + +maybeS :: r -> (a -> r) -> MaybeS a -> r +maybeS n _ NothingS = n +maybeS _ j (JustS a) = j a + +toMaybe :: MaybeS a -> Maybe a +toMaybe NothingS = Nothing +toMaybe (JustS a) = Just a + +toMaybeS :: Maybe a -> MaybeS a +toMaybeS Nothing = NothingS +toMaybeS (Just a) = JustS a diff --git a/Data/Utils/StrictPair.hs b/Data/Utils/StrictPair.hs --- a/Data/Utils/StrictPair.hs +++ b/Data/Utils/StrictPair.hs @@ -5,6 +5,22 @@ #include "containers.h" +{-# OPTIONS_HADDOCK hide #-} + +-- | A strict pair +-- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. + module Data.Utils.StrictPair (StrictPair(..), toPair) where -- | Same as regular Haskell pairs, but (x :*: _|_) = (_|_ :*: y) = diff --git a/benchmarks/IntMap.hs b/benchmarks/IntMap.hs --- a/benchmarks/IntMap.hs +++ b/benchmarks/IntMap.hs @@ -1,10 +1,9 @@ {-# LANGUAGE BangPatterns #-} module Main where -import Control.DeepSeq +import Control.DeepSeq (rnf) import Control.Exception (evaluate) -import Control.Monad.Trans (liftIO) -import Criterion.Main +import Criterion.Main (bench, defaultMain, whnf) import Data.List (foldl') import qualified Data.IntMap as M import Data.Maybe (fromMaybe) diff --git a/benchmarks/IntSet.hs b/benchmarks/IntSet.hs --- a/benchmarks/IntSet.hs +++ b/benchmarks/IntSet.hs @@ -2,10 +2,9 @@ module Main where -import Control.DeepSeq +import Control.DeepSeq (rnf) import Control.Exception (evaluate) -import Control.Monad.Trans (liftIO) -import Criterion.Main +import Criterion.Main (bench, defaultMain, whnf) import Data.List (foldl') import qualified Data.IntSet as S diff --git a/benchmarks/LookupGE/IntMap.hs b/benchmarks/LookupGE/IntMap.hs --- a/benchmarks/LookupGE/IntMap.hs +++ b/benchmarks/LookupGE/IntMap.hs @@ -1,11 +1,9 @@ {-# LANGUAGE BangPatterns #-} module Main where -import Control.DeepSeq +import Control.DeepSeq (rnf) import Control.Exception (evaluate) -import Control.Monad.Trans (liftIO) -import Criterion.Config -import Criterion.Main +import Criterion.Main (bench, defaultMain, nf) import Data.List (foldl') import qualified Data.IntMap as M import qualified LookupGE_IntMap as M @@ -14,10 +12,8 @@ main :: IO () main = do - defaultMainWith - defaultConfig - (liftIO . evaluate $ rnf [m_even, m_odd, m_large]) - [b f | b <- benches, f <- funs1] + evaluate $ rnf [m_even, m_odd, m_large] + defaultMain [b f | b <- benches, f <- funs1] where m_even = M.fromAscList elems_even :: M.IntMap Int m_odd = M.fromAscList elems_odd :: M.IntMap Int diff --git a/benchmarks/LookupGE/LookupGE_IntMap.hs b/benchmarks/LookupGE/LookupGE_IntMap.hs --- a/benchmarks/LookupGE/LookupGE_IntMap.hs +++ b/benchmarks/LookupGE/LookupGE_IntMap.hs @@ -3,9 +3,6 @@ import Prelude hiding (null) import Data.IntMap.Base -#ifdef TESTING -import Test.QuickCheck -#endif lookupGE1 :: Key -> IntMap a -> Maybe (Key,a) lookupGE1 k m = diff --git a/benchmarks/LookupGE/LookupGE_Map.hs b/benchmarks/LookupGE/LookupGE_Map.hs --- a/benchmarks/LookupGE/LookupGE_Map.hs +++ b/benchmarks/LookupGE/LookupGE_Map.hs @@ -2,9 +2,6 @@ module LookupGE_Map where import Data.Map.Base -#ifdef TESTING -import Test.QuickCheck -#endif lookupGE1 :: Ord k => k -> Map k a -> Maybe (k,a) lookupGE1 k m = diff --git a/benchmarks/LookupGE/Map.hs b/benchmarks/LookupGE/Map.hs --- a/benchmarks/LookupGE/Map.hs +++ b/benchmarks/LookupGE/Map.hs @@ -1,11 +1,9 @@ {-# LANGUAGE BangPatterns #-} module Main where -import Control.DeepSeq +import Control.DeepSeq (rnf) import Control.Exception (evaluate) -import Control.Monad.Trans (liftIO) -import Criterion.Config -import Criterion.Main +import Criterion.Main (defaultMain, bench, nf) import Data.List (foldl') import qualified Data.Map as M import qualified LookupGE_Map as M @@ -14,10 +12,8 @@ main :: IO () main = do - defaultMainWith - defaultConfig - (liftIO . evaluate $ rnf [m_even, m_odd, m_large]) - [b f | b <- benches, f <- funs1] + evaluate $ rnf [m_even, m_odd, m_large] + defaultMain [b f | b <- benches, f <- funs1] where m_even = M.fromAscList elems_even :: M.Map Int Int m_odd = M.fromAscList elems_odd :: M.Map Int Int diff --git a/benchmarks/Makefile b/benchmarks/Makefile --- a/benchmarks/Makefile +++ b/benchmarks/Makefile @@ -1,7 +1,7 @@ all: bench-%: %.hs force - ghc -O2 -DTESTING $< -I../include -i../$(TOP) -o $@ -outputdir tmp -rtsopts + ghc -O2 -DTESTING $< -I$(TOP)../include -i$(TOP).. -o $@ -outputdir tmp -rtsopts .PRECIOUS: bench-% diff --git a/benchmarks/Map.hs b/benchmarks/Map.hs --- a/benchmarks/Map.hs +++ b/benchmarks/Map.hs @@ -1,13 +1,20 @@ +{-# LANGUAGE CPP #-} {-# LANGUAGE BangPatterns #-} module Main where -import Control.DeepSeq +import Control.Applicative (Const(Const, getConst), pure) +import Control.DeepSeq (rnf) import Control.Exception (evaluate) -import Control.Monad.Trans (liftIO) -import Criterion.Main +import Criterion.Main (bench, defaultMain, whnf, nf) +import Data.Functor.Identity (Identity(..)) import Data.List (foldl') import qualified Data.Map as M +import Data.Map (alterF) import Data.Maybe (fromMaybe) +import Data.Functor ((<$)) +#if __GLASGOW_HASKELL__ >= 708 +import Data.Coerce +#endif import Prelude hiding (lookup) main = do @@ -18,8 +25,38 @@ defaultMain [ bench "lookup absent" $ whnf (lookup evens) m_odd , bench "lookup present" $ whnf (lookup evens) m_even + , bench "map" $ whnf (M.map (+ 1)) m + , bench "map really" $ nf (M.map (+ 2)) m + , bench "<$" $ whnf ((1 :: Int) <$) m + , bench "<$ really" $ nf ((2 :: Int) <$) m + , bench "alterF lookup absent" $ whnf (atLookup evens) m_odd + , bench "alterF lookup present" $ whnf (atLookup evens) m_even + , bench "alterF no rules lookup absent" $ whnf (atLookupNoRules evens) m_odd + , bench "alterF no rules lookup present" $ whnf (atLookupNoRules evens) m_even , bench "insert absent" $ whnf (ins elems_even) m_odd , bench "insert present" $ whnf (ins elems_even) m_even + , bench "alterF insert absent" $ whnf (atIns elems_even) m_odd + , bench "alterF insert present" $ whnf (atIns elems_even) m_even + , bench "alterF no rules insert absent" $ whnf (atInsNoRules elems_even) m_odd + , bench "alterF no rules insert present" $ whnf (atInsNoRules elems_even) m_even + , bench "delete absent" $ whnf (del evens) m_odd + , bench "delete present" $ whnf (del evens) m + , bench "alterF delete absent" $ whnf (atDel evens) m_odd + , bench "alterF delete present" $ whnf (atDel evens) m + , bench "alterF no rules delete absent" $ whnf (atDelNoRules evens) m_odd + , bench "alterF no rules delete present" $ whnf (atDelNoRules evens) m + , bench "alter absent" $ whnf (alt id evens) m_odd + , bench "alter insert" $ whnf (alt (const (Just 1)) evens) m_odd + , bench "alter update" $ whnf (alt id evens) m_even + , bench "alter delete" $ whnf (alt (const Nothing) evens) m + , bench "alterF alter absent" $ whnf (atAlt id evens) m_odd + , bench "alterF alter insert" $ whnf (atAlt (const (Just 1)) evens) m_odd + , bench "alterF alter update" $ whnf (atAlt id evens) m_even + , bench "alterF alter delete" $ whnf (atAlt (const Nothing) evens) m + , bench "alterF no rules alter absent" $ whnf (atAltNoRules id evens) m_odd + , bench "alterF no rules alter insert" $ whnf (atAltNoRules (const (Just 1)) evens) m_odd + , bench "alterF no rules alter update" $ whnf (atAltNoRules id evens) m_even + , bench "alterF no rules alter delete" $ whnf (atAltNoRules (const Nothing) evens) m , bench "insertWith absent" $ whnf (insWith elems_even) m_odd , bench "insertWith present" $ whnf (insWith elems_even) m_even , bench "insertWith' absent" $ whnf (insWith' elems_even) m_odd @@ -32,23 +69,16 @@ , bench "insertLookupWithKey present" $ whnf (insLookupWithKey elems_even) m_even , bench "insertLookupWithKey' absent" $ whnf (insLookupWithKey' elems_even) m_odd , bench "insertLookupWithKey' present" $ whnf (insLookupWithKey' elems_even) m_even - , bench "map" $ whnf (M.map (+ 1)) m , bench "mapWithKey" $ whnf (M.mapWithKey (+)) m , bench "foldlWithKey" $ whnf (ins elems) m -- , bench "foldlWithKey'" $ whnf (M.foldlWithKey' sum 0) m , bench "foldrWithKey" $ whnf (M.foldrWithKey consPair []) m - , bench "delete absent" $ whnf (del evens) m_odd - , bench "delete present" $ whnf (del evens) m , bench "update absent" $ whnf (upd Just evens) m_odd , bench "update present" $ whnf (upd Just evens) m_even , bench "update delete" $ whnf (upd (const Nothing) evens) m , bench "updateLookupWithKey absent" $ whnf (upd' Just evens) m_odd , bench "updateLookupWithKey present" $ whnf (upd' Just evens) m_even , bench "updateLookupWithKey delete" $ whnf (upd' (const Nothing) evens) m - , bench "alter absent" $ whnf (alt id evens) m_odd - , bench "alter insert" $ whnf (alt (const (Just 1)) evens) m_odd - , bench "alter update" $ whnf (alt id evens) m_even - , bench "alter delete" $ whnf (alt (const Nothing) evens) m , bench "mapMaybe" $ whnf (M.mapMaybe maybeDel) m , bench "mapMaybeWithKey" $ whnf (M.mapMaybeWithKey (const maybeDel)) m , bench "lookupIndex" $ whnf (lookupIndex keys) m @@ -80,12 +110,36 @@ lookup :: [Int] -> M.Map Int Int -> Int lookup xs m = foldl' (\n k -> fromMaybe n (M.lookup k m)) 0 xs +atLookup :: [Int] -> M.Map Int Int -> Int +atLookup xs m = foldl' (\n k -> fromMaybe n (getConst (alterF Const k m))) 0 xs + +newtype Consty a b = Consty { getConsty :: a } +instance Functor (Consty a) where + fmap _ (Consty a) = Consty a + +atLookupNoRules :: [Int] -> M.Map Int Int -> Int +atLookupNoRules xs m = foldl' (\n k -> fromMaybe n (getConsty (alterF Consty k m))) 0 xs + lookupIndex :: [Int] -> M.Map Int Int -> Int lookupIndex xs m = foldl' (\n k -> fromMaybe n (M.lookupIndex k m)) 0 xs ins :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int ins xs m = foldl' (\m (k, v) -> M.insert k v m) m xs +atIns :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int +atIns xs m = foldl' (\m (k, v) -> runIdentity (alterF (\_ -> Identity (Just v)) k m)) m xs + +newtype Ident a = Ident { runIdent :: a } +instance Functor Ident where +#if __GLASGOW_HASKELL__ >= 708 + fmap = coerce +#else + fmap f (Ident a) = Ident (f a) +#endif + +atInsNoRules :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int +atInsNoRules xs m = foldl' (\m (k, v) -> runIdent (alterF (\_ -> Ident (Just v)) k m)) m xs + insWith :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int insWith xs m = foldl' (\m (k, v) -> M.insertWith (+) k v m) m xs @@ -115,6 +169,12 @@ del :: [Int] -> M.Map Int Int -> M.Map Int Int del xs m = foldl' (\m k -> M.delete k m) m xs +atDel :: [Int] -> M.Map Int Int -> M.Map Int Int +atDel xs m = foldl' (\m k -> runIdentity (alterF (\_ -> Identity Nothing) k m)) m xs + +atDelNoRules :: [Int] -> M.Map Int Int -> M.Map Int Int +atDelNoRules xs m = foldl' (\m k -> runIdent (alterF (\_ -> Ident Nothing) k m)) m xs + upd :: (Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int upd f xs m = foldl' (\m k -> M.update f k m) m xs @@ -123,6 +183,12 @@ alt :: (Maybe Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int alt f xs m = foldl' (\m k -> M.alter f k m) m xs + +atAlt :: (Maybe Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int +atAlt f xs m = foldl' (\m k -> runIdentity (alterF (Identity . f) k m)) m xs + +atAltNoRules :: (Maybe Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int +atAltNoRules f xs m = foldl' (\m k -> runIdent (alterF (Ident . f) k m)) m xs maybeDel :: Int -> Maybe Int maybeDel n | n `mod` 3 == 0 = Nothing diff --git a/benchmarks/Sequence.hs b/benchmarks/Sequence.hs --- a/benchmarks/Sequence.hs +++ b/benchmarks/Sequence.hs @@ -1,14 +1,15 @@ --- > ghc -DTESTING --make -O2 -fforce-recomp -i.. Sequence.hs module Main where import Control.Applicative -import Control.DeepSeq +import Control.DeepSeq (rnf) import Control.Exception (evaluate) -import Criterion.Main -import Data.List (foldl') +import Control.Monad.Trans.State.Strict +import Criterion.Main (bench, bgroup, defaultMain, nf) +import Data.Foldable (foldl', foldr') import qualified Data.Sequence as S import qualified Data.Foldable -import System.Random +import Data.Traversable (traverse) +import System.Random (mkStdGen, randoms) main = do let s10 = S.fromList [1..10] :: S.Seq Int @@ -34,6 +35,66 @@ , bench "100" $ nf (shuffle r100) s100 , bench "1000" $ nf (shuffle r1000) s1000 ] + , bgroup "fromList" + [ bench "10" $ nf S.fromList [(0 :: Int)..9] + , bench "100" $ nf S.fromList [(0 :: Int)..99] + , bench "1000" $ nf S.fromList [(0 :: Int)..999] + , bench "10000" $ nf S.fromList [(0 :: Int)..9999] + , bench "100000" $ nf S.fromList [(0 :: Int)..99999] + ] + , bgroup "partition" + [ bench "10" $ nf (S.partition even) s10 + , bench "100" $ nf (S.partition even) s100 + , bench "1000" $ nf (S.partition even) s1000 + , bench "10000" $ nf (S.partition even) s10000 + ] + , bgroup "foldl'" + [ bench "10" $ nf (foldl' (+) 0) s10 + , bench "100" $ nf (foldl' (+) 0) s100 + , bench "1000" $ nf (foldl' (+) 0) s1000 + , bench "10000" $ nf (foldl' (+) 0) s10000 + ] + , bgroup "foldr'" + [ bench "10" $ nf (foldr' (+) 0) s10 + , bench "100" $ nf (foldr' (+) 0) s100 + , bench "1000" $ nf (foldr' (+) 0) s1000 + , bench "10000" $ nf (foldr' (+) 0) s10000 + ] + , bgroup "update" + [ bench "10" $ nf (updatePoints r10 10) s10 + , bench "100" $ nf (updatePoints r100 10) s100 + , bench "1000" $ nf (updatePoints r1000 10) s1000 + ] + , bgroup "adjust" + [ bench "10" $ nf (adjustPoints r10 (+10)) s10 + , bench "100" $ nf (adjustPoints r100 (+10)) s100 + , bench "1000" $ nf (adjustPoints r1000 (+10)) s1000 + ] + , bgroup "deleteAt" + [ bench "10" $ nf (deleteAtPoints r10) s10 + , bench "100" $ nf (deleteAtPoints r100) s100 + , bench "1000" $ nf (deleteAtPoints r1000) s1000 + ] + , bgroup "insertAt" + [ bench "10" $ nf (insertAtPoints r10 10) s10 + , bench "100" $ nf (insertAtPoints r100 10) s100 + , bench "1000" $ nf (insertAtPoints r1000 10) s1000 + ] + , bgroup "traverseWithIndex/State" + [ bench "10" $ nf multiplyDown s10 + , bench "100" $ nf multiplyDown s100 + , bench "1000" $ nf multiplyDown s1000 + ] + , bgroup "traverse/State" + [ bench "10" $ nf multiplyUp s10 + , bench "100" $ nf multiplyUp s100 + , bench "1000" $ nf multiplyUp s1000 + ] + , bgroup "replicateA/State" + [ bench "10" $ nf stateReplicate 10 + , bench "100" $ nf stateReplicate 100 + , bench "1000" $ nf stateReplicate 1000 + ] , bgroup "zip" [ bench "ix10000/5000" $ nf (\(xs,ys) -> S.zip xs ys `S.index` 5000) (s10000, u10000) , bench "nf100" $ nf (uncurry S.zip) (s100, u100) @@ -69,6 +130,51 @@ ] ] +{- +-- This is around 4.6 times as slow as insertAt +fakeInsertAt :: Int -> a -> S.Seq a -> S.Seq a +fakeInsertAt i x xs = case S.splitAt i xs of + (before, after) -> before S.>< x S.<| after +-} + +adjustPoints :: [Int] -> (a -> a) -> S.Seq a -> S.Seq a +adjustPoints points f xs = + foldl' (\acc k -> S.adjust f k acc) xs points + +insertAtPoints :: [Int] -> a -> S.Seq a -> S.Seq a +insertAtPoints points x xs = + foldl' (\acc k -> S.insertAt k x acc) xs points + +updatePoints :: [Int] -> a -> S.Seq a -> S.Seq a +updatePoints points x xs = + foldl' (\acc k -> S.update k x acc) xs points + +{- +-- For comparison. Using the old implementation of update, +-- which this simulates, can cause thunks to build up in the leaves. +fakeupdatePoints :: [Int] -> a -> S.Seq a -> S.Seq a +fakeupdatePoints points x xs = + foldl' (\acc k -> S.adjust (const x) k acc) xs points +-} + +deleteAtPoints :: [Int] -> S.Seq a -> S.Seq a +deleteAtPoints points xs = + foldl' (\acc k -> S.deleteAt k acc) xs points + +{- +fakedeleteAtPoints :: [Int] -> S.Seq a -> S.Seq a +fakedeleteAtPoints points xs = + foldl' (\acc k -> fakeDeleteAt k acc) xs points + +-- For comparison with deleteAt. deleteAt is several +-- times faster for long sequences. +fakeDeleteAt :: Int -> S.Seq a -> S.Seq a +fakeDeleteAt i xs + | 0 < i && i < S.length xs = case S.splitAt i xs of + (before, after) -> before S.>< S.drop 1 after + | otherwise = xs +-} + -- splitAt+append: repeatedly cut the sequence at a random point -- and rejoin the pieces in the opposite order. -- Finally getting the middle element forces the whole spine. @@ -76,3 +182,23 @@ shuffle ps s = case S.viewl (S.drop (S.length s `div` 2) (foldl' cut s ps)) of x S.:< _ -> x where cut xs p = let (front, back) = S.splitAt p xs in back S.>< front + +stateReplicate :: Int -> S.Seq Char +stateReplicate n = flip evalState 0 . S.replicateA n $ do + old <- get + if old > (10 :: Int) then put 0 else put (old + 1) + return $ toEnum old + +multiplyUp :: S.Seq Int -> S.Seq Int +multiplyUp = flip evalState 0 . traverse go where + go x = do + s <- get + put (s + 1) + return (s * x) + +multiplyDown :: S.Seq Int -> S.Seq Int +multiplyDown = flip evalState 0 . S.traverseWithIndex go where + go i x = do + s <- get + put (s - 1) + return (s * i * x) diff --git a/benchmarks/Set.hs b/benchmarks/Set.hs --- a/benchmarks/Set.hs +++ b/benchmarks/Set.hs @@ -1,12 +1,10 @@ {-# LANGUAGE BangPatterns #-} --- > ghc -DTESTING --make -O2 -fforce-recomp -i.. Set.hs module Main where -import Control.DeepSeq +import Control.DeepSeq (rnf) import Control.Exception (evaluate) -import Control.Monad.Trans (liftIO) -import Criterion.Main +import Criterion.Main (bench, defaultMain, whnf) import Data.List (foldl') import qualified Data.Set as S diff --git a/benchmarks/SetOperations/Makefile b/benchmarks/SetOperations/Makefile --- a/benchmarks/SetOperations/Makefile +++ b/benchmarks/SetOperations/Makefile @@ -1,3 +1,3 @@ -TOP = .. +TOP = ../ include ../Makefile diff --git a/benchmarks/SetOperations/SetOperations.hs b/benchmarks/SetOperations/SetOperations.hs --- a/benchmarks/SetOperations/SetOperations.hs +++ b/benchmarks/SetOperations/SetOperations.hs @@ -2,7 +2,7 @@ module SetOperations (benchmark) where -import Criterion.Main +import Criterion.Main (bench, defaultMain, whnf) import Data.List (partition) benchmark :: ([Int] -> container) -> Bool -> [(String, container -> container -> container)] -> IO () @@ -19,7 +19,7 @@ | (mode_str, (left, right)) <- [ ("disj_nn", disj_nn), ("disj_ns", disj_ns), ("disj_nt", disj_nt) , ("common_nn", common_nn), ("common_ns", common_ns), ("common_nt", common_nt) , ("mix_nn", mix_nn), ("mix_ns", mix_ns), ("mix_nt", mix_nt) - , ("block_nn", block_nn), ("block_sn", block_ns) + , ("block_nn", block_nn), ("block_ns", block_ns) ] , (mode_str, left, right) <- replicate 2 (mode_str, left, right) ++ @@ -35,11 +35,11 @@ !common_nn = seqPair $ (all_n, fromList [2,4..n]) !common_ns = seqPair $ (all_n, fromList [0,1+n`div`s..n]) !common_nt = seqPair $ (all_n, fromList [0,1+n`div`t..n]) - !mix_nn = seqPair $ fromLists $ partition ((== 0) . (`mod` 2)) [1..n+n] - !mix_ns = seqPair $ fromLists $ partition ((== 0) . (`mod` (1 + n`div`s))) [1..s+n] - !mix_nt = seqPair $ fromLists $ partition ((== 0) . (`mod` (1 + n`div`t))) [1..t+n] - !block_nn = seqPair $ fromLists $ partition ((< t) . (`mod` (t * 2))) [1..n+n] - !block_ns = seqPair $ fromLists $ partition ((< t) . (`mod` (t * (1 + n`div`s)))) [1..s+n] + !mix_nn = seqPair $ fromLists $ partition ((/= 0) . (`mod` 2)) [1..n+n] + !mix_ns = seqPair $ fromLists $ partition ((/= 0) . (`mod` (1 + n`div`s))) [1..s+n] + !mix_nt = seqPair $ fromLists $ partition ((/= 0) . (`mod` (1 + n`div`t))) [1..t+n] + !block_nn = seqPair $ fromLists $ partition ((>= t) . (`mod` (t * 2))) [1..n+n] + !block_ns = seqPair $ fromLists $ partition ((>= t) . (`mod` (t * (1 + n`div`s)))) [1..s+n] fromLists (xs, ys) = (fromList xs, fromList ys) seqPair pair@(xs, ys) = xs `seq` ys `seq` pair diff --git a/changelog.md b/changelog.md --- a/changelog.md +++ b/changelog.md @@ -1,12 +1,120 @@ # Changelog for [`containers` package](http://github.com/haskell/containers) +## 0.5.8.1 + +### General package changes + + * Remove all attempts to support nhc98 and any versions of GHC + before 7.0. + + * Integrate benchmarks with Cabal. (Thanks, Gabriel Gonzalez!) + + * Make Cabal report required extensions properly, and stop using + default extensions. Note that we do *not* report extensions conditionally enabled + based on GHC version, as doing so would lead to a maintenance nightmare + with no obvious benefits. + + * Use `BangPatterns` throughout to reduce noise. This extension + is now *required* to compile `containers`. + + * Improve QuickCheck properties taking arbitrary functions by using + `Test.QuickCheck.Function.Fun` instead of evil `Show` instances + for functions. + + * Expose several internal modules through Cabal (as requested by + Edward Kmett). These remain completely unsupported. + +### New exports and instances + + * Add `alterF`, `restrictKeys`, and `withoutKeys` to `Data.Map` + and `Data.IntMap`. + + * Add `take`, `drop`, `splitAt`, `takeWhileAntitone`, `dropWhileAntitone`, + and `spanAntitone` for `Data.Map` and `Data.Set`. Thanks to Cale Gibbard + for suggesting these. + + * Add `merge`, `mergeA`, and associated merge tactics for `Data.Map`. + Many thanks to Cale Gibbard, Ryan Trinkle, and Dan Doel for + inspiring the merge idea and helping refine the interface. + + * Add `fromDescList`, `fromDescListWith`, `fromDescListWithKey`, + and `fromDistinctDescList` to `Data.Map`. + + * Add `fromDescList` and `fromDistinctDescList` to `Data.Set`. + + * Add `Empty`, `:<|`, and `:|>` pattern synonyms for `Data.Sequence`. + + * Add `adjust'`, `(!?)`, `lookup`, `chunksOf`, `cycleTaking`, `insertAt`, `deleteAt`, `intersperse`, + `foldMapWithIndex`, and `traverseWithIndex` for `Data.Sequence`. + + * Derive `Generic` and `Generic1` for `Data.Tree.Tree`, `Data.Sequence.ViewL`, + and `Data.Sequence.ViewR`. + + * Add `foldTree` for `Data.Tree`. (Thanks, Daniel Wagner!) + +### Semantic changes + + * Make `Data.Sequence.splitAt` strict in its arguments. Previously, + it returned a lazy pair. + + * Fix completely erroneous definition of `length` for `Data.Sequence.ViewR`. + + * Make `Data.Map.Strict.traverseWithKey` force result values before + installing them in the new map. + + * Make `drawTree` handle newlines better. (Thanks, recursion-ninja!) + +### Deprecations + + * All functions in `Data.Map` proper that have been documented as deprecated since + version 0.5 or before now have `DEPRECATED` pragmas and will actually be + removed after another cycle or two. + + * Tree printing functions in `Data.Map` intended for library debugging are now + deprecated. They will continue to be available for the foreseeable future in + an internal module. + +### Performance changes + + * Substantially speed up `splitAt`, `zipWith`, `take`, `drop`, + `fromList`, `partition`, `foldl'`, and `foldr'` for `Data.Sequence`. + Special thanks to Lennart Spitzner for digging into the performance + problems with previous versions of `fromList` and finding a way to + make it really fast. Slightly optimize `replicateA`. Stop `traverse` + from performing many unnecessary `fmap` operations. + + * Most operations in `Data.Sequence` advertised as taking logarithmic + time (including `><` and `adjust`) now use their full allotted time + to avoid potentially building up chains of thunks in the tree. In general, + the only remaining operations that avoid doing more than they + really need are the particular bulk creation and transformation functions + that really benefit from the extra laziness. There are some situations + where this change may slow programs down, but I think having more + predictable and usually better performance more than makes up for that. + + * Add rewrite rules to fuse `fmap` with `reverse` for `Data.Sequence`. + + * Switch from *hedge* algorithms to *divide-and-conquer* algorithms + for union, intersection, difference, and merge in both `Data.Map` + and `Data.Set`. These algorithms are simpler, are known to be + asymptotically optimal, and are faster according to our benchmarks. + + * Speed up `adjust` for `Data.Map`. Allow `map` to inline, and + define a custom `(<$)`. This considerably improves mapping with + a constant function. + + * Remove non-essential laziness in `Data.Map.Lazy` implementation. + + * Speed up deletion and alteration functions for `Data.IntMap`. + + ## 0.5.7.1 *Dec 2015* * Planned to bundle with GHC 8.0.1. * Add `IsString` instance to `Data.Sequence`. - * Define `Semigroup` instances for ``Data.Map`, `Data.Set`, `Data.IntMap`, + * Define `Semigroup` instances for `Data.Map`, `Data.Set`, `Data.IntMap`, `Data.IntSet` and `Data.Sequence`. ## 0.5.6.2 *Dec 2014* diff --git a/containers.cabal b/containers.cabal --- a/containers.cabal +++ b/containers.cabal @@ -1,8 +1,8 @@ name: containers -version: 0.5.7.1 +version: 0.5.8.1 license: BSD3 license-file: LICENSE -maintainer: fox@ucw.cz +maintainer: libraries@haskell.org bug-reports: https://github.com/haskell/containers/issues synopsis: Assorted concrete container types category: Data Structures @@ -32,40 +32,174 @@ location: http://github.com/haskell/containers.git Library - build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.5 - if impl(ghc>=6.10) + build-depends: base >= 4.3 && < 5, array, deepseq >= 1.2 && < 1.5 + if impl(ghc) build-depends: ghc-prim ghc-options: -O2 -Wall + other-extensions: CPP, BangPatterns + exposed-modules: Data.IntMap Data.IntMap.Lazy Data.IntMap.Strict + Data.IntMap.Base + Data.IntSet.Base Data.IntSet Data.Map Data.Map.Lazy + Data.Map.Lazy.Merge + Data.Map.Strict.Internal Data.Map.Strict - Data.Set - if !impl(nhc98) - exposed-modules: - Data.Graph - Data.Sequence - Data.Tree - other-modules: - Data.IntMap.Base - Data.IntSet.Base + Data.Map.Strict.Merge Data.Map.Base Data.Set.Base + Data.Set + Data.Graph + Data.Sequence + Data.Sequence.Base + Data.Tree Data.Utils.BitUtil + Data.Utils.BitQueue Data.Utils.StrictFold Data.Utils.StrictPair + Data.Utils.StrictMaybe + Data.Utils.PtrEquality include-dirs: include - if impl(ghc<7.0) - extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types +----------------------------- +-- B E N C H M A R K I N G -- +----------------------------- +benchmark intmap-benchmarks + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks + main-is: IntMap.hs + ghc-options: -O2 + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2, + deepseq >= 1.1.0.0 && < 1.5 + +benchmark intset-benchmarks + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks + main-is: IntSet.hs + ghc-options: -O2 + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2, + deepseq >= 1.1.0.0 && < 1.5 + +benchmark map-benchmarks + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks + main-is: Map.hs + ghc-options: -O2 + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2, + deepseq >= 1.1.0.0 && < 1.5 + +benchmark sequence-benchmarks + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks + main-is: Sequence.hs + ghc-options: -O2 + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2, + deepseq >= 1.1.0.0 && < 1.5, + random < 1.2, + transformers + +benchmark set-benchmarks + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks + main-is: Set.hs + ghc-options: -O2 + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2, + deepseq >= 1.1.0.0 && < 1.5 + +benchmark set-operations-intmap + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks/SetOperations + main-is: SetOperations-IntMap.hs + ghc-options: -O2 + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2 + +benchmark set-operations-intset + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks/SetOperations + main-is: SetOperations-IntSet.hs + ghc-options: -O2 + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2 + +benchmark set-operations-map + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks/SetOperations + main-is: SetOperations-Map.hs + ghc-options: -O2 + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2 + +benchmark set-operations-set + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks/SetOperations + main-is: SetOperations-Set.hs + ghc-options: -O2 + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2 + +benchmark lookupge-intmap + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks/LookupGE, . + main-is: IntMap.hs + ghc-options: -O2 + cpp-options: -DTESTING + other-modules: + Data.IntMap.Base + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2, + deepseq >= 1.1.0.0 && < 1.5, + ghc-prim + +benchmark lookupge-map + type: exitcode-stdio-1.0 + hs-source-dirs: benchmarks/LookupGE, . + main-is: Map.hs + ghc-options: -O2 + cpp-options: -DTESTING + other-modules: + Data.Map.Base + build-depends: + base >= 4.2 && < 5, + containers, + criterion >= 0.4.0 && < 1.2, + deepseq >= 1.1.0.0 && < 1.5, + ghc-prim + ------------------- -- T E S T I N G -- ------------------- @@ -73,28 +207,24 @@ -- Every test-suite contains the build-depends and options of the library, -- plus the testing stuff. --- Because the test-suites cannot contain conditionals in GHC 7.0, the extensions --- are switched on for every compiler to allow GHC < 7.0 to compile the tests --- (because GHC < 7.0 cannot handle conditional LANGUAGE pragmas). --- When testing with GHC < 7.0 is not needed, the extensions should be removed. - Test-suite map-lazy-properties hs-source-dirs: tests, . main-is: map-properties.hs type: exitcode-stdio-1.0 cpp-options: -DTESTING - build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim + build-depends: base >= 4.3 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim ghc-options: -O2 + other-extensions: CPP, BangPatterns include-dirs: include - extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types build-depends: HUnit, QuickCheck, test-framework, test-framework-hunit, - test-framework-quickcheck2 + test-framework-quickcheck2, + transformers Test-suite map-strict-properties hs-source-dirs: tests, . @@ -102,16 +232,33 @@ type: exitcode-stdio-1.0 cpp-options: -DTESTING -DSTRICT - build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim + build-depends: base >= 4.3 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim ghc-options: -O2 + other-extensions: CPP, BangPatterns include-dirs: include - extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types build-depends: HUnit, QuickCheck, test-framework, test-framework-hunit, + test-framework-quickcheck2, + transformers + +Test-suite bitqueue-properties + hs-source-dirs: tests, . + main-is: bitqueue-properties.hs + type: exitcode-stdio-1.0 + cpp-options: -DTESTING + + build-depends: base >= 4.3 && < 5, ghc-prim + ghc-options: -O2 + other-extensions: CPP, BangPatterns + include-dirs: include + + build-depends: + QuickCheck, + test-framework, test-framework-quickcheck2 Test-suite set-properties @@ -120,17 +267,18 @@ type: exitcode-stdio-1.0 cpp-options: -DTESTING - build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim + build-depends: base >= 4.3 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim ghc-options: -O2 + other-extensions: CPP, BangPatterns include-dirs: include - extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types build-depends: HUnit, QuickCheck, test-framework, test-framework-hunit, - test-framework-quickcheck2 + test-framework-quickcheck2, + transformers Test-suite intmap-lazy-properties hs-source-dirs: tests, . @@ -138,10 +286,10 @@ type: exitcode-stdio-1.0 cpp-options: -DTESTING - build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim + build-depends: base >= 4.3 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim ghc-options: -O2 + other-extensions: CPP, BangPatterns include-dirs: include - extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types build-depends: HUnit, @@ -156,10 +304,10 @@ type: exitcode-stdio-1.0 cpp-options: -DTESTING -DSTRICT - build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim + build-depends: base >= 4.3 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim ghc-options: -O2 + other-extensions: CPP, BangPatterns include-dirs: include - extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types build-depends: HUnit, @@ -174,10 +322,10 @@ type: exitcode-stdio-1.0 cpp-options: -DTESTING - build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim + build-depends: base >= 4.3 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim ghc-options: -O2 + other-extensions: CPP, BangPatterns include-dirs: include - extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types build-depends: HUnit, @@ -192,10 +340,10 @@ type: exitcode-stdio-1.0 cpp-options: -DTESTING - build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim + build-depends: base >= 4.3 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim ghc-options: -O2 + other-extensions: CPP, BangPatterns include-dirs: include - extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types build-depends: QuickCheck, @@ -208,15 +356,16 @@ type: exitcode-stdio-1.0 cpp-options: -DTESTING - build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim + build-depends: base >= 4.3 && < 5, array, deepseq >= 1.2 && < 1.5, ghc-prim ghc-options: -O2 + other-extensions: CPP, BangPatterns include-dirs: include - extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types build-depends: QuickCheck, test-framework, - test-framework-quickcheck2 + test-framework-quickcheck2, + transformers test-suite map-strictness-properties hs-source-dirs: tests, . @@ -225,7 +374,7 @@ build-depends: array, - base >= 4.2 && < 5, + base >= 4.3 && < 5, ChasingBottoms, deepseq >= 1.2 && < 1.5, QuickCheck >= 2.4.0.1, @@ -234,16 +383,18 @@ test-framework-quickcheck2 >= 0.2.9 ghc-options: -Wall + other-extensions: CPP, BangPatterns include-dirs: include test-suite intmap-strictness-properties hs-source-dirs: tests, . main-is: intmap-strictness.hs type: exitcode-stdio-1.0 + other-extensions: CPP, BangPatterns build-depends: array, - base >= 4.2 && < 5, + base >= 4.3 && < 5, ChasingBottoms, deepseq >= 1.2 && < 1.5, QuickCheck >= 2.4.0.1, @@ -258,10 +409,11 @@ hs-source-dirs: tests, . main-is: intset-strictness.hs type: exitcode-stdio-1.0 + other-extensions: CPP, BangPatterns build-depends: array, - base >= 4.2 && < 5, + base >= 4.3 && < 5, ChasingBottoms, deepseq >= 1.2 && < 1.5, QuickCheck >= 2.4.0.1, diff --git a/include/containers.h b/include/containers.h --- a/include/containers.h +++ b/include/containers.h @@ -16,37 +16,18 @@ * Define INSTANCE_TYPEABLE[0-2] */ #if __GLASGOW_HASKELL__ >= 707 -#define INSTANCE_TYPEABLE0(tycon,tcname,str) deriving instance Typeable tycon -#define INSTANCE_TYPEABLE1(tycon,tcname,str) deriving instance Typeable tycon -#define INSTANCE_TYPEABLE2(tycon,tcname,str) deriving instance Typeable tycon +#define INSTANCE_TYPEABLE0(tycon) deriving instance Typeable tycon +#define INSTANCE_TYPEABLE1(tycon) deriving instance Typeable tycon +#define INSTANCE_TYPEABLE2(tycon) deriving instance Typeable tycon #elif defined(__GLASGOW_HASKELL__) -#define INSTANCE_TYPEABLE0(tycon,tcname,str) deriving instance Typeable tycon -#define INSTANCE_TYPEABLE1(tycon,tcname,str) deriving instance Typeable1 tycon -#define INSTANCE_TYPEABLE2(tycon,tcname,str) deriving instance Typeable2 tycon +#define INSTANCE_TYPEABLE0(tycon) deriving instance Typeable tycon +#define INSTANCE_TYPEABLE1(tycon) deriving instance Typeable1 tycon +#define INSTANCE_TYPEABLE2(tycon) deriving instance Typeable2 tycon #else -#define INSTANCE_TYPEABLE0(tycon,tcname,str) tcname :: TyCon; tcname = mkTyCon str; \ - instance Typeable tycon where { typeOf _ = mkTyConApp tcname [] } -#define INSTANCE_TYPEABLE1(tycon,tcname,str) tcname :: TyCon; tcname = mkTyCon str; \ - instance Typeable1 tycon where { typeOf1 _ = mkTyConApp tcname [] }; \ - instance Typeable a => Typeable (tycon a) where { typeOf = typeOfDefault } -#define INSTANCE_TYPEABLE2(tycon,tcname,str) tcname :: TyCon; tcname = mkTyCon str; \ - instance Typeable2 tycon where { typeOf2 _ = mkTyConApp tcname [] }; \ - instance Typeable a => Typeable1 (tycon a) where { typeOf1 = typeOf1Default }; \ - instance (Typeable a, Typeable b) => Typeable (tycon a b) where { typeOf = typeOfDefault } +#define INSTANCE_TYPEABLE0(tycon) +#define INSTANCE_TYPEABLE1(tycon) +#define INSTANCE_TYPEABLE2(tycon) #endif - -/* - * Use macros to define strictness of functions. - * STRICT_x_OF_y denotes an y-ary function strict in the x-th parameter. - * We do not use BangPatterns, because they are not in any standard and we - * want the compilers to be compiled by as many compilers as possible. - */ -#define STRICT_1_OF_2(fn) fn arg _ | arg `seq` False = undefined -#define STRICT_2_OF_2(fn) fn _ arg | arg `seq` False = undefined -#define STRICT_1_OF_3(fn) fn arg _ _ | arg `seq` False = undefined -#define STRICT_2_OF_3(fn) fn _ arg _ | arg `seq` False = undefined -#define STRICT_1_OF_4(fn) fn arg _ _ _ | arg `seq` False = undefined -#define STRICT_2_OF_4(fn) fn _ arg _ _ | arg `seq` False = undefined /* * We use cabal-generated MIN_VERSION_base to adapt to changes of base. diff --git a/tests/bitqueue-properties.hs b/tests/bitqueue-properties.hs new file mode 100644 --- /dev/null +++ b/tests/bitqueue-properties.hs @@ -0,0 +1,34 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} + +#if !MIN_VERSION_base(4,8,0) +import Control.Applicative ((<$>)) +#endif +import qualified Data.List as List +import Test.Framework +import Test.Framework.Providers.QuickCheck2 +import Test.QuickCheck +import Data.Utils.BitUtil (wordSize) +import Data.Utils.BitQueue + ( BitQueue + , emptyQB + , snocQB + , buildQ + , toListQ ) + +default (Int) + +main :: IO () +main = defaultMain $ map testNum [0..(wordSize - 2)] + +testNum :: Int -> Test +testNum n = testProperty ("Size "++show n) (prop_n n) + +prop_n :: Int -> Gen Bool +prop_n n = checkList <$> vectorOf n (arbitrary :: Gen Bool) + where + checkList :: [Bool] -> Bool + checkList values = toListQ q == values + where + q :: BitQueue + !q = buildQ $ List.foldl' snocQB emptyQB values diff --git a/tests/deprecated-properties.hs b/tests/deprecated-properties.hs --- a/tests/deprecated-properties.hs +++ b/tests/deprecated-properties.hs @@ -12,7 +12,7 @@ import Test.Framework import Test.Framework.Providers.QuickCheck2 -import Text.Show.Functions () +import Test.QuickCheck.Function (Fun(..), apply) default (Int) @@ -32,11 +32,16 @@ ---------- Map properties ---------- +apply2 :: Fun (a, b) c -> a -> b -> c +apply2 f a b = apply f (a, b) -prop_mapInsertWith'Strict :: [(Int, Int)] -> (Int -> Int -> Int) -> [(Int, Int)] -> Bool +apply3 :: Fun (a, b, c) d -> a -> b -> c -> d +apply3 f a b c = apply f (a, b, c) + +prop_mapInsertWith'Strict :: [(Int, Int)] -> Fun (Int, Int) Int -> [(Int, Int)] -> Bool prop_mapInsertWith'Strict xs f kxxs = let m = M.fromList xs - insertList ins = foldr (\(kx, x) -> ins f kx x) m kxxs + insertList ins = foldr (\(kx, x) -> ins (apply2 f) kx x) m kxxs in insertList M.insertWith' == insertList SM.insertWith prop_mapInsertWith'Undefined :: [(Int, Int)] -> Bool @@ -46,10 +51,10 @@ insertList ins = foldr (\(kx, _) -> ins f kx undefined) m xs in insertList M.insertWith' == insertList M.insertWith -prop_mapInsertWithKey'Strict :: [(Int, Int)] -> (Int -> Int -> Int -> Int) -> [(Int, Int)] -> Bool +prop_mapInsertWithKey'Strict :: [(Int, Int)] -> Fun (Int, Int, Int) Int -> [(Int, Int)] -> Bool prop_mapInsertWithKey'Strict xs f kxxs = let m = M.fromList xs - insertList ins = foldr (\(kx, x) -> ins f kx x) m kxxs + insertList ins = foldr (\(kx, x) -> ins (apply3 f) kx x) m kxxs in insertList M.insertWithKey' == insertList SM.insertWithKey prop_mapInsertWithKey'Undefined :: [(Int, Int)] -> Bool @@ -59,10 +64,10 @@ insertList ins = foldr (\(kx, _) -> ins f kx undefined) m xs in insertList M.insertWithKey' == insertList M.insertWithKey -prop_mapInsertLookupWithKey'Strict :: [(Int, Int)] -> (Int -> Int -> Int -> Int) -> [(Int, Int)] -> Bool +prop_mapInsertLookupWithKey'Strict :: [(Int, Int)] -> Fun (Int, Int, Int) Int -> [(Int, Int)] -> Bool prop_mapInsertLookupWithKey'Strict xs f kxxs = let m = M.fromList xs - insertLookupList insLkp = scanr (\(kx, x) (_, mp) -> insLkp f kx x mp) (Nothing, m) kxxs + insertLookupList insLkp = scanr (\(kx, x) (_, mp) -> insLkp (apply3 f) kx x mp) (Nothing, m) kxxs in insertLookupList M.insertLookupWithKey' == insertLookupList SM.insertLookupWithKey prop_mapInsertLookupWithKey'Undefined :: [(Int, Int)] -> Bool @@ -75,10 +80,10 @@ ---------- IntMap properties ---------- -prop_intmapInsertWith'Strict :: [(Int, Int)] -> (Int -> Int -> Int) -> [(Int, Int)] -> Bool +prop_intmapInsertWith'Strict :: [(Int, Int)] -> Fun (Int, Int) Int -> [(Int, Int)] -> Bool prop_intmapInsertWith'Strict xs f kxxs = let m = IM.fromList xs - insertList ins = foldr (\(kx, x) -> ins f kx x) m kxxs + insertList ins = foldr (\(kx, x) -> ins (apply2 f) kx x) m kxxs in insertList IM.insertWith' == insertList SIM.insertWith prop_intmapInsertWith'Undefined :: [(Int, Int)] -> Bool @@ -88,10 +93,10 @@ insertList ins = foldr (\(kx, _) -> ins f kx undefined) m xs in insertList IM.insertWith' == insertList IM.insertWith -prop_intmapInsertWithKey'Strict :: [(Int, Int)] -> (Int -> Int -> Int -> Int) -> [(Int, Int)] -> Bool +prop_intmapInsertWithKey'Strict :: [(Int, Int)] -> Fun (Int, Int, Int) Int -> [(Int, Int)] -> Bool prop_intmapInsertWithKey'Strict xs f kxxs = let m = IM.fromList xs - insertList ins = foldr (\(kx, x) -> ins f kx x) m kxxs + insertList ins = foldr (\(kx, x) -> ins (apply3 f) kx x) m kxxs in insertList IM.insertWithKey' == insertList SIM.insertWithKey prop_intmapInsertWithKey'Undefined :: [(Int, Int)] -> Bool diff --git a/tests/intmap-properties.hs b/tests/intmap-properties.hs --- a/tests/intmap-properties.hs +++ b/tests/intmap-properties.hs @@ -16,13 +16,13 @@ import Data.List (nub,sort) import qualified Data.List as List -import qualified Data.IntSet +import qualified Data.IntSet as IntSet import Test.Framework import Test.Framework.Providers.HUnit import Test.Framework.Providers.QuickCheck2 import Test.HUnit hiding (Test, Testable) import Test.QuickCheck -import Text.Show.Functions () +import Test.QuickCheck.Function (Fun(..), apply) default (Int) @@ -169,6 +169,13 @@ , testProperty "fromSet" prop_fromSet ] +apply2 :: Fun (a, b) c -> a -> b -> c +apply2 f a b = apply f (a, b) + +apply3 :: Fun (a, b, c) d -> a -> b -> c -> d +apply3 f a b c = apply f (a, b, c) + + {-------------------------------------------------------------------- Arbitrary, reasonably balanced trees --------------------------------------------------------------------} @@ -499,13 +506,13 @@ test_keysSet :: Assertion test_keysSet = do - keysSet (fromList [(5,"a"), (3,"b")]) @?= Data.IntSet.fromList [3,5] - keysSet (empty :: UMap) @?= Data.IntSet.empty + keysSet (fromList [(5,"a"), (3,"b")]) @?= IntSet.fromList [3,5] + keysSet (empty :: UMap) @?= IntSet.empty test_fromSet :: Assertion test_fromSet = do - fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) @?= fromList [(5,"aaaaa"), (3,"aaa")] - fromSet undefined Data.IntSet.empty @?= (empty :: IMap) + fromSet (\k -> replicate k 'a') (IntSet.fromList [3, 5]) @?= fromList [(5,"aaaaa"), (3,"aaa")] + fromSet undefined IntSet.empty @?= (empty :: IMap) ---------------------------------------------------------------- -- Lists @@ -796,6 +803,18 @@ ys' = List.nubBy ((==) `on` fst) ys f k l r = k + 2 * l + 3 * r +prop_restrictKeys :: IMap -> IMap -> Property +prop_restrictKeys m s0 = m `restrictKeys` s === filterWithKey (\k _ -> k `IntSet.member` s) m + where + s = keysSet s0 + restricted = restrictKeys m s + +prop_withoutKeys :: IMap -> IMap -> Property +prop_withoutKeys m s0 = m `withoutKeys` s === filterWithKey (\k _ -> k `IntSet.notMember` s) m + where + s = keysSet s0 + reduced = withoutKeys m s + prop_mergeWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Bool prop_mergeWithKeyModel xs ys = and [ testMergeWithKey f keep_x keep_y @@ -958,35 +977,35 @@ m = fromList xs in toAscList (deleteMax m) == init (sort xs) -prop_filter :: (Int -> Bool) -> [(Int, Int)] -> Property +prop_filter :: Fun Int Bool -> [(Int, Int)] -> Property prop_filter p ys = length ys > 0 ==> let xs = List.nubBy ((==) `on` fst) ys m = fromList xs - in filter p m == fromList (List.filter (p . snd) xs) + in filter (apply p) m == fromList (List.filter (apply p . snd) xs) -prop_partition :: (Int -> Bool) -> [(Int, Int)] -> Property +prop_partition :: Fun Int Bool -> [(Int, Int)] -> Property prop_partition p ys = length ys > 0 ==> let xs = List.nubBy ((==) `on` fst) ys m = fromList xs - in partition p m == let (a,b) = (List.partition (p . snd) xs) in (fromList a, fromList b) + in partition (apply p) m == let (a,b) = (List.partition (apply p . snd) xs) in (fromList a, fromList b) -prop_map :: (Int -> Int) -> [(Int, Int)] -> Property +prop_map :: Fun Int Int -> [(Int, Int)] -> Property prop_map f ys = length ys > 0 ==> let xs = List.nubBy ((==) `on` fst) ys m = fromList xs - in map f m == fromList [ (a, f b) | (a,b) <- xs ] + in map (apply f) m == fromList [ (a, apply f b) | (a,b) <- xs ] -prop_fmap :: (Int -> Int) -> [(Int, Int)] -> Property +prop_fmap :: Fun Int Int -> [(Int, Int)] -> Property prop_fmap f ys = length ys > 0 ==> let xs = List.nubBy ((==) `on` fst) ys m = fromList xs - in fmap f m == fromList [ (a, f b) | (a,b) <- xs ] + in fmap (apply f) m == fromList [ (a, apply f b) | (a,b) <- xs ] -prop_mapkeys :: (Int -> Int) -> [(Int, Int)] -> Property +prop_mapkeys :: Fun Int Int -> [(Int, Int)] -> Property prop_mapkeys f ys = length ys > 0 ==> let xs = List.nubBy ((==) `on` fst) ys m = fromList xs - in mapKeys f m == (fromList $ List.nubBy ((==) `on` fst) $ reverse [ (f a, b) | (a,b) <- sort xs]) + in mapKeys (apply f) m == (fromList $ List.nubBy ((==) `on` fst) $ reverse [ (apply f a, b) | (a,b) <- sort xs]) prop_splitModel :: Int -> [(Int, Int)] -> Property prop_splitModel n ys = length ys > 0 ==> @@ -1048,9 +1067,9 @@ prop_keysSet :: [(Int, Int)] -> Bool prop_keysSet xs = - keysSet (fromList xs) == Data.IntSet.fromList (List.map fst xs) + keysSet (fromList xs) == IntSet.fromList (List.map fst xs) prop_fromSet :: [(Int, Int)] -> Bool prop_fromSet ys = let xs = List.nubBy ((==) `on` fst) ys - in fromSet (\k -> fromJust $ List.lookup k xs) (Data.IntSet.fromList $ List.map fst xs) == fromList xs + in fromSet (\k -> fromJust $ List.lookup k xs) (IntSet.fromList $ List.map fst xs) == fromList xs diff --git a/tests/intmap-strictness.hs b/tests/intmap-strictness.hs --- a/tests/intmap-strictness.hs +++ b/tests/intmap-strictness.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE FlexibleInstances, GeneralizedNewtypeDeriving #-} {-# OPTIONS_GHC -fno-warn-orphans #-} module Main (main) where @@ -7,6 +6,7 @@ import Test.Framework (Test, defaultMain, testGroup) import Test.Framework.Providers.QuickCheck2 (testProperty) import Test.QuickCheck (Arbitrary(arbitrary)) +import Test.QuickCheck.Function (Fun(..), apply) import Data.IntMap.Strict (IntMap) import qualified Data.IntMap.Strict as M @@ -14,14 +14,11 @@ instance Arbitrary v => Arbitrary (IntMap v) where arbitrary = M.fromList `fmap` arbitrary -instance Show (Int -> Int) where - show _ = "" - -instance Show (Int -> Int -> Int) where - show _ = "" +apply2 :: Fun (a, b) c -> a -> b -> c +apply2 f a b = apply f (a, b) -instance Show (Int -> Int -> Int -> Int) where - show _ = "" +apply3 :: Fun (a, b, c) d -> a -> b -> c -> d +apply3 f a b c = apply f (a, b, c) ------------------------------------------------------------------------ -- * Properties @@ -42,8 +39,8 @@ pFindWithDefaultValueStrict k m = M.member k m || (isBottom $ M.findWithDefault bottom k m) -pAdjustKeyStrict :: (Int -> Int) -> IntMap Int -> Bool -pAdjustKeyStrict f m = isBottom $ M.adjust f bottom m +pAdjustKeyStrict :: Fun Int Int -> IntMap Int -> Bool +pAdjustKeyStrict f m = isBottom $ M.adjust (apply f) bottom m pAdjustValueStrict :: Int -> IntMap Int -> Bool pAdjustValueStrict k m @@ -58,26 +55,26 @@ pInsertValueStrict :: Int -> IntMap Int -> Bool pInsertValueStrict k m = isBottom $ M.insert k bottom m -pInsertWithKeyStrict :: (Int -> Int -> Int) -> Int -> IntMap Int -> Bool -pInsertWithKeyStrict f v m = isBottom $ M.insertWith f bottom v m +pInsertWithKeyStrict :: Fun (Int, Int) Int -> Int -> IntMap Int -> Bool +pInsertWithKeyStrict f v m = isBottom $ M.insertWith (apply2 f) bottom v m -pInsertWithValueStrict :: (Int -> Int -> Int) -> Int -> Int -> IntMap Int +pInsertWithValueStrict :: Fun (Int, Int) Int -> Int -> Int -> IntMap Int -> Bool pInsertWithValueStrict f k v m | M.member k m = (isBottom $ M.insertWith (const2 bottom) k v m) && not (isBottom $ M.insertWith (const2 1) k bottom m) - | otherwise = isBottom $ M.insertWith f k bottom m + | otherwise = isBottom $ M.insertWith (apply2 f) k bottom m -pInsertLookupWithKeyKeyStrict :: (Int -> Int -> Int -> Int) -> Int -> IntMap Int +pInsertLookupWithKeyKeyStrict :: Fun (Int, Int, Int) Int -> Int -> IntMap Int -> Bool -pInsertLookupWithKeyKeyStrict f v m = isBottom $ M.insertLookupWithKey f bottom v m +pInsertLookupWithKeyKeyStrict f v m = isBottom $ M.insertLookupWithKey (apply3 f) bottom v m -pInsertLookupWithKeyValueStrict :: (Int -> Int -> Int -> Int) -> Int -> Int +pInsertLookupWithKeyValueStrict :: Fun (Int, Int, Int) Int -> Int -> Int -> IntMap Int -> Bool pInsertLookupWithKeyValueStrict f k v m | M.member k m = (isBottom $ M.insertLookupWithKey (const3 bottom) k v m) && not (isBottom $ M.insertLookupWithKey (const3 1) k bottom m) - | otherwise = isBottom $ M.insertLookupWithKey f k bottom m + | otherwise = isBottom $ M.insertLookupWithKey (apply3 f) k bottom m ------------------------------------------------------------------------ -- * Test list diff --git a/tests/intset-strictness.hs b/tests/intset-strictness.hs --- a/tests/intset-strictness.hs +++ b/tests/intset-strictness.hs @@ -1,6 +1,3 @@ -{-# LANGUAGE FlexibleInstances, GeneralizedNewtypeDeriving #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} - module Main (main) where import Prelude hiding (foldl) diff --git a/tests/map-properties.hs b/tests/map-properties.hs --- a/tests/map-properties.hs +++ b/tests/map-properties.hs @@ -2,30 +2,43 @@ #ifdef STRICT import Data.Map.Strict as Data.Map +import Data.Map.Strict.Merge #else import Data.Map.Lazy as Data.Map +import Data.Map.Lazy.Merge #endif +import Data.Map.Base (Map (..), balanced, link2, link, bin) +import Control.Applicative (Const(Const, getConst), pure, (<$>), (<*>)) +import Data.Functor.Identity (Identity(runIdentity)) import Data.Monoid import Data.Maybe hiding (mapMaybe) import qualified Data.Maybe as Maybe (mapMaybe) import Data.Ord import Data.Function -import Prelude hiding (lookup, null, map, filter, foldr, foldl) -import qualified Prelude (map) +import Prelude hiding (lookup, null, map, filter, foldr, foldl, take, drop, splitAt) +import qualified Prelude import Data.List (nub,sort) import qualified Data.List as List -import qualified Data.Set +import qualified Data.Set as Set import Test.Framework import Test.Framework.Providers.HUnit import Test.Framework.Providers.QuickCheck2 import Test.HUnit hiding (Test, Testable) import Test.QuickCheck -import Text.Show.Functions () +import Test.QuickCheck.Function (Fun (..), apply) +import Test.QuickCheck.Poly (A, B) +import Control.Arrow (first) default (Int) +apply3 :: Fun (a,b,c) d -> a -> b -> c -> d +apply3 f a b c = apply f (a, b, c) + +apply2 :: Fun (a,b) c -> a -> b -> c +apply2 f a b = apply f (a, b) + main :: IO () main = defaultMain [ testCase "ticket4242" test_ticket4242 @@ -54,6 +67,7 @@ , testCase "updateWithKey" test_updateWithKey , testCase "updateLookupWithKey" test_updateLookupWithKey , testCase "alter" test_alter + , testCase "at" test_at , testCase "union" test_union , testCase "mappend" test_mappend , testCase "unionWith" test_unionWith @@ -92,6 +106,7 @@ , testCase "fromAscListWith" test_fromAscListWith , testCase "fromAscListWithKey" test_fromAscListWithKey , testCase "fromDistinctAscList" test_fromDistinctAscList + , testCase "fromDistinctDescList" test_fromDistinctDescList , testCase "filter" test_filter , testCase "filterWithKey" test_filteWithKey , testCase "partition" test_partition @@ -138,7 +153,7 @@ , testProperty "split" prop_split , testProperty "splitRoot" prop_splitRoot , testProperty "split then link" prop_link - , testProperty "split then merge" prop_merge + , testProperty "split then link2" prop_link2 , testProperty "union" prop_union , testProperty "union model" prop_unionModel , testProperty "union singleton" prop_unionSingleton @@ -148,19 +163,29 @@ , testProperty "union sum" prop_unionSum , testProperty "difference" prop_difference , testProperty "difference model" prop_differenceModel + , testProperty "withoutKeys" prop_withoutKeys , testProperty "intersection" prop_intersection + , testProperty "restrictKeys" prop_restrictKeys , testProperty "intersection model" prop_intersectionModel , testProperty "intersectionWith" prop_intersectionWith , testProperty "intersectionWithModel" prop_intersectionWithModel , testProperty "intersectionWithKey" prop_intersectionWithKey , testProperty "intersectionWithKeyModel" prop_intersectionWithKeyModel + , testProperty "differenceMerge" prop_differenceMerge + , testProperty "unionWithKeyMerge" prop_unionWithKeyMerge , testProperty "mergeWithKey model" prop_mergeWithKeyModel , testProperty "fromAscList" prop_ordered + , testProperty "fromDescList" prop_rev_ordered + , testProperty "fromDistinctDescList" prop_fromDistinctDescList , testProperty "fromList then toList" prop_list , testProperty "toDescList" prop_descList , testProperty "toAscList+toDescList" prop_ascDescList , testProperty "fromList" prop_fromList , testProperty "alter" prop_alter + , testProperty "alterF/alter" prop_alterF_alter + , testProperty "alterF/alter/noRULES" prop_alterF_alter_noRULES + , testProperty "alterF/lookup" prop_alterF_lookup + , testProperty "alterF/lookup/noRULES" prop_alterF_lookup_noRULES , testProperty "index" prop_index , testProperty "null" prop_null , testProperty "member" prop_member @@ -190,14 +215,20 @@ , testProperty "foldl'" prop_foldl' , testProperty "keysSet" prop_keysSet , testProperty "fromSet" prop_fromSet + , testProperty "takeWhileAntitone" prop_takeWhileAntitone + , testProperty "dropWhileAntitone" prop_dropWhileAntitone + , testProperty "spanAntitone" prop_spanAntitone + , testProperty "take" prop_take + , testProperty "drop" prop_drop + , testProperty "splitAt" prop_splitAt ] {-------------------------------------------------------------------- - Arbitrary, reasonably balanced trees + Arbitrary trees --------------------------------------------------------------------} instance (Enum k,Arbitrary a) => Arbitrary (Map k a) where arbitrary = sized (arbtree 0 maxkey) - where maxkey = 10^5 + where maxkey = 10^(5 :: Int) arbtree :: (Enum k, Arbitrary a) => Int -> Int -> Int -> Gen (Map k a) arbtree lo hi n = do t <- gentree lo hi n @@ -217,6 +248,18 @@ ; return (bin (toEnum i) x l r) } +-- A type with a peculiar Eq instance designed to make sure keys +-- come from where they're supposed to. +data OddEq a = OddEq a Bool deriving (Show) +getOddEq :: OddEq a -> (a, Bool) +getOddEq (OddEq a b) = (a, b) +instance Arbitrary a => Arbitrary (OddEq a) where + arbitrary = OddEq <$> arbitrary <*> arbitrary +instance Eq a => Eq (OddEq a) where + OddEq x _ == OddEq y _ = x == y +instance Ord a => Ord (OddEq a) where + OddEq x _ `compare` OddEq y _ = x `compare` y + ------------------------------------------------------------------------ type UMap = Map Int () @@ -405,6 +448,51 @@ f _ = Nothing g _ = Just "c" +test_at :: Assertion +test_at = do + employeeCurrency "John" @?= Just "Euro" + employeeCurrency "Pete" @?= Nothing + atAlter f 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a")] + atAlter f 5 (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "b" + atAlter g 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a"), (7, "c")] + atAlter g 5 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "c")] + where + f _ = Nothing + g _ = Just "c" + employeeDept = fromList([("John","Sales"), ("Bob","IT")]) + deptCountry = fromList([("IT","USA"), ("Sales","France")]) + countryCurrency = fromList([("USA", "Dollar"), ("France", "Euro")]) + employeeCurrency :: String -> Maybe String + employeeCurrency name = do + dept <- atLookup name employeeDept + country <- atLookup dept deptCountry + atLookup country countryCurrency + +-- This version of atAlter will rewrite to alterFIdentity +-- if the rules fire. +atAlter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a +atAlter f k m = runIdentity (alterF (pure . f) k m) + +-- A version of atAlter that uses a private copy of Identity +-- to ensure that the adjustF/Identity rules don't fire and +-- we use the basic implementation. +atAlterNoRULES :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a +atAlterNoRULES f k m = runIdent (alterF (Ident . f) k m) + +newtype Ident a = Ident { runIdent :: a } +instance Functor Ident where + fmap f (Ident a) = Ident (f a) + +atLookup :: Ord k => k -> Map k a -> Maybe a +atLookup k m = getConst (alterF Const k m) + +atLookupNoRULES :: Ord k => k -> Map k a -> Maybe a +atLookupNoRULES k m = getConsty (alterF Consty k m) + +newtype Consty a b = Consty { getConsty :: a} +instance Functor (Consty a) where + fmap _ (Consty a) = Consty a + ---------------------------------------------------------------- -- Combine @@ -531,13 +619,13 @@ test_keysSet :: Assertion test_keysSet = do - keysSet (fromList [(5,"a"), (3,"b")]) @?= Data.Set.fromList [3,5] - keysSet (empty :: UMap) @?= Data.Set.empty + keysSet (fromList [(5,"a"), (3,"b")]) @?= Set.fromList [3,5] + keysSet (empty :: UMap) @?= Set.empty test_fromSet :: Assertion test_fromSet = do - fromSet (\k -> replicate k 'a') (Data.Set.fromList [3, 5]) @?= fromList [(5,"aaaaa"), (3,"aaa")] - fromSet undefined Data.Set.empty @?= (empty :: IMap) + fromSet (\k -> replicate k 'a') (Set.fromList [3, 5]) @?= fromList [(5,"aaaaa"), (3,"aaa")] + fromSet undefined Set.empty @?= (empty :: IMap) ---------------------------------------------------------------- -- Lists @@ -616,6 +704,12 @@ valid (fromDistinctAscList [(3,"b"), (5,"a")]) @?= True valid (fromDistinctAscList [(3,"b"), (5,"a"), (5,"b")]) @?= False +test_fromDistinctDescList :: Assertion +test_fromDistinctDescList = do + fromDistinctDescList [(5,"a"), (3,"b")] @?= fromList [(3, "b"), (5, "a")] + valid (fromDistinctDescList [(5,"a"), (3,"b")]) @?= True + valid (fromDistinctDescList [(3,"b"), (5,"a"), (5,"b")]) @?= False + ---------------------------------------------------------------- -- Filter @@ -827,6 +921,18 @@ -- QuickCheck ---------------------------------------------------------------- +prop_differenceMerge :: Fun (Int, A, B) (Maybe A) -> Map Int A -> Map Int B -> Property +prop_differenceMerge f m1 m2 = + differenceWithKey (apply3 f) m1 m2 === merge preserveMissing dropMissing (zipWithMaybeMatched (apply3 f)) m1 m2 + +prop_unionWithKeyMerge :: Fun (Int, A, A) A -> Map Int A -> Map Int A -> Property +prop_unionWithKeyMerge f m1 m2 = + unionWithKey (apply3 f) m1 m2 === unionWithKey' (apply3 f) m1 m2 + +unionWithKey' :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a +unionWithKey' f = merge preserveMissing preserveMissing $ + zipWithMatched (\k a b -> f k a b) + prop_valid :: UMap -> Bool prop_valid t = valid t @@ -874,9 +980,9 @@ prop_link k t = let (l,r) = split k t in valid (link k () l r) -prop_merge :: Int -> UMap -> Bool -prop_merge k t = let (l,r) = split k t - in valid (merge l r) +prop_link2 :: Int -> UMap -> Bool +prop_link2 k t = let (l,r) = split k t + in valid (link2 l r) ---------------------------------------------------------------- @@ -913,6 +1019,18 @@ = sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys))) == sort ((List.\\) (nub (Prelude.map fst xs)) (nub (Prelude.map fst ys))) +prop_restrictKeys :: IMap -> IMap -> Property +prop_restrictKeys m s0 = valid restricted .&&. (m `restrictKeys` s === filterWithKey (\k _ -> k `Set.member` s) m) + where + s = keysSet s0 + restricted = restrictKeys m s + +prop_withoutKeys :: IMap -> IMap -> Property +prop_withoutKeys m s0 = valid reduced .&&. (m `withoutKeys` s === filterWithKey (\k _ -> k `Set.notMember` s) m) + where + s = keysSet s0 + reduced = withoutKeys m s + prop_intersection :: IMap -> IMap -> Bool prop_intersection t1 t2 = valid (intersection t1 t2) @@ -921,8 +1039,8 @@ = sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys))) == sort (nub ((List.intersect) (Prelude.map fst xs) (Prelude.map fst ys))) -prop_intersectionWith :: (Int -> Int -> Maybe Int) -> IMap -> IMap -> Bool -prop_intersectionWith f t1 t2 = valid (intersectionWith f t1 t2) +prop_intersectionWith :: Fun (Int, Int) (Maybe Int) -> IMap -> IMap -> Bool +prop_intersectionWith f t1 t2 = valid (intersectionWith (apply2 f) t1 t2) prop_intersectionWithModel :: [(Int,Int)] -> [(Int,Int)] -> Bool prop_intersectionWithModel xs ys @@ -932,8 +1050,8 @@ ys' = List.nubBy ((==) `on` fst) ys f l r = l + 2 * r -prop_intersectionWithKey :: (Int -> Int -> Int -> Maybe Int) -> IMap -> IMap -> Bool -prop_intersectionWithKey f t1 t2 = valid (intersectionWithKey f t1 t2) +prop_intersectionWithKey :: Fun (Int, Int, Int) (Maybe Int) -> IMap -> IMap -> Bool +prop_intersectionWithKey f t1 t2 = valid (intersectionWithKey (apply3 f) t1 t2) prop_intersectionWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Bool prop_intersectionWithKeyModel xs ys @@ -986,12 +1104,23 @@ let xs = [(x,()) | x <- [0..n::Int]] in fromAscList xs == fromList xs +prop_rev_ordered :: Property +prop_rev_ordered + = forAll (choose (5,100)) $ \n -> + let xs = [(x,()) | x <- [0..n::Int]] + in fromDescList (reverse xs) == fromList xs + prop_list :: [Int] -> Bool prop_list xs = (sort (nub xs) == [x | (x,()) <- toList (fromList [(x,()) | x <- xs])]) prop_descList :: [Int] -> Bool prop_descList xs = (reverse (sort (nub xs)) == [x | (x,()) <- toDescList (fromList [(x,()) | x <- xs])]) +prop_fromDistinctDescList :: Int -> [A] -> Property +prop_fromDistinctDescList top lst = valid converted .&&. (toList converted === reverse original) where + original = zip [top, (top-1)..0] lst + converted = fromDistinctDescList original + prop_ascDescList :: [Int] -> Bool prop_ascDescList xs = toAscList m == reverse (toDescList m) where m = fromList $ zip xs $ repeat () @@ -1016,6 +1145,21 @@ f Nothing = Just () f (Just ()) = Nothing +prop_alterF_alter :: Fun (Maybe Int) (Maybe Int) -> Int -> IMap -> Bool +prop_alterF_alter f k m = valid altered && altered == alter (apply f) k m + where altered = atAlter (apply f) k m + +prop_alterF_alter_noRULES :: Fun (Maybe Int) (Maybe Int) -> Int -> IMap -> Bool +prop_alterF_alter_noRULES f k m = valid altered && + altered == alter (apply f) k m + where altered = atAlterNoRULES (apply f) k m + +prop_alterF_lookup :: Int -> IMap -> Bool +prop_alterF_lookup k m = atLookup k m == lookup k m + +prop_alterF_lookup_noRULES :: Int -> IMap -> Bool +prop_alterF_lookup_noRULES k m = atLookupNoRULES k m == lookup k m + ------------------------------------------------------------------------ -- Compare against the list model (after nub on keys) @@ -1115,35 +1259,79 @@ m = fromList xs in toAscList (deleteMax m) == init (sort xs) -prop_filter :: (Int -> Bool) -> [(Int, Int)] -> Property +prop_filter :: Fun Int Bool -> [(Int, Int)] -> Property prop_filter p ys = length ys > 0 ==> let xs = List.nubBy ((==) `on` fst) ys m = fromList xs - in filter p m == fromList (List.filter (p . snd) xs) + in filter (apply p) m == fromList (List.filter (apply p . snd) xs) -prop_partition :: (Int -> Bool) -> [(Int, Int)] -> Property +prop_take :: Int -> Map Int Int -> Property +prop_take n xs = valid taken .&&. + taken === fromDistinctAscList (List.take n (toList xs)) + where + taken = take n xs + +prop_drop :: Int -> Map Int Int -> Property +prop_drop n xs = valid dropped .&&. + dropped === fromDistinctAscList (List.drop n (toList xs)) + where + dropped = drop n xs + +prop_splitAt :: Int -> Map Int Int -> Property +prop_splitAt n xs = valid taken .&&. + valid dropped .&&. + taken === take n xs .&&. + dropped === drop n xs + where + (taken, dropped) = splitAt n xs + +prop_takeWhileAntitone :: [(Either Int Int, Int)] -> Property +prop_takeWhileAntitone xs' = valid tw .&&. (tw === filterWithKey (\k _ -> isLeft k) xs) + where + xs = fromList xs' + tw = takeWhileAntitone isLeft xs + +prop_dropWhileAntitone :: [(Either Int Int, Int)] -> Property +prop_dropWhileAntitone xs' = valid tw .&&. (tw === filterWithKey (\k _ -> not (isLeft k)) xs) + where + xs = fromList xs' + tw = dropWhileAntitone isLeft xs + +prop_spanAntitone :: [(Either Int Int, Int)] -> Property +prop_spanAntitone xs' = valid tw .&&. valid dw + .&&. (tw === takeWhileAntitone isLeft xs) + .&&. (dw === dropWhileAntitone isLeft xs) + where + xs = fromList xs' + (tw, dw) = spanAntitone isLeft xs + +isLeft :: Either a b -> Bool +isLeft (Left _) = True +isLeft _ = False + +prop_partition :: Fun Int Bool -> [(Int, Int)] -> Property prop_partition p ys = length ys > 0 ==> let xs = List.nubBy ((==) `on` fst) ys m = fromList xs - in partition p m == let (a,b) = (List.partition (p . snd) xs) in (fromList a, fromList b) + in partition (apply p) m == let (a,b) = (List.partition (apply p . snd) xs) in (fromList a, fromList b) -prop_map :: (Int -> Int) -> [(Int, Int)] -> Property +prop_map :: Fun Int Int -> [(Int, Int)] -> Property prop_map f ys = length ys > 0 ==> let xs = List.nubBy ((==) `on` fst) ys m = fromList xs - in map f m == fromList [ (a, f b) | (a,b) <- xs ] + in map (apply f) m == fromList [ (a, apply f b) | (a,b) <- xs ] -prop_fmap :: (Int -> Int) -> [(Int, Int)] -> Property +prop_fmap :: Fun Int Int -> [(Int, Int)] -> Property prop_fmap f ys = length ys > 0 ==> let xs = List.nubBy ((==) `on` fst) ys m = fromList xs - in fmap f m == fromList [ (a, f b) | (a,b) <- xs ] + in fmap (apply f) m == fromList [ (a, (apply f) b) | (a,b) <- xs ] -prop_mapkeys :: (Int -> Int) -> [(Int, Int)] -> Property +prop_mapkeys :: Fun Int Int -> [(Int, Int)] -> Property prop_mapkeys f ys = length ys > 0 ==> let xs = List.nubBy ((==) `on` fst) ys m = fromList xs - in mapKeys f m == (fromList $ List.nubBy ((==) `on` fst) $ reverse [ (f a, b) | (a,b) <- sort xs]) + in mapKeys (apply f) m == (fromList $ List.nubBy ((==) `on` fst) $ reverse [ (apply f a, b) | (a,b) <- sort xs]) prop_splitModel :: Int -> [(Int, Int)] -> Property prop_splitModel n ys = length ys > 0 ==> @@ -1195,9 +1383,9 @@ prop_keysSet :: [(Int, Int)] -> Bool prop_keysSet xs = - keysSet (fromList xs) == Data.Set.fromList (List.map fst xs) + keysSet (fromList xs) == Set.fromList (List.map fst xs) prop_fromSet :: [(Int, Int)] -> Bool prop_fromSet ys = let xs = List.nubBy ((==) `on` fst) ys - in fromSet (\k -> fromJust $ List.lookup k xs) (Data.Set.fromList $ List.map fst xs) == fromList xs + in fromSet (\k -> fromJust $ List.lookup k xs) (Set.fromList $ List.map fst xs) == fromList xs diff --git a/tests/map-strictness.hs b/tests/map-strictness.hs --- a/tests/map-strictness.hs +++ b/tests/map-strictness.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE FlexibleInstances, GeneralizedNewtypeDeriving #-} {-# OPTIONS_GHC -fno-warn-orphans #-} module Main (main) where @@ -7,22 +6,20 @@ import Test.Framework (Test, defaultMain, testGroup) import Test.Framework.Providers.QuickCheck2 (testProperty) import Test.QuickCheck (Arbitrary(arbitrary)) +import Test.QuickCheck.Function (Fun(..), apply) import Data.Map.Strict (Map) import qualified Data.Map.Strict as M -instance (Arbitrary k, Arbitrary v, Eq k, Ord k) => +instance (Arbitrary k, Arbitrary v, Ord k) => Arbitrary (Map k v) where arbitrary = M.fromList `fmap` arbitrary -instance Show (Int -> Int) where - show _ = "" - -instance Show (Int -> Int -> Int) where - show _ = "" +apply2 :: Fun (a, b) c -> a -> b -> c +apply2 f a b = apply f (a, b) -instance Show (Int -> Int -> Int -> Int) where - show _ = "" +apply3 :: Fun (a, b, c) d -> a -> b -> c -> d +apply3 f a b c = apply f (a, b, c) ------------------------------------------------------------------------ -- * Properties @@ -43,8 +40,8 @@ pFindWithDefaultValueStrict k m = M.member k m || (isBottom $ M.findWithDefault bottom k m) -pAdjustKeyStrict :: (Int -> Int) -> Map Int Int -> Bool -pAdjustKeyStrict f m = isBottom $ M.adjust f bottom m +pAdjustKeyStrict :: Fun Int Int -> Map Int Int -> Bool +pAdjustKeyStrict f m = isBottom $ M.adjust (apply f) bottom m pAdjustValueStrict :: Int -> Map Int Int -> Bool pAdjustValueStrict k m @@ -59,26 +56,26 @@ pInsertValueStrict :: Int -> Map Int Int -> Bool pInsertValueStrict k m = isBottom $ M.insert k bottom m -pInsertWithKeyStrict :: (Int -> Int -> Int) -> Int -> Map Int Int -> Bool -pInsertWithKeyStrict f v m = isBottom $ M.insertWith f bottom v m +pInsertWithKeyStrict :: Fun (Int, Int) Int -> Int -> Map Int Int -> Bool +pInsertWithKeyStrict f v m = isBottom $ M.insertWith (apply2 f) bottom v m -pInsertWithValueStrict :: (Int -> Int -> Int) -> Int -> Int -> Map Int Int +pInsertWithValueStrict :: Fun (Int, Int) Int -> Int -> Int -> Map Int Int -> Bool pInsertWithValueStrict f k v m | M.member k m = (isBottom $ M.insertWith (const2 bottom) k v m) && not (isBottom $ M.insertWith (const2 1) k bottom m) - | otherwise = isBottom $ M.insertWith f k bottom m + | otherwise = isBottom $ M.insertWith (apply2 f) k bottom m -pInsertLookupWithKeyKeyStrict :: (Int -> Int -> Int -> Int) -> Int +pInsertLookupWithKeyKeyStrict :: Fun (Int, Int, Int) Int -> Int -> Map Int Int -> Bool -pInsertLookupWithKeyKeyStrict f v m = isBottom $ M.insertLookupWithKey f bottom v m +pInsertLookupWithKeyKeyStrict f v m = isBottom $ M.insertLookupWithKey (apply3 f) bottom v m -pInsertLookupWithKeyValueStrict :: (Int -> Int -> Int -> Int) -> Int -> Int +pInsertLookupWithKeyValueStrict :: Fun (Int, Int, Int) Int -> Int -> Int -> Map Int Int -> Bool pInsertLookupWithKeyValueStrict f k v m | M.member k m = (isBottom $ M.insertLookupWithKey (const3 bottom) k v m) && not (isBottom $ M.insertLookupWithKey (const3 1) k bottom m) - | otherwise = isBottom $ M.insertLookupWithKey f k bottom m + | otherwise = isBottom $ M.insertLookupWithKey (apply3 f) k bottom m ------------------------------------------------------------------------ -- * Test list diff --git a/tests/seq-properties.hs b/tests/seq-properties.hs --- a/tests/seq-properties.hs +++ b/tests/seq-properties.hs @@ -1,15 +1,16 @@ -import Data.Sequence -- needs to be compiled with -DTESTING for use here +import Data.Sequence.Base import Control.Applicative (Applicative(..)) import Control.Arrow ((***)) +import Control.Monad.Trans.State.Strict import Data.Array (listArray) -import Data.Foldable (Foldable(foldl, foldl1, foldr, foldr1, foldMap), toList, all, sum) +import Data.Foldable (Foldable(foldl, foldl1, foldr, foldr1, foldMap, fold), toList, all, sum, foldl', foldr') import Data.Functor ((<$>), (<$)) import Data.Maybe -import Data.Monoid (Monoid(..)) +import Data.Monoid (Monoid(..), All(..), Endo(..), Dual(..)) import Data.Traversable (Traversable(traverse), sequenceA) import Prelude hiding ( - null, length, take, drop, splitAt, + lookup, null, length, take, drop, splitAt, foldl, foldl1, foldr, foldr1, scanl, scanl1, scanr, scanr1, filter, reverse, replicate, zip, zipWith, zip3, zipWith3, all, sum) @@ -27,8 +28,10 @@ [ testProperty "fmap" prop_fmap , testProperty "(<$)" prop_constmap , testProperty "foldr" prop_foldr + , testProperty "foldr'" prop_foldr' , testProperty "foldr1" prop_foldr1 , testProperty "foldl" prop_foldl + , testProperty "foldl'" prop_foldl' , testProperty "foldl1" prop_foldl1 , testProperty "(==)" prop_equals , testProperty "compare" prop_compare @@ -71,11 +74,15 @@ , testProperty "unstableSort" prop_unstableSort , testProperty "unstableSortBy" prop_unstableSortBy , testProperty "index" prop_index + , testProperty "(!?)" prop_safeIndex , testProperty "adjust" prop_adjust + , testProperty "insertAt" prop_insertAt + , testProperty "deleteAt" prop_deleteAt , testProperty "update" prop_update , testProperty "take" prop_take , testProperty "drop" prop_drop , testProperty "splitAt" prop_splitAt + , testProperty "chunksOf" prop_chunksOf , testProperty "elemIndexL" prop_elemIndexL , testProperty "elemIndicesL" prop_elemIndicesL , testProperty "elemIndexR" prop_elemIndexR @@ -87,6 +94,9 @@ , testProperty "foldlWithIndex" prop_foldlWithIndex , testProperty "foldrWithIndex" prop_foldrWithIndex , testProperty "mapWithIndex" prop_mapWithIndex + , testProperty "foldMapWithIndex/foldlWithIndex" prop_foldMapWithIndexL + , testProperty "foldMapWithIndex/foldrWithIndex" prop_foldMapWithIndexR + , testProperty "traverseWithIndex" prop_traverseWithIndex , testProperty "reverse" prop_reverse , testProperty "zip" prop_zip , testProperty "zipWith" prop_zipWith @@ -96,6 +106,8 @@ , testProperty "zipWith4" prop_zipWith4 , testProperty "<*>" prop_ap , testProperty "*>" prop_then + , testProperty "cycleTaking" prop_cycleTaking + , testProperty "intersperse" prop_intersperse , testProperty ">>=" prop_bind ] @@ -113,8 +125,8 @@ instance (Arbitrary a, Sized a) => Arbitrary (FingerTree a) where arbitrary = sized arb where - arb :: (Arbitrary a, Sized a) => Int -> Gen (FingerTree a) - arb 0 = return Empty + arb :: (Arbitrary b, Sized b) => Int -> Gen (FingerTree b) + arb 0 = return EmptyT arb 1 = Single <$> arbitrary arb n = do pr <- arbitrary @@ -126,13 +138,13 @@ m <- arb n_m return $ deep pr m sf - shrink (Deep _ (One a) Empty (One b)) = [Single a, Single b] + shrink (Deep _ (One a) EmptyT (One b)) = [Single a, Single b] shrink (Deep _ pr m sf) = [deep pr' m sf | pr' <- shrink pr] ++ [deep pr m' sf | m' <- shrink m] ++ [deep pr m sf' | sf' <- shrink sf] shrink (Single x) = map Single (shrink x) - shrink Empty = [] + shrink EmptyT = [] instance (Arbitrary a, Sized a) => Arbitrary (Node a) where arbitrary = oneof [ @@ -174,7 +186,7 @@ valid (Seq xs) = valid xs instance (Sized a, Valid a) => Valid (FingerTree a) where - valid Empty = True + valid EmptyT = True valid (Single x) = valid x valid (Deep s pr m sf) = s == size pr + size m + size sf && valid pr && valid m && valid sf @@ -223,25 +235,39 @@ prop_constmap x xs = toList' (x <$ xs) ~= map (const x) (toList xs) -prop_foldr :: Seq A -> Bool +prop_foldr :: Seq A -> Property prop_foldr xs = - foldr f z xs == Prelude.foldr f z (toList xs) + foldr f z xs === Prelude.foldr f z (toList xs) where f = (:) z = [] +prop_foldr' :: Seq A -> Property +prop_foldr' xs = + foldr' f z xs === foldr' f z (toList xs) + where + f = (:) + z = [] + prop_foldr1 :: Seq Int -> Property prop_foldr1 xs = not (null xs) ==> foldr1 f xs == Data.List.foldr1 f (toList xs) where f = (-) -prop_foldl :: Seq A -> Bool +prop_foldl :: Seq A -> Property prop_foldl xs = - foldl f z xs == Prelude.foldl f z (toList xs) + foldl f z xs === Prelude.foldl f z (toList xs) where f = flip (:) z = [] +prop_foldl' :: Seq A -> Property +prop_foldl' xs = + foldl' f z xs === foldl' f z (toList xs) + where + f = flip (:) + z = [] + prop_foldl1 :: Seq Int -> Property prop_foldl1 xs = not (null xs) ==> foldl1 f xs == Data.List.foldl1 f (toList xs) @@ -475,6 +501,26 @@ not (null xs) ==> forAll (choose (0, length xs-1)) $ \ i -> index xs i == toList xs !! i +prop_safeIndex :: Seq A -> Property +prop_safeIndex xs = + forAll (choose (-3, length xs + 3)) $ \i -> + ((i < 0 || i >= length xs) .&&. lookup i xs === Nothing) .||. + lookup i xs === Just (toList xs !! i) + +prop_insertAt :: A -> Seq A -> Property +prop_insertAt x xs = + forAll (choose (-3, length xs + 3)) $ \i -> + let res = insertAt i x xs + in valid res .&&. res === case splitAt i xs of (front, back) -> front >< x <| back + +prop_deleteAt :: Seq A -> Property +prop_deleteAt xs = + forAll (choose (-3, length xs + 3)) $ \i -> + let res = deleteAt i xs + in valid res .&&. + (((0 <= i && i < length xs) .&&. res === case splitAt i xs of (front, back) -> front >< drop 1 back) + .||. ((i < 0 || i >= length xs) .&&. res === xs)) + prop_adjust :: Int -> Int -> Seq Int -> Bool prop_adjust n i xs = toList' (adjust f i xs) ~= adjustList f i (toList xs) @@ -496,6 +542,14 @@ prop_splitAt n xs = toListPair' (splitAt n xs) ~= Prelude.splitAt n (toList xs) +prop_chunksOf :: Seq A -> Property +prop_chunksOf xs = + forAll (choose (1, length xs + 3)) $ \n -> + let chunks = chunksOf n xs + in valid chunks .&&. + conjoin [valid c .&&. 1 <= length c && length c <= n | c <- toList chunks] .&&. + fold chunks === xs + adjustList :: (a -> a) -> Int -> [a] -> [a] adjustList f i xs = [if j == i then f x else x | (j, x) <- Prelude.zip [0..] xs] @@ -552,6 +606,16 @@ foldrWithIndex f z xs == Data.List.foldr (uncurry f) z (Data.List.zip [0..] (toList xs)) where f n y ys = (n,y):ys +prop_foldMapWithIndexL :: (Fun (B, Int, A) B) -> B -> Seq A -> Bool +prop_foldMapWithIndexL (Fun _ f) z t = foldlWithIndex f' z t == + appEndo (getDual (foldMapWithIndex (\i -> Dual . Endo . flip (flip f' i)) t)) z + where f' b i a = f (b, i, a) + +prop_foldMapWithIndexR :: (Fun (Int, A, B) B) -> B -> Seq A -> Bool +prop_foldMapWithIndexR (Fun _ f) z t = foldrWithIndex f' z t == + appEndo (foldMapWithIndex (\i -> Endo . f' i) t) z + where f' i a b = f (i, a, b) + -- * Transformations prop_mapWithIndex :: Seq A -> Bool @@ -559,6 +623,11 @@ toList' (mapWithIndex f xs) ~= map (uncurry f) (Data.List.zip [0..] (toList xs)) where f = (,) +prop_traverseWithIndex :: Seq Int -> Bool +prop_traverseWithIndex xs = + runState (traverseWithIndex (\i x -> modify ((i,x) :)) xs) [] == + runState (sequenceA . mapWithIndex (\i x -> modify ((i,x) :)) $ xs) [] + prop_reverse :: Seq A -> Bool prop_reverse xs = toList' (reverse xs) ~= Prelude.reverse (toList xs) @@ -601,6 +670,14 @@ prop_then :: Seq A -> Seq B -> Bool prop_then xs ys = toList' (xs *> ys) ~= (toList xs *> toList ys) + +prop_intersperse :: A -> Seq A -> Bool +prop_intersperse x xs = + toList' (intersperse x xs) ~= Data.List.intersperse x (toList xs) + +prop_cycleTaking :: Int -> Seq A -> Property +prop_cycleTaking n xs = + (n <= 0 || not (null xs)) ==> toList' (cycleTaking n xs) ~= Data.List.take n (Data.List.cycle (toList xs)) -- Monad operations diff --git a/tests/set-properties.hs b/tests/set-properties.hs --- a/tests/set-properties.hs +++ b/tests/set-properties.hs @@ -1,15 +1,26 @@ +{-# LANGUAGE CPP #-} import qualified Data.IntSet as IntSet import Data.List (nub,sort) import qualified Data.List as List import Data.Monoid (mempty) import Data.Maybe import Data.Set -import Prelude hiding (lookup, null, map, filter, foldr, foldl) +import Prelude hiding (lookup, null, map, filter, foldr, foldl, all, take, drop, splitAt) import Test.Framework import Test.Framework.Providers.HUnit import Test.Framework.Providers.QuickCheck2 import Test.HUnit hiding (Test, Testable) import Test.QuickCheck +import Test.QuickCheck.Function +import Test.QuickCheck.Poly +import Control.Monad.Trans.State.Strict +import Control.Monad.Trans.Class +import Control.Monad (liftM, liftM3) +import Data.Functor.Identity +import Data.Foldable (all) +#if !MIN_VERSION_base(4,8,0) +import Control.Applicative (Applicative (..), (<$>)) +#endif main :: IO () main = defaultMain [ testCase "lookupLT" test_lookupLT @@ -30,6 +41,7 @@ , testProperty "prop_LookupGE" prop_LookupGE , testProperty "prop_InsertValid" prop_InsertValid , testProperty "prop_InsertDelete" prop_InsertDelete + , testProperty "prop_InsertBiased" prop_InsertBiased , testProperty "prop_DeleteValid" prop_DeleteValid , testProperty "prop_Link" prop_Link , testProperty "prop_Merge" prop_Merge @@ -37,15 +49,19 @@ , testProperty "prop_UnionInsert" prop_UnionInsert , testProperty "prop_UnionAssoc" prop_UnionAssoc , testProperty "prop_UnionComm" prop_UnionComm + , testProperty "prop_UnionBiased" prop_UnionBiased , testProperty "prop_DiffValid" prop_DiffValid , testProperty "prop_Diff" prop_Diff , testProperty "prop_IntValid" prop_IntValid , testProperty "prop_Int" prop_Int + , testProperty "prop_IntBiased" prop_IntBiased , testProperty "prop_Ordered" prop_Ordered + , testProperty "prop_DescendingOrdered" prop_DescendingOrdered , testProperty "prop_List" prop_List , testProperty "prop_DescList" prop_DescList , testProperty "prop_AscDescList" prop_AscDescList , testProperty "prop_fromList" prop_fromList + , testProperty "prop_fromListDesc" prop_fromListDesc , testProperty "prop_isProperSubsetOf" prop_isProperSubsetOf , testProperty "prop_isProperSubsetOf2" prop_isProperSubsetOf2 , testProperty "prop_isSubsetOf" prop_isSubsetOf @@ -60,6 +76,8 @@ , testProperty "prop_foldL" prop_foldL , testProperty "prop_foldL'" prop_foldL' , testProperty "prop_map" prop_map + , testProperty "prop_map2" prop_map2 + , testProperty "prop_mapMonotonic" prop_mapMonotonic , testProperty "prop_maxView" prop_maxView , testProperty "prop_minView" prop_minView , testProperty "prop_split" prop_split @@ -67,8 +85,27 @@ , testProperty "prop_splitRoot" prop_splitRoot , testProperty "prop_partition" prop_partition , testProperty "prop_filter" prop_filter + , testProperty "takeWhileAntitone" prop_takeWhileAntitone + , testProperty "dropWhileAntitone" prop_dropWhileAntitone + , testProperty "spanAntitone" prop_spanAntitone + , testProperty "take" prop_take + , testProperty "drop" prop_drop + , testProperty "splitAt" prop_splitAt ] +-- A type with a peculiar Eq instance designed to make sure keys +-- come from where they're supposed to. +data OddEq a = OddEq a Bool deriving (Show) + +getOddEq :: OddEq a -> (a, Bool) +getOddEq (OddEq b a) = (b, a) +instance Arbitrary a => Arbitrary (OddEq a) where + arbitrary = OddEq <$> arbitrary <*> arbitrary +instance Eq a => Eq (OddEq a) where + OddEq x _ == OddEq y _ = x == y +instance Ord a => Ord (OddEq a) where + OddEq x _ `compare` OddEq y _ = x `compare` y + ---------------------------------------------------------------- -- Unit tests ---------------------------------------------------------------- @@ -124,37 +161,138 @@ {-------------------------------------------------------------------- Arbitrary, reasonably balanced trees --------------------------------------------------------------------} -instance (Enum a) => Arbitrary (Set a) where - arbitrary = sized (arbtree 0 maxkey) - where maxkey = 10000 - arbtree :: (Enum a) => Int -> Int -> Int -> Gen (Set a) - arbtree lo hi n = do t <- gentree lo hi n - if balanced t then return t else arbtree lo hi n - where gentree lo hi n - | n <= 0 = return Tip - | lo >= hi = return Tip - | otherwise = do i <- choose (lo,hi) - m <- choose (1,70) - let (ml,mr) | m==(1::Int) = (1,2) - | m==2 = (2,1) - | m==3 = (1,1) - | otherwise = (2,2) - l <- gentree lo (i-1) (n `div` ml) - r <- gentree (i+1) hi (n `div` mr) - return (bin (toEnum i) l r) +-- | The IsInt class lets us constrain a type variable to be Int in an entirely +-- standard way. The constraint @ IsInt a @ is essentially equivalent to the +-- GHC-only constraint @ a ~ Int @, but @ IsInt @ requires manual intervention +-- to use. If ~ is ever standardized, we should certainly use it instead. +-- Earlier versions used an Enum constraint, but this is confusing because +-- not all Enum instances will work properly for the Arbitrary instance here. +class (Show a, Read a, Integral a, Arbitrary a) => IsInt a where + fromIntF :: f Int -> f a +instance IsInt Int where + fromIntF = id + +-- | Convert an Int to any instance of IsInt +fromInt :: IsInt a => Int -> a +fromInt = runIdentity . fromIntF . Identity + +{- We don't actually need this, but we can add it if we ever do +toIntF :: IsInt a => g a -> g Int +toIntF = unf . fromIntF . F $ id + +newtype F g a b = F {unf :: g b -> a} + +toInt :: IsInt a => a -> Int +toInt = runIdentity . toIntF . Identity -} + + +-- How much the minimum value of an arbitrary set should vary +positionFactor :: Int +positionFactor = 1 + +-- How much the gap between consecutive elements in an arbitrary +-- set should vary +gapRange :: Int +gapRange = 5 + +instance IsInt a => Arbitrary (Set a) where + arbitrary = sized (\sz0 -> do + sz <- choose (0, sz0) + middle <- choose (-positionFactor * (sz + 1), positionFactor * (sz + 1)) + let shift = (sz * (gapRange) + 1) `quot` 2 + start = middle - shift + t <- evalStateT (mkArb step sz) start + if valid t then pure t else error "Test generated invalid tree!") + where + step = do + i <- get + diff <- lift $ choose (1, gapRange) + let i' = i + diff + put i' + pure (fromInt i') + +class Monad m => MonadGen m where + liftGen :: Gen a -> m a +instance MonadGen Gen where + liftGen = id +instance MonadGen m => MonadGen (StateT s m) where + liftGen = lift . liftGen + +-- | Given an action that produces successively larger elements and +-- a size, produce a set of arbitrary shape with exactly that size. +mkArb :: MonadGen m => m a -> Int -> m (Set a) +mkArb step n + | n <= 0 = return Tip + | n == 1 = singleton `liftM` step + | n == 2 = do + dir <- liftGen arbitrary + p <- step + q <- step + if dir + then return (Bin 2 q (singleton p) Tip) + else return (Bin 2 p Tip (singleton q)) + | otherwise = do + -- This assumes a balance factor of delta = 3 + let upper = (3*(n - 1)) `quot` 4 + let lower = (n + 2) `quot` 4 + ln <- liftGen $ choose (lower, upper) + let rn = n - ln - 1 + liftM3 (\lt x rt -> Bin n x lt rt) (mkArb step ln) step (mkArb step rn) + +-- | Given a strictly increasing list of elements, produce an arbitrarily +-- shaped set with exactly those elements. +setFromList :: [a] -> Gen (Set a) +setFromList xs = flip evalStateT xs $ mkArb step (length xs) + where + step = do + x : xs <- get + put xs + pure x + +data TwoSets = TwoSets (Set Int) (Set Int) deriving (Show) + +data TwoLists a = TwoLists [a] [a] + +data Options2 = One2 | Two2 | Both2 deriving (Bounded, Enum) +instance Arbitrary Options2 where + arbitrary = arbitraryBoundedEnum + +-- We produce two lists from a simple "universe". This instance +-- is intended to give good results when the two lists are then +-- combined with each other; if other elements are used with them, +-- they may or may not behave particularly well. +instance IsInt a => Arbitrary (TwoLists a) where + arbitrary = sized $ \sz0 -> do + sz <- choose (0, sz0) + let universe = [0,3..3*(fromInt sz - 1)] + divide2Gen universe + +instance Arbitrary TwoSets where + arbitrary = do + TwoLists l r <- arbitrary + TwoSets <$> setFromList l <*> setFromList r + +divide2Gen :: [a] -> Gen (TwoLists a) +divide2Gen [] = pure (TwoLists [] []) +divide2Gen (x : xs) = do + way <- arbitrary + TwoLists ls rs <- divide2Gen xs + case way of + One2 -> pure (TwoLists (x : ls) rs) + Two2 -> pure (TwoLists ls (x : rs)) + Both2 -> pure (TwoLists (x : ls) (x : rs)) + {-------------------------------------------------------------------- - Valid tree's + Valid trees --------------------------------------------------------------------} -forValid :: (Enum a,Show a,Testable b) => (Set a -> b) -> Property +forValid :: (IsInt a,Testable b) => (Set a -> b) -> Property forValid f = forAll arbitrary $ \t -> --- classify (balanced t) "balanced" $ classify (size t == 0) "empty" $ classify (size t > 0 && size t <= 10) "small" $ classify (size t > 10 && size t <= 64) "medium" $ - classify (size t > 64) "large" $ - balanced t ==> f t + classify (size t > 64) "large" $ f t forValidUnitTree :: Testable a => (Set Int -> a) -> Property forValidUnitTree f = forValid f @@ -210,6 +348,13 @@ prop_InsertDelete :: Int -> Set Int -> Property prop_InsertDelete k t = not (member k t) ==> delete k (insert k t) == t +prop_InsertBiased :: Int -> Set Int -> Bool +prop_InsertBiased k t = (k, True) `member` kt + where + t' = mapMonotonic (`OddEq` False) t + kt' = insert (OddEq k True) t' + kt = mapMonotonic getOddEq kt' + prop_DeleteValid :: Int -> Property prop_DeleteValid k = forValidUnitTree $ \t -> valid (delete k (insert k t)) @@ -241,9 +386,22 @@ prop_UnionAssoc :: Set Int -> Set Int -> Set Int -> Bool prop_UnionAssoc t1 t2 t3 = union t1 (union t2 t3) == union (union t1 t2) t3 -prop_UnionComm :: Set Int -> Set Int -> Bool -prop_UnionComm t1 t2 = (union t1 t2 == union t2 t1) +prop_UnionComm :: TwoSets -> Bool +prop_UnionComm (TwoSets t1 t2) = (union t1 t2 == union t2 t1) +prop_UnionBiased :: TwoSets -> Property +prop_UnionBiased (TwoSets l r) = union l' r' === union l' (difference r' l') + where + l' = mapMonotonic (`OddEq` False) l + r' = mapMonotonic (`OddEq` True) r + +prop_IntBiased :: TwoSets -> Bool +prop_IntBiased (TwoSets l r) = all (\(OddEq _ b) -> not b) l'r' + where + l' = mapMonotonic (`OddEq` False) l + r' = mapMonotonic (`OddEq` True) r + l'r' = intersection l' r' + prop_DiffValid :: Property prop_DiffValid = forValidUnitTree $ \t1 -> forValidUnitTree $ \t2 -> @@ -268,8 +426,13 @@ prop_Ordered :: Property prop_Ordered = forAll (choose (5,100)) $ \n -> let xs = [0..n::Int] - in fromAscList xs == fromList xs + in fromAscList xs === fromList xs +prop_DescendingOrdered :: Property +prop_DescendingOrdered = forAll (choose (5,100)) $ \n -> + let xs = [n,n-1..0::Int] + in fromDescList xs === fromList xs + prop_List :: [Int] -> Bool prop_List xs = (sort (nub xs) == toList (fromList xs)) @@ -280,15 +443,24 @@ prop_AscDescList xs = toAscList s == reverse (toDescList s) where s = fromList xs -prop_fromList :: [Int] -> Bool -prop_fromList xs - = case fromList xs of - t -> t == fromAscList sort_xs && - t == fromDistinctAscList nub_sort_xs && - t == List.foldr insert empty xs - where sort_xs = sort xs +prop_fromList :: [Int] -> Property +prop_fromList xs = + t === fromAscList sort_xs .&&. + t === fromDistinctAscList nub_sort_xs .&&. + t === List.foldr insert empty xs + where t = fromList xs + sort_xs = sort xs nub_sort_xs = List.map List.head $ List.group sort_xs +prop_fromListDesc :: [Int] -> Property +prop_fromListDesc xs = + t === fromDescList sort_xs .&&. + t === fromDistinctDescList nub_sort_xs .&&. + t === List.foldr insert empty xs + where t = fromList xs + sort_xs = reverse (sort xs) + nub_sort_xs = List.map List.head $ List.group sort_xs + {-------------------------------------------------------------------- Set operations are like IntSet operations --------------------------------------------------------------------} @@ -296,21 +468,21 @@ toIntSet = IntSet.fromList . toList -- Check that Set Int.isProperSubsetOf is the same as Set.isProperSubsetOf. -prop_isProperSubsetOf :: Set Int -> Set Int -> Bool -prop_isProperSubsetOf a b = isProperSubsetOf a b == IntSet.isProperSubsetOf (toIntSet a) (toIntSet b) +prop_isProperSubsetOf :: TwoSets -> Bool +prop_isProperSubsetOf (TwoSets a b) = isProperSubsetOf a b == IntSet.isProperSubsetOf (toIntSet a) (toIntSet b) -- In the above test, isProperSubsetOf almost always returns False (since a -- random set is almost never a subset of another random set). So this second -- test checks the True case. -prop_isProperSubsetOf2 :: Set Int -> Set Int -> Bool -prop_isProperSubsetOf2 a b = isProperSubsetOf a c == (a /= c) where +prop_isProperSubsetOf2 :: TwoSets -> Bool +prop_isProperSubsetOf2 (TwoSets a b) = isProperSubsetOf a c == (a /= c) where c = union a b -prop_isSubsetOf :: Set Int -> Set Int -> Bool -prop_isSubsetOf a b = isSubsetOf a b == IntSet.isSubsetOf (toIntSet a) (toIntSet b) +prop_isSubsetOf :: TwoSets -> Bool +prop_isSubsetOf (TwoSets a b) = isSubsetOf a b == IntSet.isSubsetOf (toIntSet a) (toIntSet b) -prop_isSubsetOf2 :: Set Int -> Set Int -> Bool -prop_isSubsetOf2 a b = isSubsetOf a (union a b) +prop_isSubsetOf2 :: TwoSets -> Bool +prop_isSubsetOf2 (TwoSets a b) = isSubsetOf a (union a b) prop_size :: Set Int -> Bool prop_size s = size s == List.length (toList s) @@ -321,8 +493,8 @@ prop_findMin :: Set Int -> Property prop_findMin s = not (null s) ==> findMin s == minimum (toList s) -prop_ord :: Set Int -> Set Int -> Bool -prop_ord s1 s2 = s1 `compare` s2 == toList s1 `compare` toList s2 +prop_ord :: TwoSets -> Bool +prop_ord (TwoSets s1 s2) = s1 `compare` s2 == toList s1 `compare` toList s2 prop_readShow :: Set Int -> Bool prop_readShow s = s == read (show s) @@ -342,6 +514,12 @@ prop_map :: Set Int -> Bool prop_map s = map id s == s +prop_map2 :: Fun Int Int -> Fun Int Int -> Set Int -> Property +prop_map2 f g s = map (apply f) (map (apply g) s) === map (apply f . apply g) s + +prop_mapMonotonic :: Set Int -> Property +prop_mapMonotonic s = mapMonotonic id s === s + prop_maxView :: Set Int -> Bool prop_maxView s = case maxView s of Nothing -> null s @@ -376,3 +554,47 @@ prop_filter :: Set Int -> Int -> Bool prop_filter s i = partition odd s == (filter odd s, filter even s) + +prop_take :: Int -> Set Int -> Property +prop_take n xs = valid taken .&&. + taken === fromDistinctAscList (List.take n (toList xs)) + where + taken = take n xs + +prop_drop :: Int -> Set Int -> Property +prop_drop n xs = valid dropped .&&. + dropped === fromDistinctAscList (List.drop n (toList xs)) + where + dropped = drop n xs + +prop_splitAt :: Int -> Set Int -> Property +prop_splitAt n xs = valid taken .&&. + valid dropped .&&. + taken === take n xs .&&. + dropped === drop n xs + where + (taken, dropped) = splitAt n xs + +prop_takeWhileAntitone :: [Either Int Int] -> Property +prop_takeWhileAntitone xs' = valid tw .&&. tw === filter isLeft xs + where + xs = fromList xs' + tw = takeWhileAntitone isLeft xs + +prop_dropWhileAntitone :: [Either Int Int] -> Property +prop_dropWhileAntitone xs' = valid tw .&&. tw === filter (not . isLeft) xs + where + xs = fromList xs' + tw = dropWhileAntitone isLeft xs + +prop_spanAntitone :: [Either Int Int] -> Property +prop_spanAntitone xs' = valid tw .&&. valid dw + .&&. tw === takeWhileAntitone isLeft xs + .&&. dw === dropWhileAntitone isLeft xs + where + xs = fromList xs' + (tw, dw) = spanAntitone isLeft xs + +isLeft :: Either a b -> Bool +isLeft (Left _) = True +isLeft _ = False