diff --git a/Data/Graph.hs b/Data/Graph.hs --- a/Data/Graph.hs +++ b/Data/Graph.hs @@ -1,3 +1,7 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE Rank2Types #-} +#endif #if __GLASGOW_HASKELL__ >= 703 {-# LANGUAGE Trustworthy #-} #endif @@ -6,7 +10,7 @@ -- Module : Data.Graph -- Copyright : (c) The University of Glasgow 2002 -- License : BSD-style (see the file libraries/base/LICENSE) --- +-- -- Maintainer : libraries@haskell.org -- Stability : experimental -- Portability : portable @@ -20,36 +24,36 @@ module Data.Graph( - -- * External interface + -- * External interface - -- At present the only one with a "nice" external interface - stronglyConnComp, stronglyConnCompR, SCC(..), flattenSCC, flattenSCCs, + -- At present the only one with a "nice" external interface + stronglyConnComp, stronglyConnCompR, SCC(..), flattenSCC, flattenSCCs, - -- * Graphs + -- * Graphs - Graph, Table, Bounds, Edge, Vertex, + Graph, Table, Bounds, Edge, Vertex, - -- ** Building graphs + -- ** Building graphs - graphFromEdges, graphFromEdges', buildG, transposeG, - -- reverseE, + graphFromEdges, graphFromEdges', buildG, transposeG, + -- reverseE, - -- ** Graph properties + -- ** Graph properties - vertices, edges, - outdegree, indegree, + vertices, edges, + outdegree, indegree, - -- * Algorithms + -- * Algorithms - dfs, dff, - topSort, - components, - scc, - bcc, - -- tree, back, cross, forward, - reachable, path, + dfs, dff, + topSort, + components, + scc, + bcc, + -- tree, back, cross, forward, + reachable, path, - module Data.Tree + module Data.Tree ) where @@ -68,22 +72,27 @@ import Data.Tree (Tree(Node), Forest) -- std interfaces +import Control.DeepSeq (NFData(rnf)) import Data.Maybe import Data.Array import Data.List ------------------------------------------------------------------------- --- - --- External interface --- - +-- - +-- External interface +-- - ------------------------------------------------------------------------- -- | Strongly connected component. -data SCC vertex = AcyclicSCC vertex -- ^ A single vertex that is not - -- in any cycle. - | CyclicSCC [vertex] -- ^ A maximal set of mutually - -- reachable vertices. +data SCC vertex = AcyclicSCC vertex -- ^ A single vertex that is not + -- in any cycle. + | CyclicSCC [vertex] -- ^ A maximal set of mutually + -- reachable vertices. +instance NFData a => NFData (SCC a) where + rnf (AcyclicSCC v) = rnf v + rnf (CyclicSCC vs) = rnf vs + -- | The vertices of a list of strongly connected components. flattenSCCs :: [SCC a] -> [a] flattenSCCs = concatMap flattenSCC @@ -96,13 +105,13 @@ -- | The strongly connected components of a directed graph, topologically -- sorted. stronglyConnComp - :: Ord key - => [(node, key, [key])] - -- ^ The graph: a list of nodes uniquely identified by keys, - -- with a list of keys of nodes this node has edges to. - -- The out-list may contain keys that don't correspond to - -- nodes of the graph; such edges are ignored. - -> [SCC node] + :: Ord key + => [(node, key, [key])] + -- ^ The graph: a list of nodes uniquely identified by keys, + -- with a list of keys of nodes this node has edges to. + -- The out-list may contain keys that don't correspond to + -- nodes of the graph; such edges are ignored. + -> [SCC node] stronglyConnComp edges0 = map get_node (stronglyConnCompR edges0) @@ -117,31 +126,31 @@ -- (some of) the result of 'SCC', so you don't want to lose the -- dependency information. stronglyConnCompR - :: Ord key - => [(node, key, [key])] - -- ^ The graph: a list of nodes uniquely identified by keys, - -- with a list of keys of nodes this node has edges to. - -- The out-list may contain keys that don't correspond to - -- nodes of the graph; such edges are ignored. - -> [SCC (node, key, [key])] -- ^ Topologically sorted + :: Ord key + => [(node, key, [key])] + -- ^ The graph: a list of nodes uniquely identified by keys, + -- with a list of keys of nodes this node has edges to. + -- The out-list may contain keys that don't correspond to + -- nodes of the graph; such edges are ignored. + -> [SCC (node, key, [key])] -- ^ Topologically sorted stronglyConnCompR [] = [] -- added to avoid creating empty array in graphFromEdges -- SOF stronglyConnCompR edges0 = map decode forest where (graph, vertex_fn,_) = graphFromEdges edges0 - forest = scc graph + forest = scc graph decode (Node v []) | mentions_itself v = CyclicSCC [vertex_fn v] - | otherwise = AcyclicSCC (vertex_fn v) + | otherwise = AcyclicSCC (vertex_fn v) decode other = CyclicSCC (dec other []) - where - dec (Node v ts) vs = vertex_fn v : foldr dec vs ts + where + dec (Node v ts) vs = vertex_fn v : foldr dec vs ts mentions_itself v = v `elem` (graph ! v) ------------------------------------------------------------------------- --- - --- Graphs --- - +-- - +-- Graphs +-- - ------------------------------------------------------------------------- -- | Abstract representation of vertices. @@ -191,9 +200,9 @@ -- does not include the function which maps keys to vertices. This -- version of 'graphFromEdges' is for backwards compatibility. graphFromEdges' - :: Ord key - => [(node, key, [key])] - -> (Graph, Vertex -> (node, key, [key])) + :: Ord key + => [(node, key, [key])] + -> (Graph, Vertex -> (node, key, [key])) graphFromEdges' x = (a,b) where (a,b,_) = graphFromEdges x @@ -202,40 +211,40 @@ -- The out-list may contain keys that don't correspond to -- nodes of the graph; they are ignored. graphFromEdges - :: Ord key - => [(node, key, [key])] - -> (Graph, Vertex -> (node, key, [key]), key -> Maybe Vertex) + :: Ord key + => [(node, key, [key])] + -> (Graph, Vertex -> (node, key, [key]), key -> Maybe Vertex) graphFromEdges edges0 = (graph, \v -> vertex_map ! v, key_vertex) where - max_v = length edges0 - 1 + max_v = length edges0 - 1 bounds0 = (0,max_v) :: (Vertex, Vertex) sorted_edges = sortBy lt edges0 - edges1 = zipWith (,) [0..] sorted_edges + edges1 = zipWith (,) [0..] sorted_edges - graph = array bounds0 [(,) v (mapMaybe key_vertex ks) | (,) v (_, _, ks) <- edges1] - key_map = array bounds0 [(,) v k | (,) v (_, k, _ ) <- edges1] - vertex_map = array bounds0 edges1 + graph = array bounds0 [(,) v (mapMaybe key_vertex ks) | (,) v (_, _, ks) <- edges1] + key_map = array bounds0 [(,) v k | (,) v (_, k, _ ) <- edges1] + vertex_map = array bounds0 edges1 (_,k1,_) `lt` (_,k2,_) = k1 `compare` k2 -- key_vertex :: key -> Maybe Vertex - -- returns Nothing for non-interesting vertices + -- returns Nothing for non-interesting vertices key_vertex k = findVertex 0 max_v - where - findVertex a b | a > b - = Nothing - findVertex a b = case compare k (key_map ! mid) of - LT -> findVertex a (mid-1) - EQ -> Just mid - GT -> findVertex (mid+1) b - where - mid = (a + b) `div` 2 + where + findVertex a b | a > b + = Nothing + findVertex a b = case compare k (key_map ! mid) of + LT -> findVertex a (mid-1) + EQ -> Just mid + GT -> findVertex (mid+1) b + where + mid = (a + b) `div` 2 ------------------------------------------------------------------------- --- - --- Depth first search --- - +-- - +-- Depth first search +-- - ------------------------------------------------------------------------- -- | A spanning forest of the graph, obtained from a depth-first search of @@ -310,9 +319,9 @@ #endif /* !USE_ST_MONAD */ ------------------------------------------------------------------------- --- - --- Algorithms --- - +-- - +-- Algorithms +-- - ------------------------------------------------------------------------- ------------------------------------------------------------ diff --git a/Data/IntMap.hs b/Data/IntMap.hs --- a/Data/IntMap.hs +++ b/Data/IntMap.hs @@ -1,2010 +1,95 @@ -{-# LANGUAGE NoBangPatterns, ScopedTypeVariables #-} -#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 -{-# LANGUAGE Trustworthy #-} -#endif ------------------------------------------------------------------------------ --- | --- Module : Data.IntMap --- 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 integer keys to values. --- --- Since many function names (but not the type name) clash with --- "Prelude" names, this module is usually imported @qualified@, e.g. --- --- > import Data.IntMap (IntMap) --- > import qualified Data.IntMap as IntMap --- --- The implementation is based on /big-endian patricia trees/. This data --- structure performs especially well on binary operations like 'union' --- and 'intersection'. However, my benchmarks show that it is also --- (much) faster on insertions and deletions when compared to a generic --- size-balanced map implementation (see "Data.Map"). --- --- * 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), --- October 1968, pages 514-534. --- --- Operation comments contain the operation time complexity in --- the Big-O notation . --- Many operations have a worst-case complexity of /O(min(n,W))/. --- This means that the operation can become linear in the number of --- elements with a maximum of /W/ -- the number of bits in an 'Int' --- (32 or 64). ------------------------------------------------------------------------------ - --- It is essential that the bit fiddling functions like mask, zero, branchMask --- etc are inlined. If they do not, the memory allocation skyrockets. The GHC --- usually gets it right, but it is disastrous if it does not. Therefore we --- explicitly mark these functions INLINE. - -module Data.IntMap ( - -- * Map type -#if !defined(TESTING) - IntMap, Key -- instance Eq,Show -#else - IntMap(..), Key -- instance Eq,Show -#endif - - -- * Operators - , (!), (\\) - - -- * Query - , null - , size - , member - , notMember - , lookup - , findWithDefault - - -- * Construction - , empty - , singleton - - -- ** Insertion - , insert - , insertWith - , insertWith' - , insertWithKey - , 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 - - -- * Traversal - -- ** Map - , map - , mapWithKey - , mapAccum - , mapAccumWithKey - , mapAccumRWithKey - - -- * Folds - , foldr - , foldl - , foldrWithKey - , foldlWithKey - -- ** Strict folds - , foldr' - , foldl' - , foldrWithKey' - , foldlWithKey' - -- ** Legacy folds - , fold - , foldWithKey - - -- * Conversion - , elems - , keys - , keysSet - , assocs - - -- ** Lists - , toList - , fromList - , fromListWith - , fromListWithKey - - -- ** Ordered lists - , toAscList - , fromAscList - , fromAscListWith - , fromAscListWithKey - , fromDistinctAscList - - -- * Filter - , filter - , filterWithKey - , partition - , partitionWithKey - - , mapMaybe - , mapMaybeWithKey - , mapEither - , mapEitherWithKey - - , split - , splitLookup - - -- * Submap - , isSubmapOf, isSubmapOfBy - , isProperSubmapOf, isProperSubmapOfBy - - -- * Min\/Max - , findMin - , findMax - , deleteMin - , deleteMax - , deleteFindMin - , deleteFindMax - , updateMin - , updateMax - , updateMinWithKey - , updateMaxWithKey - , minView - , maxView - , minViewWithKey - , maxViewWithKey - - -- * Debugging - , showTree - , showTreeWith - ) where - -import Prelude hiding (lookup,map,filter,foldr,foldl,null) -import Data.Bits -import qualified Data.IntSet as IntSet -import Data.Monoid (Monoid(..)) -import Data.Maybe (fromMaybe) -import Data.Typeable -import qualified Data.Foldable as Foldable -import Data.Traversable (Traversable(traverse)) -import Control.Applicative (Applicative(pure,(<*>)),(<$>)) -import Control.Monad ( liftM ) -import Control.DeepSeq (NFData(rnf)) -{- --- just for testing -import qualified Prelude -import Test.QuickCheck -import List (nub,sort) -import qualified List --} - -#if __GLASGOW_HASKELL__ -import Text.Read -import Data.Data (Data(..), mkNoRepType) -#endif - -#if __GLASGOW_HASKELL__ >= 503 -import GHC.Exts ( Word(..), Int(..), shiftRL# ) -#elif __GLASGOW_HASKELL__ -import Word -import GlaExts ( Word(..), Int(..), shiftRL# ) -#else -import Data.Word -#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 - -infixl 9 \\{-This comment teaches CPP correct behaviour -} - --- A "Nat" is a natural machine word (an unsigned Int) -type Nat = Word - -natFromInt :: Key -> Nat -natFromInt = fromIntegral -{-# INLINE natFromInt #-} - -intFromNat :: Nat -> Key -intFromNat = fromIntegral -{-# INLINE intFromNat #-} - -shiftRL :: Nat -> Key -> Nat -#if __GLASGOW_HASKELL__ -{-------------------------------------------------------------------- - GHC: use unboxing to get @shiftRL@ inlined. ---------------------------------------------------------------------} -shiftRL (W# x) (I# i) - = W# (shiftRL# x i) -#else -shiftRL x i = shiftR x i -{-# INLINE shiftRL #-} -#endif - -{-------------------------------------------------------------------- - Operators ---------------------------------------------------------------------} - --- | /O(min(n,W))/. 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' - -(!) :: IntMap a -> Key -> a -m ! k = find k m - --- | Same as 'difference'. -(\\) :: IntMap a -> IntMap b -> IntMap a -m1 \\ m2 = difference m1 m2 - -{-------------------------------------------------------------------- - Types ---------------------------------------------------------------------} - --- The order of constructors of IntMap matters when considering performance. --- Currently in GHC 7.0, when type has 3 constructors, they are matched from --- the first to the last -- the best performance is achieved when the --- constructors are ordered by frequency. --- On GHC 7.0, reordering constructors from Nil | Tip | Bin to Bin | Tip | Nil --- improves the containers_benchmark by 9.5% on x86 and by 8% on x86_64. - --- | A map of integers to values @a@. -data IntMap a = Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a) - | Tip {-# UNPACK #-} !Key a - | Nil - -type Prefix = Int -type Mask = Int -type Key = Int - -instance Monoid (IntMap a) where - mempty = empty - mappend = union - mconcat = unions - -instance Foldable.Foldable IntMap where - fold Nil = mempty - fold (Tip _ v) = v - fold (Bin _ _ l r) = Foldable.fold l `mappend` Foldable.fold r - foldr = foldr - foldl = foldl - foldMap _ Nil = mempty - foldMap f (Tip _k v) = f v - foldMap f (Bin _ _ l r) = Foldable.foldMap f l `mappend` Foldable.foldMap f r - -instance Traversable IntMap where - traverse _ Nil = pure Nil - traverse f (Tip k v) = Tip k <$> f v - traverse f (Bin p m l r) = Bin p m <$> traverse f l <*> traverse f r - -instance NFData a => NFData (IntMap a) where - rnf Nil = () - rnf (Tip _ v) = rnf v - rnf (Bin _ _ l r) = rnf l `seq` rnf r - -#if __GLASGOW_HASKELL__ - -{-------------------------------------------------------------------- - A Data instance ---------------------------------------------------------------------} - --- This instance preserves data abstraction at the cost of inefficiency. --- We omit reflection services for the sake of data abstraction. - -instance Data a => Data (IntMap a) where - gfoldl f z im = z fromList `f` (toList im) - toConstr _ = error "toConstr" - gunfold _ _ = error "gunfold" - dataTypeOf _ = mkNoRepType "Data.IntMap.IntMap" - dataCast1 f = gcast1 f - -#endif - -{-------------------------------------------------------------------- - Query ---------------------------------------------------------------------} --- | /O(1)/. Is the map empty? --- --- > Data.IntMap.null (empty) == True --- > Data.IntMap.null (singleton 1 'a') == False - -null :: IntMap a -> Bool -null Nil = True -null _ = False - --- | /O(n)/. Number of elements in the map. --- --- > size empty == 0 --- > 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 - --- | /O(min(n,W))/. Is the key a member of the map? --- --- > member 5 (fromList [(5,'a'), (3,'b')]) == True --- > member 1 (fromList [(5,'a'), (3,'b')]) == False - -member :: Key -> IntMap a -> Bool -member k m - = case lookup k m of - Nothing -> False - Just _ -> True - --- | /O(log n)/. Is the key not a member of the map? --- --- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False --- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True - -notMember :: Key -> IntMap a -> Bool -notMember k m = not $ member k m - --- The 'go' function in the lookup causes 10% speedup, but also an increased --- memory allocation. It does not cause speedup with other methods like insert --- and delete, so it is present only in lookup. - --- | /O(min(n,W))/. Lookup the value at a key in the map. See also 'Data.Map.lookup'. -lookup :: Key -> IntMap a -> Maybe a -lookup k = k `seq` go - where - go (Bin _ m l r) - | zero k m = go l - | otherwise = go r - go (Tip kx x) - | k == kx = Just x - | otherwise = Nothing - go Nil = Nothing - - -find :: Key -> IntMap a -> a -find k m - = case lookup k m of - Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map") - Just x -> x - --- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@ --- returns the value at key @k@ or returns @def@ when the key is not an --- element of the map. --- --- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' --- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' - -findWithDefault :: a -> Key -> IntMap a -> a -findWithDefault def k m - = case lookup k m of - Nothing -> def - Just x -> x - -{-------------------------------------------------------------------- - Construction ---------------------------------------------------------------------} --- | /O(1)/. The empty map. --- --- > empty == fromList [] --- > size empty == 0 - -empty :: IntMap a -empty - = Nil - --- | /O(1)/. A map of one element. --- --- > singleton 1 'a' == fromList [(1, 'a')] --- > size (singleton 1 'a') == 1 - -singleton :: Key -> a -> IntMap a -singleton k x - = Tip k x - -{-------------------------------------------------------------------- - Insert ---------------------------------------------------------------------} --- | /O(min(n,W))/. Insert a new key\/value pair in the map. --- If the key is already present in the map, the associated value is --- replaced with the supplied value, i.e. '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' - -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 -> join 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 -> join k (Tip k x) ky t - Nil -> Tip k x - --- right-biased insertion, used by 'union' --- | /O(min(n,W))/. Insert with a combining function. --- @'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 @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 :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a -insertWith f k x t - = insertWithKey (\_ x' y' -> f x' y') k x t - --- | Same as 'insertWith', but the combining function is applied strictly. -insertWith' :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a -insertWith' f k x t - = insertWithKey' (\_ x' y' -> f x' y') k x t - --- | /O(min(n,W))/. Insert with a combining function. --- @'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 @f key new_value old_value@. --- --- > 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" - -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 -> join 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 -> join k (Tip k x) ky t - Nil -> Tip k x - --- | Same as 'insertWithKey', but the combining function is applied strictly. -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 -> join 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 -> let x' = f k x y in seq x' (Tip k x') - | otherwise -> join k (Tip k x) ky t - 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@) --- 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")]) - -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,join 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,join k (Tip k x) ky t) - Nil -> (Nothing,Tip k x) - - -{-------------------------------------------------------------------- - Deletion - [delete] is the inlined version of [deleteWith (\k x -> Nothing)] ---------------------------------------------------------------------} --- | /O(min(n,W))/. 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 - -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 - --- | /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. --- --- > 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 :: (a -> a) -> Key -> IntMap a -> IntMap a -adjust f k m - = adjustWithKey (\_ x -> f x) k m - --- | /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. --- --- > 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 :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a -adjustWithKey f - = updateWithKey (\k' x -> Just (f k' x)) - --- | /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@. --- --- > 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 :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a -update f - = updateWithKey (\_ x -> f x) - --- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@ --- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is --- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. --- --- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing --- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] --- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] --- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -updateWithKey :: (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 - Just y' -> Tip ky y' - Nothing -> Nil - | otherwise -> t - Nil -> Nil - --- | /O(min(n,W))/. Lookup and update. --- The function returns original value, if it is updated. --- This is different behavior than 'Data.Map.updateLookupWithKey'. --- 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 "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") - -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) - - - --- | /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 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 -> join 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 -> join k (Tip k x) ky t - Nothing -> Tip ky y - Nil -> case f Nothing of - Just x -> Tip k x - Nothing -> Nil - - -{-------------------------------------------------------------------- - Union ---------------------------------------------------------------------} --- | The union of a list of maps. --- --- > 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 :: [IntMap a] -> IntMap a -unions xs - = foldlStrict union empty xs - --- | The union of a list of maps, with a combining operation. --- --- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] --- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] - -unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a -unionsWith f ts - = foldlStrict (unionWith f) empty ts - --- | /O(n+m)/. The (left-biased) union of two maps. --- It prefers the first map 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 :: IntMap a -> IntMap a -> IntMap a -union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = union1 - | shorter m2 m1 = union2 - | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2) - | otherwise = join p1 t1 p2 t2 - where - union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2 - | zero p2 m1 = Bin p1 m1 (union l1 t2) r1 - | otherwise = Bin p1 m1 l1 (union r1 t2) - - union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2 - | zero p1 m2 = Bin p2 m2 (union t1 l2) r2 - | otherwise = Bin p2 m2 l2 (union t1 r2) - -union (Tip k x) t = insert k x t -union t (Tip k x) = insertWith (\_ y -> y) k x t -- right bias -union Nil t = t -union t Nil = t - --- | /O(n+m)/. The union with a combining function. --- --- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] - -unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a -unionWith f m1 m2 - = unionWithKey (\_ x y -> f x y) m1 m2 - --- | /O(n+m)/. The 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 :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a -unionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = union1 - | shorter m2 m1 = union2 - | p1 == p2 = Bin p1 m1 (unionWithKey f l1 l2) (unionWithKey f r1 r2) - | otherwise = join p1 t1 p2 t2 - where - union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2 - | zero p2 m1 = Bin p1 m1 (unionWithKey f l1 t2) r1 - | otherwise = Bin p1 m1 l1 (unionWithKey f r1 t2) - - union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2 - | zero p1 m2 = Bin p2 m2 (unionWithKey f t1 l2) r2 - | otherwise = Bin p2 m2 l2 (unionWithKey f t1 r2) - -unionWithKey f (Tip k x) t = insertWithKey f k x t -unionWithKey f t (Tip k x) = insertWithKey (\k' x' y' -> f k' y' x') k x t -- right bias -unionWithKey _ Nil t = t -unionWithKey _ t Nil = t - -{-------------------------------------------------------------------- - Difference ---------------------------------------------------------------------} --- | /O(n+m)/. Difference between two maps (based on keys). --- --- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b" - -difference :: IntMap a -> IntMap b -> IntMap a -difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = difference1 - | shorter m2 m1 = difference2 - | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2) - | otherwise = t1 - where - difference1 | nomatch p2 p1 m1 = t1 - | zero p2 m1 = bin p1 m1 (difference l1 t2) r1 - | otherwise = bin p1 m1 l1 (difference r1 t2) - - difference2 | nomatch p1 p2 m2 = t1 - | zero p1 m2 = difference t1 l2 - | otherwise = difference t1 r2 - -difference t1@(Tip k _) t2 - | member k t2 = Nil - | otherwise = t1 - -difference Nil _ = Nil -difference t (Tip k _) = delete k t -difference t Nil = t - --- | /O(n+m)/. Difference with a combining function. --- --- > 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 :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a -differenceWith f m1 m2 - = differenceWithKey (\_ x y -> f x y) m1 m2 - --- | /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 :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a -differenceWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = difference1 - | shorter m2 m1 = difference2 - | p1 == p2 = bin p1 m1 (differenceWithKey f l1 l2) (differenceWithKey f r1 r2) - | otherwise = t1 - where - difference1 | nomatch p2 p1 m1 = t1 - | zero p2 m1 = bin p1 m1 (differenceWithKey f l1 t2) r1 - | otherwise = bin p1 m1 l1 (differenceWithKey f r1 t2) - - difference2 | nomatch p1 p2 m2 = t1 - | zero p1 m2 = differenceWithKey f t1 l2 - | otherwise = differenceWithKey f t1 r2 - -differenceWithKey f t1@(Tip k x) t2 - = case lookup k t2 of - Just y -> case f k x y of - Just y' -> Tip k y' - Nothing -> Nil - Nothing -> t1 - -differenceWithKey _ Nil _ = Nil -differenceWithKey f t (Tip k y) = updateWithKey (\k' x -> f k' x y) k t -differenceWithKey _ t Nil = t - - -{-------------------------------------------------------------------- - Intersection ---------------------------------------------------------------------} --- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys). --- --- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a" - -intersection :: IntMap a -> IntMap b -> IntMap a -intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = intersection1 - | shorter m2 m1 = intersection2 - | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2) - | otherwise = Nil - where - intersection1 | nomatch p2 p1 m1 = Nil - | zero p2 m1 = intersection l1 t2 - | otherwise = intersection r1 t2 - - intersection2 | nomatch p1 p2 m2 = Nil - | zero p1 m2 = intersection t1 l2 - | otherwise = intersection t1 r2 - -intersection t1@(Tip k _) t2 - | member k t2 = t1 - | otherwise = Nil -intersection t (Tip k _) - = case lookup k t of - Just y -> Tip k y - Nothing -> Nil -intersection Nil _ = Nil -intersection _ 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" - -intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c -intersectionWith f m1 m2 - = intersectionWithKey (\_ x y -> f x y) m1 m2 - --- | /O(n+m)/. The 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 :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c -intersectionWithKey f t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = intersection1 - | shorter m2 m1 = intersection2 - | p1 == p2 = bin p1 m1 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2) - | otherwise = Nil - where - intersection1 | nomatch p2 p1 m1 = Nil - | zero p2 m1 = intersectionWithKey f l1 t2 - | otherwise = intersectionWithKey f r1 t2 - - intersection2 | nomatch p1 p2 m2 = Nil - | zero p1 m2 = intersectionWithKey f t1 l2 - | otherwise = intersectionWithKey f t1 r2 - -intersectionWithKey f (Tip k x) t2 - = case lookup k t2 of - Just y -> Tip k (f k x y) - Nothing -> Nil -intersectionWithKey f t1 (Tip k y) - = case lookup k t1 of - Just x -> Tip k (f k x y) - Nothing -> Nil -intersectionWithKey _ Nil _ = Nil -intersectionWithKey _ _ Nil = Nil - - -{-------------------------------------------------------------------- - Min\/Max ---------------------------------------------------------------------} - --- | /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 :: (Key -> a -> a) -> IntMap a -> IntMap a -updateMinWithKey f t - = case t of - Bin p m l r | m < 0 -> let t' = updateMinWithKeyUnsigned f r in Bin p m l t' - Bin p m l r -> let t' = updateMinWithKeyUnsigned f l in Bin p m t' r - Tip k y -> Tip k (f k y) - Nil -> error "maxView: empty map has no maximal element" - -updateMinWithKeyUnsigned :: (Key -> a -> a) -> IntMap a -> IntMap a -updateMinWithKeyUnsigned f t - = case t of - Bin p m l r -> let t' = updateMinWithKeyUnsigned f l in Bin p m t' r - Tip k y -> Tip k (f k y) - Nil -> error "updateMinWithKeyUnsigned Nil" - --- | /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 :: (Key -> a -> a) -> IntMap a -> IntMap a -updateMaxWithKey f t - = case t of - Bin p m l r | m < 0 -> let t' = updateMaxWithKeyUnsigned f l in Bin p m t' r - Bin p m l r -> let t' = updateMaxWithKeyUnsigned f r in Bin p m l t' - Tip k y -> Tip k (f k y) - Nil -> error "maxView: empty map has no maximal element" - -updateMaxWithKeyUnsigned :: (Key -> a -> a) -> IntMap a -> IntMap a -updateMaxWithKeyUnsigned f t - = case t of - Bin p m l r -> let t' = updateMaxWithKeyUnsigned f r in Bin p m l t' - Tip k y -> Tip k (f k y) - Nil -> error "updateMaxWithKeyUnsigned Nil" - - --- | /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 :: IntMap a -> Maybe ((Key, a), IntMap a) -maxViewWithKey t - = case t of - Bin p m l r | m < 0 -> let (result, t') = maxViewUnsigned l in Just (result, bin p m t' r) - Bin p m l r -> let (result, t') = maxViewUnsigned r in Just (result, bin p m l t') - Tip k y -> Just ((k,y), Nil) - Nil -> Nothing - -maxViewUnsigned :: IntMap a -> ((Key, a), IntMap a) -maxViewUnsigned t - = case t of - Bin p m l r -> let (result,t') = maxViewUnsigned r in (result,bin p m l t') - Tip k y -> ((k,y), Nil) - Nil -> error "maxViewUnsigned Nil" - --- | /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 :: IntMap a -> Maybe ((Key, a), IntMap a) -minViewWithKey t - = case t of - Bin p m l r | m < 0 -> let (result, t') = minViewUnsigned r in Just (result, bin p m l t') - Bin p m l r -> let (result, t') = minViewUnsigned l in Just (result, bin p m t' r) - Tip k y -> Just ((k,y),Nil) - Nil -> Nothing - -minViewUnsigned :: IntMap a -> ((Key, a), IntMap a) -minViewUnsigned t - = case t of - Bin p m l r -> let (result,t') = minViewUnsigned l in (result,bin p m t' r) - Tip k y -> ((k,y),Nil) - Nil -> error "minViewUnsigned Nil" - - --- | /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 -> a) -> IntMap a -> IntMap a -updateMax f = updateMaxWithKey (const f) - --- | /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 -> a) -> IntMap a -> IntMap a -updateMin f = updateMinWithKey (const f) - --- Similar to the Arrow instance. -first :: (a -> c) -> (a, b) -> (c, b) -first f (x,y) = (f x,y) - --- | /O(log n)/. Retrieves the maximal key of the map, and the map --- stripped of that element, or 'Nothing' if passed an empty map. -maxView :: IntMap a -> Maybe (a, IntMap a) -maxView t = liftM (first snd) (maxViewWithKey t) - --- | /O(log n)/. Retrieves the minimal key of the map, and the map --- stripped of that element, or 'Nothing' if passed an empty map. -minView :: IntMap a -> Maybe (a, IntMap a) -minView t = liftM (first snd) (minViewWithKey t) - --- | /O(log n)/. Delete and find the maximal element. -deleteFindMax :: IntMap a -> (a, IntMap a) -deleteFindMax = fromMaybe (error "deleteFindMax: empty map has no maximal element") . maxView - --- | /O(log n)/. Delete and find the minimal element. -deleteFindMin :: IntMap a -> (a, IntMap a) -deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") . minView - --- | /O(log n)/. The minimal key of the map. -findMin :: IntMap a -> (Key, a) -findMin Nil = error $ "findMin: empty map has no minimal element" -findMin (Tip k v) = (k,v) -findMin (Bin _ m l r) - | m < 0 = go r - | otherwise = go l - where go (Tip k v) = (k,v) - go (Bin _ _ l' _) = go l' - go Nil = error "findMax Nil" - --- | /O(log n)/. The maximal key of the map. -findMax :: IntMap a -> (Key, a) -findMax Nil = error $ "findMax: empty map has no maximal element" -findMax (Tip k v) = (k,v) -findMax (Bin _ m l r) - | m < 0 = go l - | otherwise = go r - where go (Tip k v) = (k,v) - go (Bin _ _ _ r') = go r' - go Nil = error "findMax Nil" - --- | /O(log n)/. Delete the minimal key. An error is thrown if the IntMap is already empty. --- Note, this is not the same behavior Map. -deleteMin :: IntMap a -> IntMap a -deleteMin = maybe (error "deleteMin: empty map has no minimal element") snd . minView - --- | /O(log n)/. Delete the maximal key. An error is thrown if the IntMap is already empty. --- Note, this is not the same behavior Map. -deleteMax :: IntMap a -> IntMap a -deleteMax = maybe (error "deleteMax: empty map has no maximal element") snd . maxView - - -{-------------------------------------------------------------------- - Submap ---------------------------------------------------------------------} --- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). --- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@). -isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool -isProperSubmapOf m1 m2 - = isProperSubmapOfBy (==) m1 m2 - -{- | /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 :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool -isProperSubmapOfBy predicate t1 t2 - = case submapCmp predicate t1 t2 of - LT -> True - _ -> False - -submapCmp :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Ordering -submapCmp predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - | shorter m1 m2 = GT - | shorter m2 m1 = submapCmpLt - | p1 == p2 = submapCmpEq - | otherwise = GT -- disjoint - where - submapCmpLt | nomatch p1 p2 m2 = GT - | zero p1 m2 = submapCmp predicate t1 l2 - | otherwise = submapCmp predicate t1 r2 - submapCmpEq = case (submapCmp predicate l1 l2, submapCmp predicate r1 r2) of - (GT,_ ) -> GT - (_ ,GT) -> GT - (EQ,EQ) -> EQ - _ -> LT - -submapCmp _ (Bin _ _ _ _) _ = GT -submapCmp predicate (Tip kx x) (Tip ky y) - | (kx == ky) && predicate x y = EQ - | otherwise = GT -- disjoint -submapCmp predicate (Tip k x) t - = case lookup k t of - Just y | predicate x y -> LT - _ -> GT -- disjoint -submapCmp _ Nil Nil = EQ -submapCmp _ Nil _ = LT - --- | /O(n+m)/. Is this a submap? --- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). -isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool -isSubmapOf m1 m2 - = isSubmapOfBy (==) m1 m2 - -{- | /O(n+m)/. - The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if - 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': - - > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) - > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) - > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) - - But the following are all 'False': - - > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)]) - > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) - > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)]) --} -isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool -isSubmapOfBy predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - | shorter m1 m2 = False - | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy predicate t1 l2 - else isSubmapOfBy predicate t1 r2) - | otherwise = (p1==p2) && isSubmapOfBy predicate l1 l2 && isSubmapOfBy predicate r1 r2 -isSubmapOfBy _ (Bin _ _ _ _) _ = False -isSubmapOfBy predicate (Tip k x) t = case lookup k t of - Just y -> predicate x y - Nothing -> False -isSubmapOfBy _ Nil _ = True - -{-------------------------------------------------------------------- - 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) -> IntMap a -> IntMap b -map f = mapWithKey (\_ x -> f x) - --- | /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 :: (Key -> a -> b) -> IntMap a -> IntMap b -mapWithKey f t - = case t of - Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r) - Tip k x -> Tip k (f k x) - Nil -> Nil - --- | /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 -> IntMap b -> (a,IntMap c) -mapAccum f = mapAccumWithKey (\a' _ x -> f a' x) - --- | /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 -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap 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 -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) -mapAccumL f a t - = case t of - Bin p m l r -> let (a1,l') = mapAccumL f a l - (a2,r') = mapAccumL f a1 r - in (a2,Bin p m l' r') - 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 --- argument through the map in descending order of keys. -mapAccumRWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) -mapAccumRWithKey f a t - = case t of - Bin p m l r -> let (a1,r') = mapAccumRWithKey f a r - (a2,l') = mapAccumRWithKey f a1 l - in (a2,Bin p m l' r') - Tip k x -> let (a',x') = f a k x in (a',Tip k x') - Nil -> (a,Nil) - -{-------------------------------------------------------------------- - Filter ---------------------------------------------------------------------} --- | /O(n)/. Filter all values that satisfy some 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) -> IntMap a -> IntMap a -filter p m - = filterWithKey (\_ x -> p x) m - --- | /O(n)/. Filter all keys\/values that satisfy some predicate. --- --- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a -filterWithKey predicate t - = case t of - Bin p m l r - -> bin p m (filterWithKey predicate l) (filterWithKey predicate r) - Tip k x - | predicate k x -> t - | otherwise -> Nil - Nil -> Nil - --- | /O(n)/. Partition the map according to some 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) -> IntMap a -> (IntMap a,IntMap a) -partition p m - = partitionWithKey (\_ x -> p x) m - --- | /O(n)/. Partition the map according to some 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 :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a) -partitionWithKey predicate t - = case t of - Bin p m l r - -> let (l1,l2) = partitionWithKey predicate l - (r1,r2) = partitionWithKey predicate r - in (bin p m l1 r1, bin p m l2 r2) - Tip k x - | predicate k x -> (t,Nil) - | otherwise -> (Nil,t) - Nil -> (Nil,Nil) - --- | /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) -> IntMap a -> IntMap 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 :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b -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 -> Tip k y - Nothing -> Nil -mapMaybeWithKey _ Nil = Nil - --- | /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) -> IntMap a -> (IntMap b, IntMap 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 :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) -mapEitherWithKey f (Bin p m l r) - = (bin p m l1 r1, bin p m l2 r2) - where - (l1,l2) = mapEitherWithKey f l - (r1,r2) = mapEitherWithKey f r -mapEitherWithKey f (Tip k x) = case f k x of - Left y -> (Tip k y, Nil) - Right z -> (Nil, Tip k z) -mapEitherWithKey _ Nil = (Nil, Nil) - --- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@ --- where all keys in @map1@ are lower than @k@ and all 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 :: Key -> IntMap a -> (IntMap a,IntMap a) -split k t - = case t of - Bin _ m l r - | m < 0 -> (if k >= 0 -- handle negative numbers. - then let (lt,gt) = split' k l in (union r lt, gt) - else let (lt,gt) = split' k r in (lt, union gt l)) - | otherwise -> split' k t - Tip ky _ - | k>ky -> (t,Nil) - | k (Nil,t) - | otherwise -> (Nil,Nil) - Nil -> (Nil,Nil) - -split' :: Key -> IntMap a -> (IntMap a,IntMap a) -split' k t - = case t of - Bin p m l r - | nomatch k p m -> if k>p then (t,Nil) else (Nil,t) - | zero k m -> let (lt,gt) = split k l in (lt,union gt r) - | otherwise -> let (lt,gt) = split k r in (union l lt,gt) - Tip ky _ - | k>ky -> (t,Nil) - | k (Nil,t) - | otherwise -> (Nil,Nil) - Nil -> (Nil,Nil) - --- | /O(log n)/. Performs a 'split' but also returns whether the pivot --- key was found in the original 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 :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a) -splitLookup k t - = case t of - Bin _ m l r - | m < 0 -> (if k >= 0 -- handle negative numbers. - then let (lt,found,gt) = splitLookup' k l in (union r lt,found, gt) - else let (lt,found,gt) = splitLookup' k r in (lt,found, union gt l)) - | otherwise -> splitLookup' k t - Tip ky y - | k>ky -> (t,Nothing,Nil) - | k (Nil,Nothing,t) - | otherwise -> (Nil,Just y,Nil) - Nil -> (Nil,Nothing,Nil) - -splitLookup' :: Key -> IntMap a -> (IntMap a,Maybe a,IntMap a) -splitLookup' k t - = case t of - Bin p m l r - | nomatch k p m -> if k>p then (t,Nothing,Nil) else (Nil,Nothing,t) - | zero k m -> let (lt,found,gt) = splitLookup k l in (lt,found,union gt r) - | otherwise -> let (lt,found,gt) = splitLookup k r in (union l lt,found,gt) - Tip ky y - | k>ky -> (t,Nothing,Nil) - | k (Nil,Nothing,t) - | otherwise -> (Nil,Just y,Nil) - Nil -> (Nil,Nothing,Nil) - -{-------------------------------------------------------------------- - Fold ---------------------------------------------------------------------} --- | /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. --- --- /Please note that fold will be deprecated in the future and removed./ -fold :: (a -> b -> b) -> b -> IntMap a -> b -fold = foldr -{-# INLINE fold #-} - --- | /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 -> IntMap a -> b -foldr f z t = - case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before - _ -> go z t - where - go z' Nil = z' - go z' (Tip _ x) = f x z' - go z' (Bin _ _ l r) = go (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 -> IntMap a -> b -foldr' f z t = - case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before - _ -> go z t - where - STRICT_1_OF_2(go) - go z' Nil = z' - go z' (Tip _ x) = f x z' - go z' (Bin _ _ l r) = go (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 -> IntMap b -> a -foldl f z t = - case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before - _ -> go z t - where - go z' Nil = z' - go z' (Tip _ x) = f z' x - go z' (Bin _ _ l r) = go (go z' l) 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 -> IntMap b -> a -foldl' f z t = - case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before - _ -> go z t - where - STRICT_1_OF_2(go) - go z' Nil = z' - go z' (Tip _ x) = f z' x - go z' (Bin _ _ l r) = go (go z' l) r -{-# INLINE foldl' #-} - --- | /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. --- --- /Please note that foldWithKey will be deprecated in the future and removed./ -foldWithKey :: (Int -> a -> b -> b) -> b -> IntMap a -> b -foldWithKey = foldrWithKey -{-# INLINE foldWithKey #-} - --- | /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 :: (Int -> a -> b -> b) -> b -> IntMap a -> b -foldrWithKey f z t = - case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before - _ -> go z t - where - 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 #-} - --- | /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' :: (Int -> a -> b -> b) -> b -> IntMap a -> b -foldrWithKey' f z t = - case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before - _ -> go z t - where - STRICT_1_OF_2(go) - 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' #-} - --- | /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 -> Int -> b -> a) -> a -> IntMap b -> a -foldlWithKey f z t = - case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before - _ -> go z t - where - 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 #-} - --- | /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 -> Int -> b -> a) -> a -> IntMap b -> a -foldlWithKey' f z t = - case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before - _ -> go z t - where - STRICT_1_OF_2(go) - 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' #-} - -{-------------------------------------------------------------------- - List variations ---------------------------------------------------------------------} --- | /O(n)/. --- Return all elements of the map in the ascending order of their keys. --- --- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"] --- > elems empty == [] - -elems :: IntMap a -> [a] -elems - = foldr (:) [] - --- | /O(n)/. Return all keys of the map in ascending order. --- --- > keys (fromList [(5,"a"), (3,"b")]) == [3,5] --- > keys empty == [] - -keys :: IntMap a -> [Key] -keys - = foldrWithKey (\k _ ks -> k:ks) [] - --- | /O(n*min(n,W))/. The set of all keys of the map. --- --- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5] --- > keysSet empty == Data.IntSet.empty - -keysSet :: IntMap a -> IntSet.IntSet -keysSet m = IntSet.fromDistinctAscList (keys m) - - --- | /O(n)/. Return all key\/value pairs in the map in ascending key order. --- --- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] --- > assocs empty == [] - -assocs :: IntMap a -> [(Key,a)] -assocs m - = toList m - - -{-------------------------------------------------------------------- - Lists ---------------------------------------------------------------------} --- | /O(n)/. Convert the map to a list of key\/value pairs. --- --- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] --- > toList empty == [] - -toList :: IntMap a -> [(Key,a)] -toList - = foldrWithKey (\k x xs -> (k,x):xs) [] - --- | /O(n)/. Convert the map to a list of key\/value pairs where the --- keys are in ascending order. --- --- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] - -toAscList :: IntMap a -> [(Key,a)] -toAscList t - = -- NOTE: the following algorithm only works for big-endian trees - let (pos,neg) = span (\(k,_) -> k >=0) (foldrWithKey (\k x xs -> (k,x):xs) [] t) in neg ++ pos - --- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs. --- --- > 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")] - -fromList :: [(Key,a)] -> IntMap a -fromList xs - = foldlStrict ins empty xs - where - ins t (k,x) = insert k x t - --- | /O(n*min(n,W))/. Create 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,"c")] == fromList [(3, "ab"), (5, "cba")] --- > fromListWith (++) [] == empty - -fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a -fromListWith f xs - = fromListWithKey (\_ x y -> f x y) xs - --- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'. --- --- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value --- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")] --- > fromListWithKey f [] == empty - -fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromListWithKey f xs - = foldlStrict ins empty xs - where - ins t (k,x) = insertWithKey f k x t - --- | /O(n)/. Build a map from a list of key\/value pairs where --- the keys are in ascending order. --- --- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] --- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] - -fromAscList :: [(Key,a)] -> IntMap a -fromAscList xs - = fromAscListWithKey (\_ x _ -> x) xs - --- | /O(n)/. Build a map from a list of key\/value pairs where --- the keys are in ascending order, with a combining function on equal keys. --- /The precondition (input list is ascending) is not checked./ --- --- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] - -fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWith f xs - = fromAscListWithKey (\_ x y -> f x y) xs - --- | /O(n)/. Build a map from a list of key\/value pairs where --- the keys are in ascending order, with a combining function on equal keys. --- /The precondition (input list is ascending) is not checked./ --- --- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value --- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "5:b|a")] - -fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWithKey _ [] = Nil -fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0) - where - -- [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 combineEq (kx,yy) xs - | otherwise = z:combineEq x xs - --- | /O(n)/. Build a map from a list of key\/value pairs where --- the keys are in ascending order and all distinct. --- /The precondition (input list is strictly ascending) is not checked./ --- --- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] - -#ifdef __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 - -#ifdef __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 - pxy = mask px mxy - in if shorter m mxy - then reduce z zs m pxy (Bin pxy mxy ty tx) stk' - else work z zs (Push px tx stk) - - finish _ t Nada = t - finish px tx (Push py ty stk) = finish p (join py ty px tx) stk - where m = branchMask px py - p = mask px m - -data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada - - -{-------------------------------------------------------------------- - Eq ---------------------------------------------------------------------} -instance Eq a => Eq (IntMap a) where - t1 == t2 = equal t1 t2 - t1 /= t2 = nequal t1 t2 - -equal :: Eq a => IntMap a -> IntMap a -> Bool -equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2) -equal (Tip kx x) (Tip ky y) - = (kx == ky) && (x==y) -equal Nil Nil = True -equal _ _ = False - -nequal :: Eq a => IntMap a -> IntMap a -> Bool -nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2) -nequal (Tip kx x) (Tip ky y) - = (kx /= ky) || (x/=y) -nequal Nil Nil = False -nequal _ _ = True - -{-------------------------------------------------------------------- - Ord ---------------------------------------------------------------------} - -instance Ord a => Ord (IntMap a) where - compare m1 m2 = compare (toList m1) (toList m2) - -{-------------------------------------------------------------------- - Functor ---------------------------------------------------------------------} - -instance Functor IntMap where - fmap = map - -{-------------------------------------------------------------------- - Show ---------------------------------------------------------------------} - -instance Show a => Show (IntMap a) where - showsPrec d m = showParen (d > 10) $ - showString "fromList " . shows (toList m) - -{- -XXX unused code - -showMap :: (Show a) => [(Key,a)] -> ShowS -showMap [] - = showString "{}" -showMap (x:xs) - = showChar '{' . showElem x . showTail xs - where - showTail [] = showChar '}' - showTail (x':xs') = showChar ',' . showElem x' . showTail xs' - - showElem (k,v) = shows k . showString ":=" . shows v --} - -{-------------------------------------------------------------------- - Read ---------------------------------------------------------------------} -instance (Read e) => Read (IntMap 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 - -{-------------------------------------------------------------------- - Typeable ---------------------------------------------------------------------} - -#include "Typeable.h" -INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap") - -{-------------------------------------------------------------------- - Debugging ---------------------------------------------------------------------} --- | /O(n)/. Show the tree that implements the map. The tree is shown --- in a compressed, hanging format. -showTree :: Show a => IntMap a -> String -showTree s - = showTreeWith True False s - - -{- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows - the tree that implements the map. 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. --} -showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String -showTreeWith hang wide t - | hang = (showsTreeHang wide [] t) "" - | otherwise = (showsTree wide [] [] t) "" - -showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS -showsTree wide lbars rbars t - = case t of - Bin p m l r - -> showsTree wide (withBar rbars) (withEmpty rbars) r . - showWide wide rbars . - showsBars lbars . showString (showBin p m) . showString "\n" . - showWide wide lbars . - showsTree wide (withEmpty lbars) (withBar lbars) l - Tip k x - -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n" - Nil -> showsBars lbars . showString "|\n" - -showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS -showsTreeHang wide bars t - = case t of - Bin p m l r - -> showsBars bars . showString (showBin p m) . showString "\n" . - showWide wide bars . - showsTreeHang wide (withBar bars) l . - showWide wide bars . - showsTreeHang wide (withEmpty bars) r - Tip k x - -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n" - Nil -> showsBars bars . showString "|\n" - -showBin :: Prefix -> Mask -> String -showBin _ _ - = "*" -- ++ show (p,m) - -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 - - -{-------------------------------------------------------------------- - Helpers ---------------------------------------------------------------------} -{-------------------------------------------------------------------- - Join ---------------------------------------------------------------------} -join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a -join p1 t1 p2 t2 - | zero p1 m = Bin p m t1 t2 - | otherwise = Bin p m t2 t1 - where - m = branchMask p1 p2 - p = mask p1 m -{-# INLINE join #-} - -{-------------------------------------------------------------------- - @bin@ assures that we never have empty trees within a tree. ---------------------------------------------------------------------} -bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a -bin _ _ l Nil = l -bin _ _ Nil r = r -bin p m l r = Bin p m l r -{-# INLINE bin #-} - - -{-------------------------------------------------------------------- - Endian independent bit twiddling ---------------------------------------------------------------------} -zero :: Key -> Mask -> Bool -zero i m - = (natFromInt i) .&. (natFromInt m) == 0 -{-# INLINE zero #-} - -nomatch,match :: Key -> Prefix -> Mask -> Bool -nomatch i p m - = (mask i m) /= p -{-# INLINE nomatch #-} - -match i p m - = (mask i m) == p -{-# INLINE match #-} - -mask :: Key -> Mask -> Prefix -mask i m - = maskW (natFromInt i) (natFromInt m) -{-# INLINE mask #-} - - -{-------------------------------------------------------------------- - Big endian operations ---------------------------------------------------------------------} -maskW :: Nat -> Nat -> Prefix -maskW i m - = intFromNat (i .&. (complement (m-1) `xor` m)) -{-# INLINE maskW #-} - -shorter :: Mask -> Mask -> Bool -shorter m1 m2 - = (natFromInt m1) > (natFromInt m2) -{-# INLINE shorter #-} - -branchMask :: Prefix -> Prefix -> Mask -branchMask p1 p2 - = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2)) -{-# INLINE branchMask #-} - -{---------------------------------------------------------------------- - Finding the highest bit (mask) in a word [x] can be done efficiently in - three ways: - * convert to a floating point value and the mantissa tells us the - [log2(x)] that corresponds with the highest bit position. The mantissa - is retrieved either via the standard C function [frexp] or by some bit - twiddling on IEEE compatible numbers (float). Note that one needs to - use at least [double] precision for an accurate mantissa of 32 bit - numbers. - * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit). - * use processor specific assembler instruction (asm). - - The most portable way would be [bit], but is it efficient enough? - I have measured the cycle counts of the different methods on an AMD - Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction: - - highestBitMask: method cycles - -------------- - frexp 200 - float 33 - bit 11 - asm 12 - - highestBit: method cycles - -------------- - frexp 195 - float 33 - bit 11 - asm 11 - - Wow, the bit twiddling is on today's RISC like machines even faster - than a single CISC instruction (BSR)! -----------------------------------------------------------------------} - -{---------------------------------------------------------------------- - [highestBitMask] returns a word where only the highest bit is set. - It is found by first setting all bits in lower positions than the - highest bit and than taking an exclusive or with the original value. - Allthough the function may look expensive, GHC compiles this into - excellent C code that subsequently compiled into highly efficient - machine code. The algorithm is derived from Jorg Arndt's FXT library. -----------------------------------------------------------------------} -highestBitMask :: Nat -> Nat -highestBitMask x0 - = case (x0 .|. shiftRL x0 1) of - x1 -> case (x1 .|. shiftRL x1 2) of - x2 -> case (x2 .|. shiftRL x2 4) of - x3 -> case (x3 .|. shiftRL x3 8) of - x4 -> case (x4 .|. shiftRL x4 16) of - x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms - x6 -> (x6 `xor` (shiftRL x6 1)) -{-# INLINE highestBitMask #-} - - -{-------------------------------------------------------------------- - Utilities ---------------------------------------------------------------------} - -foldlStrict :: (a -> b -> a) -> a -> [b] -> a -foldlStrict f = go - where - go z [] = z - go z (x:xs) = let z' = f z x in z' `seq` go z' xs -{-# INLINE foldlStrict #-} +{-# LANGUAGE CPP #-} +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.IntMap +-- 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 integer keys to values +-- (dictionaries). +-- +-- This module re-exports the value lazy 'Data.IntMap.Lazy' API, plus +-- several value strict functions from 'Data.IntMap.Strict'. +-- +-- These modules are intended to be imported qualified, to avoid name +-- clashes with Prelude functions, e.g. +-- +-- > import Data.IntMap (IntMap) +-- > import qualified Data.IntMap as IntMap +-- +-- The implementation is based on /big-endian patricia trees/. This data +-- structure performs especially well on binary operations like 'union' +-- and 'intersection'. However, my benchmarks show that it is also +-- (much) faster on insertions and deletions when compared to a generic +-- size-balanced map implementation (see "Data.Map"). +-- +-- * 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), +-- October 1968, pages 514-534. +-- +-- Operation comments contain the operation time complexity in +-- the Big-O notation . +-- Many operations have a worst-case complexity of /O(min(n,W))/. +-- This means that the operation can become linear in the number of +-- elements with a maximum of /W/ -- the number of bits in an 'Int' +-- (32 or 64). +----------------------------------------------------------------------------- + +module Data.IntMap + ( module Data.IntMap.Lazy + , insertWith' + , insertWithKey' + , fold + , foldWithKey + ) where + +import Prelude hiding (lookup,map,filter,foldr,foldl,null) +import Data.IntMap.Lazy +import qualified Data.IntMap.Strict as S + +-- | /Deprecated./ As of version 0.5, replaced by 'S.insertWith'. +-- +-- /O(log n)/. Same as 'insertWith', but the combining function is +-- applied strictly. This function is deprecated, use 'insertWith' in +-- "Data.IntMap.Strict" instead. +insertWith' :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWith' = S.insertWith +{-# INLINE insertWith' #-} + +-- | /Deprecated./ As of version 0.5, replaced by 'S.insertWithKey'. +-- +-- /O(log n)/. Same as 'insertWithKey', but the combining function is +-- applied strictly. This function is deprecated, use 'insertWithKey' +-- in "Data.IntMap.Strict" instead. +insertWithKey' :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWithKey' = S.insertWithKey +{-# INLINE insertWithKey' #-} + +-- | /Deprecated./ As of version 0.5, replaced by 'foldr'. +-- +-- /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. +fold :: (a -> b -> b) -> b -> IntMap a -> b +fold = foldr +{-# INLINE fold #-} + +-- | /Deprecated./ As of version 0.5, replaced by 'foldrWithKey'. +-- +-- /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. +foldWithKey :: (Int -> a -> b -> b) -> b -> IntMap a -> b +foldWithKey = foldrWithKey +{-# INLINE foldWithKey #-} diff --git a/Data/IntMap/Base.hs b/Data/IntMap/Base.hs new file mode 100644 --- /dev/null +++ b/Data/IntMap/Base.hs @@ -0,0 +1,2171 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE MagicHash, DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.IntMap.Base +-- Copyright : (c) Daan Leijen 2002 +-- (c) Andriy Palamarchuk 2008 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- This defines the data structures and core (hidden) manipulations +-- on representations. +----------------------------------------------------------------------------- + +-- [Note: INLINE bit fiddling] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- It is essential that the bit fiddling functions like mask, zero, branchMask +-- etc are inlined. If they do not, the memory allocation skyrockets. The GHC +-- usually gets it right, but it is disastrous if it does not. Therefore we +-- explicitly mark these functions INLINE. + + +-- [Note: Local 'go' functions and capturing] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- Care must be taken when using 'go' function which captures an argument. +-- Sometimes (for example when the argument is passed to a data constructor, +-- as in insert), GHC heap-allocates more than necessary. Therefore C-- code +-- must be checked for increased allocation when creating and modifying such +-- functions. + + +-- [Note: Order of constructors] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- The order of constructors of IntMap matters when considering performance. +-- Currently in GHC 7.0, when type has 3 constructors, they are matched from +-- the first to the last -- the best performance is achieved when the +-- constructors are ordered by frequency. +-- On GHC 7.0, reordering constructors from Nil | Tip | Bin to Bin | Tip | Nil +-- improves the benchmark by circa 10%. + +module Data.IntMap.Base ( + -- * Map type + IntMap(..), Key -- instance Eq,Show + + -- * 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 + , mergeWithKey' + + -- * Traversal + -- ** Map + , map + , mapWithKey + , traverseWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeys + , mapKeysWith + , mapKeysMonotonic + + -- * Folds + , foldr + , foldl + , foldrWithKey + , foldlWithKey + -- ** 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 + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + + -- * Internal types + , Mask, Prefix, Nat + + -- * Utility + , natFromInt + , intFromNat + , shiftRL + , shiftLL + , join + , bin + , zero + , nomatch + , match + , mask + , maskW + , shorter + , branchMask + , highestBitMask + , foldlStrict + ) where + +import Data.Bits + +import Prelude hiding (lookup,map,filter,foldr,foldl,null) +import qualified Data.IntSet.Base as IntSet +import Data.Monoid (Monoid(..)) +import Data.Maybe (fromMaybe) +import Data.Typeable +import qualified Data.Foldable as Foldable +import Data.Traversable (Traversable(traverse)) +import Control.Applicative (Applicative(pure,(<*>)),(<$>)) +import Control.Monad ( liftM ) +import Control.DeepSeq (NFData(rnf)) + +#if __GLASGOW_HASKELL__ +import Text.Read +import Data.Data (Data(..), mkNoRepType) +#endif + +#if __GLASGOW_HASKELL__ +import GHC.Exts ( Word(..), Int(..), build ) +import GHC.Prim ( uncheckedShiftL#, uncheckedShiftRL# ) +#else +import Data.Word +#endif + +-- On GHC, include MachDeps.h to get WORD_SIZE_IN_BITS macro. +#if defined(__GLASGOW_HASKELL__) +#include "MachDeps.h" +#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 + +-- A "Nat" is a natural machine word (an unsigned Int) +type Nat = Word + +natFromInt :: Key -> Nat +natFromInt = fromIntegral +{-# INLINE natFromInt #-} + +intFromNat :: Nat -> Key +intFromNat = fromIntegral +{-# INLINE intFromNat #-} + +-- Right and left logical shifts. +shiftRL, shiftLL :: Nat -> Key -> Nat +#if __GLASGOW_HASKELL__ +{-------------------------------------------------------------------- + GHC: use unboxing to get @shiftRL@ inlined. +--------------------------------------------------------------------} +shiftRL (W# x) (I# i) = W# (uncheckedShiftRL# x i) +shiftLL (W# x) (I# i) = W# (uncheckedShiftL# x i) +#else +shiftRL x i = shiftR x i +shiftLL x i = shiftL x i +#endif +{-# INLINE shiftRL #-} +{-# INLINE shiftLL #-} + +{-------------------------------------------------------------------- + Types +--------------------------------------------------------------------} + + +-- | A map of integers to values @a@. + +-- See Note: Order of constructors +data IntMap a = Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(IntMap a) !(IntMap a) + | Tip {-# UNPACK #-} !Key a + | Nil + +type Prefix = Int +type Mask = Int +type Key = Int + +{-------------------------------------------------------------------- + Operators +--------------------------------------------------------------------} + +-- | /O(min(n,W))/. 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' + +(!) :: IntMap a -> Key -> a +m ! k = find k m + +-- | Same as 'difference'. +(\\) :: IntMap a -> IntMap b -> IntMap a +m1 \\ m2 = difference m1 m2 + +infixl 9 \\{-This comment teaches CPP correct behaviour -} + +{-------------------------------------------------------------------- + Types +--------------------------------------------------------------------} + +instance Monoid (IntMap a) where + mempty = empty + mappend = union + mconcat = unions + +instance Foldable.Foldable IntMap where + fold Nil = mempty + fold (Tip _ v) = v + fold (Bin _ _ l r) = Foldable.fold l `mappend` Foldable.fold r + foldr = foldr + foldl = foldl + foldMap _ Nil = mempty + foldMap f (Tip _k v) = f v + foldMap f (Bin _ _ l r) = Foldable.foldMap f l `mappend` Foldable.foldMap f r + +instance Traversable IntMap where + traverse f = traverseWithKey (\_ -> f) + +instance NFData a => NFData (IntMap a) where + rnf Nil = () + rnf (Tip _ v) = rnf v + rnf (Bin _ _ l r) = rnf l `seq` rnf r + +#if __GLASGOW_HASKELL__ + +{-------------------------------------------------------------------- + A Data instance +--------------------------------------------------------------------} + +-- This instance preserves data abstraction at the cost of inefficiency. +-- We omit reflection services for the sake of data abstraction. + +instance Data a => Data (IntMap a) where + gfoldl f z im = z fromList `f` (toList im) + toConstr _ = error "toConstr" + gunfold _ _ = error "gunfold" + dataTypeOf _ = mkNoRepType "Data.IntMap.IntMap" + dataCast1 f = gcast1 f + +#endif + +{-------------------------------------------------------------------- + Query +--------------------------------------------------------------------} +-- | /O(1)/. Is the map empty? +-- +-- > Data.IntMap.null (empty) == True +-- > Data.IntMap.null (singleton 1 'a') == False + +null :: IntMap a -> Bool +null Nil = True +null _ = False +{-# INLINE null #-} + +-- | /O(n)/. Number of elements in the map. +-- +-- > size empty == 0 +-- > 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 + +-- | /O(min(n,W))/. Is the key a member of the map? +-- +-- > member 5 (fromList [(5,'a'), (3,'b')]) == True +-- > member 1 (fromList [(5,'a'), (3,'b')]) == False + +-- See Note: Local 'go' functions and capturing] +member :: Key -> IntMap a -> Bool +member k = k `seq` go + where + go (Bin p m l r) | nomatch k p m = False + | zero k m = go l + | otherwise = go r + go (Tip kx _) = k == kx + go Nil = False + +-- | /O(min(n,W))/. Is the key not a member of the map? +-- +-- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False +-- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True + +notMember :: Key -> IntMap a -> Bool +notMember k m = not $ member k m + +-- | /O(min(n,W))/. Lookup the value at a key in the map. See also 'Data.Map.lookup'. + +-- See Note: Local 'go' functions and capturing] +lookup :: Key -> IntMap a -> Maybe a +lookup k = k `seq` go + where + go (Bin p m l r) | nomatch k p m = Nothing + | zero k m = go l + | otherwise = go r + go (Tip kx x) | k == kx = Just x + | otherwise = Nothing + go Nil = Nothing + + +-- See Note: Local 'go' functions and capturing] +find :: Key -> IntMap a -> a +find k = k `seq` go + where + go (Bin p m l r) | nomatch k p m = not_found + | zero k m = go l + | otherwise = go r + go (Tip kx x) | k == kx = x + | otherwise = not_found + go Nil = not_found + + not_found = error ("IntMap.!: key " ++ show k ++ " is not an element of the map") + +-- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@ +-- returns the value at key @k@ or returns @def@ when the key is not an +-- element of the map. +-- +-- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' +-- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' + +-- See Note: Local 'go' functions and capturing] +findWithDefault :: a -> Key -> IntMap a -> a +findWithDefault def k = k `seq` go + where + go (Bin p m l r) | nomatch k p m = def + | zero k m = go l + | otherwise = go r + go (Tip kx x) | k == kx = x + | otherwise = def + go Nil = def + +-- | /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') + +-- See Note: Local 'go' functions and capturing. +lookupLT :: Key -> IntMap a -> Maybe (Key, a) +lookupLT k t = k `seq` case t of + Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r + _ -> go Nil t + where + go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r + | zero k m = go def l + | otherwise = go l r + go def (Tip ky y) | k <= ky = unsafeFindMax def + | otherwise = Just (ky, y) + go def Nil = unsafeFindMax def + +-- | /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 + +-- See Note: Local 'go' functions and capturing. +lookupGT :: Key -> IntMap a -> Maybe (Key, a) +lookupGT k t = k `seq` case t of + Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r + _ -> go Nil t + where + go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def + | zero k m = go r l + | otherwise = go def r + go def (Tip ky y) | k >= ky = unsafeFindMin def + | otherwise = Just (ky, y) + go def Nil = unsafeFindMin def + +-- | /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') + +-- See Note: Local 'go' functions and capturing. +lookupLE :: Key -> IntMap a -> Maybe (Key, a) +lookupLE k t = k `seq` case t of + Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r + _ -> go Nil t + where + go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r + | zero k m = go def l + | otherwise = go l r + go def (Tip ky y) | k < ky = unsafeFindMax def + | otherwise = Just (ky, y) + go def Nil = unsafeFindMax def + +-- | /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 + +-- See Note: Local 'go' functions and capturing. +lookupGE :: Key -> IntMap a -> Maybe (Key, a) +lookupGE k t = k `seq` case t of + Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r + _ -> go Nil t + where + go def (Bin p m l r) | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def + | zero k m = go r l + | otherwise = go def r + go def (Tip ky y) | k > ky = unsafeFindMin def + | otherwise = Just (ky, y) + go def Nil = unsafeFindMin def + + +-- Helper function for lookupGE and lookupGT. It assumes that if a Bin node is +-- given, it has m > 0. +unsafeFindMin :: IntMap a -> Maybe (Key, a) +unsafeFindMin Nil = Nothing +unsafeFindMin (Tip ky y) = Just (ky, y) +unsafeFindMin (Bin _ _ l _) = unsafeFindMin l + +-- Helper function for lookupLE and lookupLT. It assumes that if a Bin node is +-- given, it has m > 0. +unsafeFindMax :: IntMap a -> Maybe (Key, a) +unsafeFindMax Nil = Nothing +unsafeFindMax (Tip ky y) = Just (ky, y) +unsafeFindMax (Bin _ _ _ r) = unsafeFindMax r + +{-------------------------------------------------------------------- + Construction +--------------------------------------------------------------------} +-- | /O(1)/. The empty map. +-- +-- > empty == fromList [] +-- > size empty == 0 + +empty :: IntMap a +empty + = Nil +{-# INLINE empty #-} + +-- | /O(1)/. A map of one element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +-- > size (singleton 1 'a') == 1 + +singleton :: Key -> a -> IntMap a +singleton k x + = Tip k x +{-# INLINE singleton #-} + +{-------------------------------------------------------------------- + Insert +--------------------------------------------------------------------} +-- | /O(min(n,W))/. Insert a new key\/value pair in the map. +-- If the key is already present in the map, the associated value is +-- replaced with the supplied value, i.e. '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' + +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 -> join 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 -> join k (Tip k x) ky t + Nil -> Tip k x + +-- right-biased insertion, used by 'union' +-- | /O(min(n,W))/. Insert with a combining function. +-- @'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 @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 :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWith f k x t + = insertWithKey (\_ x' y' -> f x' y') k x t + +-- | /O(min(n,W))/. Insert with a combining function. +-- @'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 @f key new_value old_value@. +-- +-- > 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" + +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 -> join 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 -> join k (Tip k x) ky t + 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@) +-- 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")]) + +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,join 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,join k (Tip k x) ky t) + Nil -> (Nothing,Tip k x) + + +{-------------------------------------------------------------------- + Deletion +--------------------------------------------------------------------} +-- | /O(min(n,W))/. 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 + +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 + +-- | /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. +-- +-- > 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 :: (a -> a) -> Key -> IntMap a -> IntMap a +adjust f k m + = adjustWithKey (\_ x -> f x) k m + +-- | /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. +-- +-- > 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 :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a +adjustWithKey f + = updateWithKey (\k' x -> Just (f k' x)) + +-- | /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@. +-- +-- > 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 :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a +update f + = updateWithKey (\_ x -> f x) + +-- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +updateWithKey :: (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 + Just y' -> Tip ky y' + Nothing -> Nil + | otherwise -> t + Nil -> Nil + +-- | /O(min(n,W))/. Lookup and update. +-- The function returns original value, if it is updated. +-- This is different behavior than 'Data.Map.updateLookupWithKey'. +-- 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 "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") + +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) + + + +-- | /O(min(n,W))/. 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 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 -> join 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 -> join k (Tip k x) ky t + Nothing -> Tip ky y + Nil -> case f Nothing of + Just x -> Tip k x + Nothing -> Nil + + +{-------------------------------------------------------------------- + Union +--------------------------------------------------------------------} +-- | The union of a list of maps. +-- +-- > 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 :: [IntMap a] -> IntMap a +unions xs + = foldlStrict union empty xs + +-- | The union of a list of maps, with a combining operation. +-- +-- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] +-- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] + +unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a +unionsWith f ts + = foldlStrict (unionWith f) empty ts + +-- | /O(n+m)/. The (left-biased) union of two maps. +-- It prefers the first map 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 :: IntMap a -> IntMap a -> IntMap a +union m1 m2 + = mergeWithKey' Bin const id id m1 m2 + +-- | /O(n+m)/. The union with a combining function. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] + +unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWith f m1 m2 + = unionWithKey (\_ x y -> f x y) m1 m2 + +-- | /O(n+m)/. The 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 :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWithKey f m1 m2 + = mergeWithKey' Bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 (f k1 x1 x2)) id id m1 m2 + +{-------------------------------------------------------------------- + Difference +--------------------------------------------------------------------} +-- | /O(n+m)/. Difference between two maps (based on keys). +-- +-- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b" + +difference :: IntMap a -> IntMap b -> IntMap a +difference m1 m2 + = mergeWithKey (\_ _ _ -> Nothing) id (const Nil) m1 m2 + +-- | /O(n+m)/. Difference with a combining function. +-- +-- > 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 :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a +differenceWith f m1 m2 + = differenceWithKey (\_ x y -> f x y) m1 m2 + +-- | /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 :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a +differenceWithKey f m1 m2 + = mergeWithKey f id (const Nil) m1 m2 + + +{-------------------------------------------------------------------- + Intersection +--------------------------------------------------------------------} +-- | /O(n+m)/. The (left-biased) intersection of two maps (based on keys). +-- +-- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a" + +intersection :: IntMap a -> IntMap b -> IntMap a +intersection m1 m2 + = mergeWithKey' bin const (const Nil) (const Nil) m1 m2 + +-- | /O(n+m)/. The intersection with a combining function. +-- +-- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" + +intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c +intersectionWith f m1 m2 + = intersectionWithKey (\_ x y -> f x y) m1 m2 + +-- | /O(n+m)/. The 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 :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c +intersectionWithKey f m1 m2 + = mergeWithKey' bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 (f k1 x1 x2)) (const Nil) (const Nil) m1 m2 + +{-------------------------------------------------------------------- + MergeWithKey +--------------------------------------------------------------------} + +-- | /O(n+m)/. A high-performance universal combining function. Using +-- 'mergeWithKey', all combining functions can be defined without any loss of +-- efficiency (with exception of 'union', 'difference' and 'intersection', +-- where sharing of some nodes is lost with 'mergeWithKey'). +-- +-- 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 :: (Key -> a -> b -> Maybe c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c) + -> IntMap a -> IntMap b -> IntMap c +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 -> Tip k1 x + {-# INLINE combine #-} +{-# INLINE mergeWithKey #-} + +-- Slightly more general version of mergeWithKey. It differs in the following: +-- +-- * the combining function operates on maps instead of keys and values. The +-- reason is to enable sharing in union, difference and intersection. +-- +-- * mergeWithKey' is given an equivalent of bin. The reason is that in union*, +-- Bin constructor can be used, because we know both subtrees are nonempty. + +mergeWithKey' :: (Prefix -> Mask -> IntMap c -> IntMap c -> IntMap c) + -> (IntMap a -> IntMap b -> IntMap c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c) + -> IntMap a -> IntMap b -> IntMap c +mergeWithKey' bin' f g1 g2 = go + where + go t1@(Bin p1 m1 l1 r1) t2@(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 = maybe_join p1 (g1 t1) p2 (g2 t2) + where + merge1 | nomatch p2 p1 m1 = maybe_join p1 (g1 t1) p2 (g2 t2) + | zero p2 m1 = bin' p1 m1 (go l1 t2) (g1 r1) + | otherwise = bin' p1 m1 (g1 l1) (go r1 t2) + merge2 | nomatch p1 p2 m2 = maybe_join p1 (g1 t1) p2 (g2 t2) + | zero p1 m2 = bin' p2 m2 (go t1 l2) (g2 r2) + | otherwise = bin' p2 m2 (g2 l2) (go t1 r2) + + go t1'@(Bin _ _ _ _) t2'@(Tip k2' _) = merge t2' k2' t1' + where merge t2 k2 t1@(Bin p1 m1 l1 r1) | nomatch k2 p1 m1 = maybe_join p1 (g1 t1) k2 (g2 t2) + | zero k2 m1 = bin' p1 m1 (merge t2 k2 l1) (g1 r1) + | otherwise = bin' p1 m1 (g1 l1) (merge t2 k2 r1) + merge t2 k2 t1@(Tip k1 _) | k1 == k2 = f t1 t2 + | otherwise = maybe_join k1 (g1 t1) k2 (g2 t2) + merge t2 _ Nil = g2 t2 + + go t1@(Bin _ _ _ _) Nil = g1 t1 + + go t1'@(Tip k1' _) t2' = merge t1' k1' t2' + where merge t1 k1 t2@(Bin p2 m2 l2 r2) | nomatch k1 p2 m2 = maybe_join k1 (g1 t1) p2 (g2 t2) + | zero k1 m2 = bin' p2 m2 (merge t1 k1 l2) (g2 r2) + | otherwise = bin' p2 m2 (g2 l2) (merge t1 k1 r2) + merge t1 k1 t2@(Tip k2 _) | k1 == k2 = f t1 t2 + | otherwise = maybe_join k1 (g1 t1) k2 (g2 t2) + merge t1 _ Nil = g1 t1 + + go Nil t2 = g2 t2 + + maybe_join _ Nil _ t2 = t2 + maybe_join _ t1 _ Nil = t1 + maybe_join p1 t1 p2 t2 = join p1 t1 p2 t2 + {-# INLINE maybe_join #-} +{-# INLINE mergeWithKey' #-} + +{-------------------------------------------------------------------- + Min\/Max +--------------------------------------------------------------------} + +-- | /O(min(n,W))/. 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 :: (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) + _ -> go f t + where + go f' (Bin p m l r) = bin p m (go f' l) r + go f' (Tip k y) = case f' k y of + Just y' -> Tip k y' + Nothing -> Nil + go _ Nil = error "updateMinWithKey Nil" + +-- | /O(min(n,W))/. 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 :: (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 + _ -> go f t + where + go f' (Bin p m l r) = bin p m l (go f' r) + go f' (Tip k y) = case f' k y of + Just y' -> Tip k y' + Nothing -> Nil + go _ Nil = error "updateMaxWithKey Nil" + +-- | /O(min(n,W))/. 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 :: 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) + _ -> Just (go t) + where + go (Bin p m l r) = case go r of (result, r') -> (result, bin p m l r') + go (Tip k y) = ((k, y), Nil) + go Nil = error "maxViewWithKey Nil" + +-- | /O(min(n,W))/. 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 :: 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') + _ -> Just (go t) + where + go (Bin p m l r) = case go l of (result, l') -> (result, bin p m l' r) + go (Tip k y) = ((k, y), Nil) + go Nil = error "minViewWithKey Nil" + +-- | /O(min(n,W))/. 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) -> IntMap a -> IntMap a +updateMax f = updateMaxWithKey (const f) + +-- | /O(min(n,W))/. 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) -> IntMap a -> IntMap a +updateMin f = updateMinWithKey (const f) + +-- Similar to the Arrow instance. +first :: (a -> c) -> (a, b) -> (c, b) +first f (x,y) = (f x,y) + +-- | /O(min(n,W))/. Retrieves the maximal key of the map, and the map +-- stripped of that element, or 'Nothing' if passed an empty map. +maxView :: IntMap a -> Maybe (a, IntMap a) +maxView t = liftM (first snd) (maxViewWithKey t) + +-- | /O(min(n,W))/. Retrieves the minimal key of the map, and the map +-- stripped of that element, or 'Nothing' if passed an empty map. +minView :: IntMap a -> Maybe (a, IntMap a) +minView t = liftM (first snd) (minViewWithKey t) + +-- | /O(min(n,W))/. Delete and find the maximal element. +deleteFindMax :: IntMap a -> ((Key, a), IntMap a) +deleteFindMax = fromMaybe (error "deleteFindMax: empty map has no maximal element") . maxViewWithKey + +-- | /O(min(n,W))/. Delete and find the minimal element. +deleteFindMin :: IntMap a -> ((Key, a), IntMap a) +deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") . minViewWithKey + +-- | /O(min(n,W))/. The minimal key of the map. +findMin :: IntMap a -> (Key, a) +findMin Nil = error $ "findMin: empty map has no minimal element" +findMin (Tip k v) = (k,v) +findMin (Bin _ m l r) + | m < 0 = go r + | otherwise = go l + where go (Tip k v) = (k,v) + go (Bin _ _ l' _) = go l' + go Nil = error "findMax Nil" + +-- | /O(min(n,W))/. The maximal key of the map. +findMax :: IntMap a -> (Key, a) +findMax Nil = error $ "findMax: empty map has no maximal element" +findMax (Tip k v) = (k,v) +findMax (Bin _ m l r) + | m < 0 = go l + | otherwise = go r + where go (Tip k v) = (k,v) + go (Bin _ _ _ r') = go r' + go Nil = error "findMax Nil" + +-- | /O(min(n,W))/. Delete the minimal key. An error is thrown if the IntMap is already empty. +-- Note, this is not the same behavior Map. +deleteMin :: IntMap a -> IntMap a +deleteMin = maybe Nil snd . minView + +-- | /O(min(n,W))/. Delete the maximal key. An error is thrown if the IntMap is already empty. +-- Note, this is not the same behavior Map. +deleteMax :: IntMap a -> IntMap a +deleteMax = maybe Nil snd . maxView + + +{-------------------------------------------------------------------- + Submap +--------------------------------------------------------------------} +-- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). +-- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@). +isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool +isProperSubmapOf m1 m2 + = isProperSubmapOfBy (==) m1 m2 + +{- | /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 :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool +isProperSubmapOfBy predicate t1 t2 + = case submapCmp predicate t1 t2 of + LT -> True + _ -> False + +submapCmp :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Ordering +submapCmp predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) + | shorter m1 m2 = GT + | shorter m2 m1 = submapCmpLt + | p1 == p2 = submapCmpEq + | otherwise = GT -- disjoint + where + submapCmpLt | nomatch p1 p2 m2 = GT + | zero p1 m2 = submapCmp predicate t1 l2 + | otherwise = submapCmp predicate t1 r2 + submapCmpEq = case (submapCmp predicate l1 l2, submapCmp predicate r1 r2) of + (GT,_ ) -> GT + (_ ,GT) -> GT + (EQ,EQ) -> EQ + _ -> LT + +submapCmp _ (Bin _ _ _ _) _ = GT +submapCmp predicate (Tip kx x) (Tip ky y) + | (kx == ky) && predicate x y = EQ + | otherwise = GT -- disjoint +submapCmp predicate (Tip k x) t + = case lookup k t of + Just y | predicate x y -> LT + _ -> GT -- disjoint +submapCmp _ Nil Nil = EQ +submapCmp _ Nil _ = LT + +-- | /O(n+m)/. Is this a submap? +-- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). +isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool +isSubmapOf m1 m2 + = isSubmapOfBy (==) m1 m2 + +{- | /O(n+m)/. + The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if + 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': + + > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) + + But the following are all 'False': + + > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)]) + > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)]) +-} +isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool +isSubmapOfBy predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) + | shorter m1 m2 = False + | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubmapOfBy predicate t1 l2 + else isSubmapOfBy predicate t1 r2) + | otherwise = (p1==p2) && isSubmapOfBy predicate l1 l2 && isSubmapOfBy predicate r1 r2 +isSubmapOfBy _ (Bin _ _ _ _) _ = False +isSubmapOfBy predicate (Tip k x) t = case lookup k t of + Just y -> predicate x y + Nothing -> False +isSubmapOfBy _ Nil _ = True + +{-------------------------------------------------------------------- + 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) -> 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 + +-- | /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 :: (Key -> a -> b) -> IntMap a -> IntMap b +mapWithKey f t + = case t of + Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r) + Tip k x -> Tip k (f k x) + Nil -> Nil + +-- | /O(n)/. +-- @'traverseWithKey' f s == '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 +{-# INLINE traverseWithKey #-} +traverseWithKey :: Applicative t => (Key -> a -> t b) -> IntMap a -> t (IntMap b) +traverseWithKey f = go + where + go Nil = pure Nil + go (Tip k v) = Tip k <$> f k v + go (Bin p m l r) = Bin p m <$> go l <*> go r + +-- | /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 -> IntMap b -> (a,IntMap c) +mapAccum f = mapAccumWithKey (\a' _ x -> f a' x) + +-- | /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 -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap 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 -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) +mapAccumL f a t + = case t of + Bin p m l r -> let (a1,l') = mapAccumL f a l + (a2,r') = mapAccumL f a1 r + in (a2,Bin p m l' r') + 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 +-- argument through the map in descending order of keys. +mapAccumRWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) +mapAccumRWithKey f a t + = case t of + Bin p m l r -> let (a1,r') = mapAccumRWithKey f a r + (a2,l') = mapAccumRWithKey f a1 l + in (a2,Bin p m l' r') + Tip k x -> let (a',x') = f a k x in (a',Tip k x') + Nil -> (a,Nil) + +-- | /O(n*min(n,W))/. +-- @'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 :: (Key->Key) -> IntMap a -> IntMap a +mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) [] + +-- | /O(n*min(n,W))/. +-- @'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 :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a +mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] + +-- | /O(n*min(n,W))/. +-- @'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 slightly better performance than 'mapKeys'. +-- +-- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")] + +mapKeysMonotonic :: (Key->Key) -> IntMap a -> IntMap a +mapKeysMonotonic f = fromDistinctAscList . foldrWithKey (\k x xs -> (f k, x) : xs) [] + +{-------------------------------------------------------------------- + Filter +--------------------------------------------------------------------} +-- | /O(n)/. Filter all values that satisfy some 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) -> IntMap a -> IntMap a +filter p m + = filterWithKey (\_ x -> p x) m + +-- | /O(n)/. Filter all keys\/values that satisfy some predicate. +-- +-- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a +filterWithKey predicate t + = case t of + Bin p m l r + -> bin p m (filterWithKey predicate l) (filterWithKey predicate r) + Tip k x + | predicate k x -> t + | otherwise -> Nil + Nil -> Nil + +-- | /O(n)/. Partition the map according to some 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) -> IntMap a -> (IntMap a,IntMap a) +partition p m + = partitionWithKey (\_ x -> p x) m + +-- | /O(n)/. Partition the map according to some 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 :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a) +partitionWithKey predicate t + = case t of + Bin p m l r + -> let (l1,l2) = partitionWithKey predicate l + (r1,r2) = partitionWithKey predicate r + in (bin p m l1 r1, bin p m l2 r2) + Tip k x + | predicate k x -> (t,Nil) + | otherwise -> (Nil,t) + Nil -> (Nil,Nil) + +-- | /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) -> IntMap a -> IntMap 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 :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b +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 -> Tip k y + Nothing -> Nil +mapMaybeWithKey _ Nil = Nil + +-- | /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) -> IntMap a -> (IntMap b, IntMap 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 :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) +mapEitherWithKey f (Bin p m l r) + = (bin p m l1 r1, bin p m l2 r2) + where + (l1,l2) = mapEitherWithKey f l + (r1,r2) = mapEitherWithKey f r +mapEitherWithKey f (Tip k x) = case f k x of + Left y -> (Tip k y, Nil) + Right z -> (Nil, Tip k z) +mapEitherWithKey _ Nil = (Nil, Nil) + +-- | /O(min(n,W))/. The expression (@'split' k map@) is a pair @(map1,map2)@ +-- where all keys in @map1@ are lower than @k@ and all 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 :: Key -> IntMap a -> (IntMap a, IntMap a) +split k t = + case t of Bin _ m l r | m < 0 -> if k >= 0 -- handle negative numbers. + then case go k l of (lt, gt) -> (union r lt, gt) + else case go k r of (lt, gt) -> (lt, union gt l) + _ -> go k t + where + go k' t'@(Bin p m l r) | nomatch k' p m = if k' > p then (t', Nil) else (Nil, t') + | zero k' m = case go k' l of (lt, gt) -> (lt, union gt r) + | otherwise = case go k' r of (lt, gt) -> (union l lt, gt) + go k' t'@(Tip ky _) | k' > ky = (t', Nil) + | k' < ky = (Nil, t') + | otherwise = (Nil, Nil) + go _ Nil = (Nil, Nil) + +-- | /O(min(n,W))/. Performs a 'split' but also returns whether the pivot +-- key was found in the original 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 :: Key -> IntMap a -> (IntMap a, Maybe a, IntMap a) +splitLookup k t = + case t of Bin _ m l r | m < 0 -> if k >= 0 -- handle negative numbers. + then case go k l of (lt, fnd, gt) -> (union r lt, fnd, gt) + else case go k r of (lt, fnd, gt) -> (lt, fnd, union gt l) + _ -> 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) -> (lt, fnd, union gt r) + | otherwise = case go k' r of (lt, fnd, gt) -> (union l lt, fnd, gt) + go k' t'@(Tip ky y) | k' > ky = (t', Nothing, Nil) + | k' < ky = (Nil, Nothing, t') + | otherwise = (Nil, Just y, Nil) + go _ Nil = (Nil, Nothing, Nil) + +{-------------------------------------------------------------------- + Fold +--------------------------------------------------------------------} +-- | /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 -> IntMap a -> b +foldr f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t + where + go z' Nil = z' + go z' (Tip _ x) = f x z' + go z' (Bin _ _ l r) = go (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 -> IntMap a -> b +foldr' f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t + where + STRICT_1_OF_2(go) + go z' Nil = z' + go z' (Tip _ x) = f x z' + go z' (Bin _ _ l r) = go (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 -> IntMap b -> a +foldl f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t + where + go z' Nil = z' + go z' (Tip _ x) = f z' x + go z' (Bin _ _ l r) = go (go z' l) 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 -> IntMap b -> a +foldl' f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t + where + STRICT_1_OF_2(go) + go z' Nil = z' + go z' (Tip _ x) = f z' x + go z' (Bin _ _ l r) = go (go z' l) 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 :: (Int -> a -> b -> b) -> b -> IntMap a -> b +foldrWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t + where + 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 #-} + +-- | /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' :: (Int -> a -> b -> b) -> b -> IntMap a -> b +foldrWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t + where + STRICT_1_OF_2(go) + 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' #-} + +-- | /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 -> Int -> b -> a) -> a -> IntMap b -> a +foldlWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t + where + 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 #-} + +-- | /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 -> Int -> b -> a) -> a -> IntMap b -> a +foldlWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t + where + STRICT_1_OF_2(go) + 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' #-} + +{-------------------------------------------------------------------- + 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 :: IntMap 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 :: IntMap a -> [Key] +keys = foldrWithKey (\k _ ks -> k : ks) [] + +-- | /O(n)/. An alias for 'toAscList'. Returns 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 :: IntMap a -> [(Key,a)] +assocs = toAscList + +-- | /O(n*min(n,W))/. The set of all keys of the map. +-- +-- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5] +-- > keysSet empty == Data.IntSet.empty + +keysSet :: IntMap a -> IntSet.IntSet +keysSet Nil = IntSet.Nil +keysSet (Tip kx _) = IntSet.singleton kx +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' + computeBm acc (Tip kx _) = acc .|. IntSet.bitmapOf kx + computeBm _ Nil = error "Data.IntSet.keysSet: Nil" + +-- | /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.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")] +-- > fromSet undefined Data.IntSet.empty == empty + +fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a +fromSet _ IntSet.Nil = Nil +fromSet f (IntSet.Bin p m l r) = Bin p m (fromSet f l) (fromSet f r) +fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1) + where -- This is slightly complicated, as we to convert the dense + -- representation of IntSet into tree representation of IntMap. + -- + -- We are given a nonzero bit mask 'bmask' of 'bits' bits with prefix 'prefix'. + -- 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 + 0 -> Tip prefix (g prefix) + _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of + bits2 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 -> + buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2 + | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 -> + buildTree g prefix bmask bits2 + | otherwise -> + Bin prefix bits2 (buildTree g prefix bmask bits2) (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2) + +{-------------------------------------------------------------------- + Lists +--------------------------------------------------------------------} +-- | /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 :: IntMap a -> [(Key,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 :: IntMap a -> [(Key,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 :: IntMap a -> [(Key,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 :: (Key -> a -> b -> b) -> b -> IntMap a -> b +foldrFB = foldrWithKey +{-# INLINE[0] foldrFB #-} +foldlFB :: (a -> Key -> b -> a) -> a -> IntMap 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 "IntMap.elems" [~1] forall m . elems m = build (\c n -> foldrFB (\_ x xs -> c x xs) n m) #-} +{-# RULES "IntMap.elemsBack" [1] foldrFB (\_ x xs -> x : xs) [] = elems #-} +{-# RULES "IntMap.keys" [~1] forall m . keys m = build (\c n -> foldrFB (\k _ xs -> c k xs) n m) #-} +{-# RULES "IntMap.keysBack" [1] foldrFB (\k _ xs -> k : xs) [] = keys #-} +{-# RULES "IntMap.toAscList" [~1] forall m . toAscList m = build (\c n -> foldrFB (\k x xs -> c (k,x) xs) n m) #-} +{-# RULES "IntMap.toAscListBack" [1] foldrFB (\k x xs -> (k, x) : xs) [] = toAscList #-} +{-# RULES "IntMap.toDescList" [~1] forall m . toDescList m = build (\c n -> foldlFB (\xs k x -> c (k,x) xs) n m) #-} +{-# RULES "IntMap.toDescListBack" [1] foldlFB (\xs k x -> (k, x) : xs) [] = toDescList #-} +#endif + + +-- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs. +-- +-- > 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")] + +fromList :: [(Key,a)] -> IntMap a +fromList xs + = foldlStrict ins empty xs + where + ins t (k,x) = insert k x t + +-- | /O(n*min(n,W))/. Create 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,"c")] == fromList [(3, "ab"), (5, "cba")] +-- > fromListWith (++) [] == empty + +fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a +fromListWith f xs + = fromListWithKey (\_ x y -> f x y) xs + +-- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")] +-- > fromListWithKey f [] == empty + +fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a +fromListWithKey f xs + = foldlStrict ins empty xs + where + ins t (k,x) = insertWithKey f k x t + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order. +-- +-- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] + +fromAscList :: [(Key,a)] -> IntMap a +fromAscList xs + = fromAscListWithKey (\_ x _ -> x) xs + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order, with a combining function on equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] + +fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a +fromAscListWith f xs + = fromAscListWithKey (\_ x y -> f x y) xs + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order, with a combining function on equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "5:b|a")] + +fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a +fromAscListWithKey _ [] = Nil +fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0) + where + -- [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 combineEq (kx,yy) xs + | otherwise = z:combineEq x xs + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order and all distinct. +-- /The precondition (input list is strictly ascending) is not checked./ +-- +-- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] + +fromDistinctAscList :: [(Key,a)] -> IntMap a +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 + + 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) + reduce z zs m px tx stk@(Push py ty stk') = + let mxy = branchMask px py + pxy = mask px mxy + in if shorter m mxy + then reduce z zs m pxy (Bin pxy mxy ty tx) stk' + else work z zs (Push px tx stk) + + finish _ t Nada = t + finish px tx (Push py ty stk) = finish p (join py ty px tx) stk + where m = branchMask px py + p = mask px m + +data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada + + +{-------------------------------------------------------------------- + Eq +--------------------------------------------------------------------} +instance Eq a => Eq (IntMap a) where + t1 == t2 = equal t1 t2 + t1 /= t2 = nequal t1 t2 + +equal :: Eq a => IntMap a -> IntMap a -> Bool +equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) + = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2) +equal (Tip kx x) (Tip ky y) + = (kx == ky) && (x==y) +equal Nil Nil = True +equal _ _ = False + +nequal :: Eq a => IntMap a -> IntMap a -> Bool +nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) + = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2) +nequal (Tip kx x) (Tip ky y) + = (kx /= ky) || (x/=y) +nequal Nil Nil = False +nequal _ _ = True + +{-------------------------------------------------------------------- + Ord +--------------------------------------------------------------------} + +instance Ord a => Ord (IntMap a) where + compare m1 m2 = compare (toList m1) (toList m2) + +{-------------------------------------------------------------------- + Functor +--------------------------------------------------------------------} + +instance Functor IntMap where + fmap = map + +{-------------------------------------------------------------------- + Show +--------------------------------------------------------------------} + +instance Show a => Show (IntMap a) where + showsPrec d m = showParen (d > 10) $ + showString "fromList " . shows (toList m) + +{-------------------------------------------------------------------- + Read +--------------------------------------------------------------------} +instance (Read e) => Read (IntMap 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 + +{-------------------------------------------------------------------- + Typeable +--------------------------------------------------------------------} + +#include "Typeable.h" +INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap") + +{-------------------------------------------------------------------- + Helpers +--------------------------------------------------------------------} +{-------------------------------------------------------------------- + Join +--------------------------------------------------------------------} +join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a +join p1 t1 p2 t2 + | zero p1 m = Bin p m t1 t2 + | otherwise = Bin p m t2 t1 + where + m = branchMask p1 p2 + p = mask p1 m +{-# INLINE join #-} + +{-------------------------------------------------------------------- + @bin@ assures that we never have empty trees within a tree. +--------------------------------------------------------------------} +bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a +bin _ _ l Nil = l +bin _ _ Nil r = r +bin p m l r = Bin p m l r +{-# INLINE bin #-} + + +{-------------------------------------------------------------------- + Endian independent bit twiddling +--------------------------------------------------------------------} +zero :: Key -> Mask -> Bool +zero i m + = (natFromInt i) .&. (natFromInt m) == 0 +{-# INLINE zero #-} + +nomatch,match :: Key -> Prefix -> Mask -> Bool +nomatch i p m + = (mask i m) /= p +{-# INLINE nomatch #-} + +match i p m + = (mask i m) == p +{-# INLINE match #-} + +mask :: Key -> Mask -> Prefix +mask i m + = maskW (natFromInt i) (natFromInt m) +{-# INLINE mask #-} + + +{-------------------------------------------------------------------- + Big endian operations +--------------------------------------------------------------------} +maskW :: Nat -> Nat -> Prefix +maskW i m + = intFromNat (i .&. (complement (m-1) `xor` m)) +{-# INLINE maskW #-} + +shorter :: Mask -> Mask -> Bool +shorter m1 m2 + = (natFromInt m1) > (natFromInt m2) +{-# INLINE shorter #-} + +branchMask :: Prefix -> Prefix -> Mask +branchMask p1 p2 + = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2)) +{-# INLINE branchMask #-} + +{---------------------------------------------------------------------- + Finding the highest bit (mask) in a word [x] can be done efficiently in + three ways: + * convert to a floating point value and the mantissa tells us the + [log2(x)] that corresponds with the highest bit position. The mantissa + is retrieved either via the standard C function [frexp] or by some bit + twiddling on IEEE compatible numbers (float). Note that one needs to + use at least [double] precision for an accurate mantissa of 32 bit + numbers. + * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit). + * use processor specific assembler instruction (asm). + + The most portable way would be [bit], but is it efficient enough? + I have measured the cycle counts of the different methods on an AMD + Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction: + + highestBitMask: method cycles + -------------- + frexp 200 + float 33 + bit 11 + asm 12 + + highestBit: method cycles + -------------- + frexp 195 + float 33 + bit 11 + asm 11 + + Wow, the bit twiddling is on today's RISC like machines even faster + than a single CISC instruction (BSR)! +----------------------------------------------------------------------} + +{---------------------------------------------------------------------- + [highestBitMask] returns a word where only the highest bit is set. + It is found by first setting all bits in lower positions than the + highest bit and than taking an exclusive or with the original value. + Allthough the function may look expensive, GHC compiles this into + excellent C code that subsequently compiled into highly efficient + machine code. The algorithm is derived from Jorg Arndt's FXT library. +----------------------------------------------------------------------} +highestBitMask :: Nat -> Nat +highestBitMask x0 + = case (x0 .|. shiftRL x0 1) of + x1 -> case (x1 .|. shiftRL x1 2) of + x2 -> case (x2 .|. shiftRL x2 4) of + x3 -> case (x3 .|. shiftRL x3 8) of + x4 -> case (x4 .|. shiftRL x4 16) of +#if !(defined(__GLASGOW_HASKELL__) && WORD_SIZE_IN_BITS==32) + x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms +#endif + x6 -> (x6 `xor` (shiftRL x6 1)) +{-# INLINE highestBitMask #-} + + +{-------------------------------------------------------------------- + Utilities +--------------------------------------------------------------------} + +foldlStrict :: (a -> b -> a) -> a -> [b] -> a +foldlStrict f = go + where + go z [] = z + go z (x:xs) = let z' = f z x in z' `seq` go z' xs +{-# INLINE foldlStrict #-} + +{-------------------------------------------------------------------- + Debugging +--------------------------------------------------------------------} +-- | /O(n)/. Show the tree that implements the map. The tree is shown +-- in a compressed, hanging format. +showTree :: Show a => IntMap a -> String +showTree s + = showTreeWith True False s + + +{- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows + the tree that implements the map. 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. +-} +showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String +showTreeWith hang wide t + | hang = (showsTreeHang wide [] t) "" + | otherwise = (showsTree wide [] [] t) "" + +showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS +showsTree wide lbars rbars t + = case t of + Bin p m l r + -> showsTree wide (withBar rbars) (withEmpty rbars) r . + showWide wide rbars . + showsBars lbars . showString (showBin p m) . showString "\n" . + showWide wide lbars . + showsTree wide (withEmpty lbars) (withBar lbars) l + Tip k x + -> showsBars lbars . showString " " . shows k . showString ":=" . shows x . showString "\n" + Nil -> showsBars lbars . showString "|\n" + +showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS +showsTreeHang wide bars t + = case t of + Bin p m l r + -> showsBars bars . showString (showBin p m) . showString "\n" . + showWide wide bars . + showsTreeHang wide (withBar bars) l . + showWide wide bars . + showsTreeHang wide (withEmpty bars) r + Tip k x + -> showsBars bars . showString " " . shows k . showString ":=" . shows x . showString "\n" + Nil -> showsBars bars . showString "|\n" + +showBin :: Prefix -> Mask -> String +showBin _ _ + = "*" -- ++ show (p,m) + +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 diff --git a/Data/IntMap/Lazy.hs b/Data/IntMap/Lazy.hs new file mode 100644 --- /dev/null +++ b/Data/IntMap/Lazy.hs @@ -0,0 +1,214 @@ +{-# LANGUAGE CPP #-} +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.IntMap.Lazy +-- 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 integer keys to values +-- (dictionaries). +-- +-- API of this module is strict in the keys, but lazy in the values. +-- If you need value-strict maps, use 'Data.IntMap.Strict' instead. +-- The 'IntMap' type itself is shared between the lazy and strict modules, +-- meaning that the same 'IntMap' 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 Data.IntMap.Lazy (IntMap) +-- > import qualified Data.IntMap.Lazy as IntMap +-- +-- The implementation is based on /big-endian patricia trees/. This data +-- structure performs especially well on binary operations like 'union' +-- and 'intersection'. However, my benchmarks show that it is also +-- (much) faster on insertions and deletions when compared to a generic +-- size-balanced map implementation (see "Data.Map"). +-- +-- * 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), +-- October 1968, pages 514-534. +-- +-- Operation comments contain the operation time complexity in +-- the Big-O notation . +-- Many operations have a worst-case complexity of /O(min(n,W))/. +-- This means that the operation can become linear in the number of +-- elements with a maximum of /W/ -- the number of bits in an 'Int' +-- (32 or 64). +----------------------------------------------------------------------------- + +module Data.IntMap.Lazy ( + -- * Strictness properties + -- $strictness + + -- * Map type +#if !defined(TESTING) + IntMap, Key -- instance Eq,Show +#else + IntMap(..), Key -- instance Eq,Show +#endif + + -- * Operators + , (!), (\\) + + -- * Query + , IM.null + , size + , member + , notMember + , IM.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 + , IM.map + , mapWithKey + , traverseWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeys + , mapKeysWith + , mapKeysMonotonic + + -- * Folds + , IM.foldr + , IM.foldl + , foldrWithKey + , foldlWithKey + -- ** 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 + , IM.filter + , filterWithKey + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + ) where + +import Data.IntMap.Base as IM + +-- $strictness +-- +-- This module satisfies the following strictness property: +-- +-- * Key arguments are evaluated to WHNF +-- +-- Here are some examples that illustrate the property: +-- +-- > insertWith (\ new old -> old) undefined v m == undefined +-- > insertWith (\ new old -> old) k undefined m == OK +-- > delete undefined m == undefined diff --git a/Data/IntMap/Strict.hs b/Data/IntMap/Strict.hs new file mode 100644 --- /dev/null +++ b/Data/IntMap/Strict.hs @@ -0,0 +1,964 @@ +{-# LANGUAGE CPP #-} +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.IntMap.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 maps from integer 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.IntMap.Lazy' instead. +-- The 'IntMap' type itself is shared between the lazy and strict modules, +-- meaning that the same 'IntMap' 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 Data.IntMap.Strict (IntMap) +-- > import qualified Data.IntMap.Strict as IntMap +-- +-- The implementation is based on /big-endian patricia trees/. This data +-- structure performs especially well on binary operations like 'union' +-- and 'intersection'. However, my benchmarks show that it is also +-- (much) faster on insertions and deletions when compared to a generic +-- size-balanced map implementation (see "Data.Map"). +-- +-- * 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), +-- October 1968, pages 514-534. +-- +-- Operation comments contain the operation time complexity in +-- the Big-O notation . +-- Many operations have a worst-case complexity of /O(min(n,W))/. +-- This means that the operation can become linear in the number of +-- elements with a maximum of /W/ -- the number of bits in an 'Int' +-- (32 or 64). +-- +-- Be aware that the 'Functor', 'Traversable' and 'Data' instances +-- are the same as for the 'Data.IntMap.Lazy' module, so if they are used +-- on strict maps, the resulting maps will be lazy. +----------------------------------------------------------------------------- + +-- See the notes at the beginning of Data.IntMap.Base. + +module Data.IntMap.Strict ( + -- * Strictness properties + -- $strictness + + -- * Map type +#if !defined(TESTING) + IntMap, Key -- instance Eq,Show +#else + IntMap(..), Key -- instance Eq,Show +#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 + -- ** 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 + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + ) where + +import Prelude hiding (lookup,map,filter,foldr,foldl,null) + +import Data.Bits +import Data.IntMap.Base hiding + ( findWithDefault + , singleton + , insert + , insertWith + , insertWithKey + , insertLookupWithKey + , adjust + , adjustWithKey + , update + , updateWithKey + , updateLookupWithKey + , alter + , unionsWith + , unionWith + , unionWithKey + , differenceWith + , differenceWithKey + , intersectionWith + , intersectionWithKey + , mergeWithKey + , updateMinWithKey + , updateMaxWithKey + , updateMax + , updateMin + , map + , mapWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeysWith + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + , fromSet + , fromList + , fromListWith + , fromListWithKey + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + ) +import qualified Data.IntSet.Base as IntSet +import Data.StrictPair + +-- $strictness +-- +-- This module satisfies the following strictness properties: +-- +-- 1. Key and value arguments are evaluated to WHNF; +-- +-- 2. Keys and values are evaluated to WHNF before they are stored in +-- the map. +-- +-- Here are some examples that illustrate the first property: +-- +-- > insertWith (\ new old -> old) k undefined m == undefined +-- > 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(min(n,W))/. The expression @('findWithDefault' def k map)@ +-- returns the value at key @k@ or returns @def@ when the key is not an +-- element of the map. +-- +-- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' +-- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' + +-- See IntMap.Base.Note: Local 'go' functions and capturing] +findWithDefault :: a -> Key -> IntMap a -> a +findWithDefault def k = def `seq` k `seq` go + where + go (Bin p m l r) | nomatch k p m = def + | zero k m = go l + | otherwise = go r + go (Tip kx x) | k == kx = x + | otherwise = def + go Nil = def + +{-------------------------------------------------------------------- + Construction +--------------------------------------------------------------------} +-- | /O(1)/. A map of one element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +-- > size (singleton 1 'a') == 1 + +singleton :: Key -> a -> IntMap a +singleton k x + = x `seq` Tip k x +{-# INLINE singleton #-} + +{-------------------------------------------------------------------- + Insert +--------------------------------------------------------------------} +-- | /O(min(n,W))/. Insert a new key\/value pair in the map. +-- If the key is already present in the map, the associated value is +-- replaced with the supplied value, i.e. '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' + +insert :: Key -> a -> IntMap a -> IntMap a +insert k x t = k `seq` x `seq` + case t of + Bin p m l r + | nomatch k p m -> join 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 -> join k (Tip k x) ky t + Nil -> Tip k x + +-- right-biased insertion, used by 'union' +-- | /O(min(n,W))/. Insert with a combining function. +-- @'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 @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 :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWith f k x t + = insertWithKey (\_ x' y' -> f x' y') k x t + +-- | /O(min(n,W))/. Insert with a combining function. +-- @'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 @f key new_value old_value@. +-- +-- > 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" +-- +-- If the key exists in the map, this function is lazy in @x@ but strict +-- in the result of @f@. + +insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWithKey f k x t = k `seq` x `seq` + case t of + Bin p m l r + | nomatch k p m -> join 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 -> join k (Tip k x) ky t + 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@) +-- 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")]) + +insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a) +insertLookupWithKey f k x t = k `seq` x `seq` + case t of + Bin p m l r + | nomatch k p m -> Nothing `strictPair` join k (Tip k x) p t + | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found `strictPair` Bin p m l' r) + | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found `strictPair` Bin p m l r') + Tip ky y + | k==ky -> (Just y `strictPair` (Tip k $! f k x y)) + | otherwise -> (Nothing `strictPair` join k (Tip k x) ky t) + Nil -> Nothing `strictPair` Tip k x + + +{-------------------------------------------------------------------- + Deletion +--------------------------------------------------------------------} +-- | /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. +-- +-- > 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 :: (a -> a) -> Key -> IntMap a -> IntMap a +adjust f k m + = adjustWithKey (\_ x -> f x) k m + +-- | /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. +-- +-- > 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 :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a +adjustWithKey f + = updateWithKey (\k' x -> Just (f k' x)) + +-- | /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@. +-- +-- > 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 :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a +update f + = updateWithKey (\_ x -> f x) + +-- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +updateWithKey :: (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 + Just y' -> y' `seq` Tip ky y' + Nothing -> Nil + | otherwise -> t + Nil -> Nil + +-- | /O(min(n,W))/. Lookup and update. +-- The function returns original value, if it is updated. +-- This is different behavior than 'Data.Map.updateLookupWithKey'. +-- 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 "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") + +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 `strictPair` bin p m l' r) + | otherwise -> let (found,r') = updateLookupWithKey f k r in (found `strictPair` bin p m l r') + Tip ky y + | k==ky -> case f k y of + Just y' -> y' `seq` (Just y `strictPair` Tip ky y') + Nothing -> (Just y, Nil) + | otherwise -> (Nothing,t) + Nil -> (Nothing,Nil) + + + +-- | /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 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 -> x `seq` join 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 -> x `seq` Tip ky x + Nothing -> Nil + | otherwise -> case f Nothing of + Just x -> x `seq` join k (Tip k x) ky t + Nothing -> t + Nil -> case f Nothing of + Just x -> x `seq` Tip k x + Nothing -> Nil + + +{-------------------------------------------------------------------- + Union +--------------------------------------------------------------------} +-- | The union of a list of maps, with a combining operation. +-- +-- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] +-- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] + +unionsWith :: (a->a->a) -> [IntMap a] -> IntMap a +unionsWith f ts + = foldlStrict (unionWith f) empty ts + +-- | /O(n+m)/. The union with a combining function. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] + +unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWith f m1 m2 + = unionWithKey (\_ x y -> f x y) m1 m2 + +-- | /O(n+m)/. The 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 :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWithKey f m1 m2 + = mergeWithKey' Bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 $! f k1 x1 x2) id id m1 m2 + +{-------------------------------------------------------------------- + Difference +--------------------------------------------------------------------} + +-- | /O(n+m)/. Difference with a combining function. +-- +-- > 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 :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a +differenceWith f m1 m2 + = differenceWithKey (\_ x y -> f x y) m1 m2 + +-- | /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 :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a +differenceWithKey f m1 m2 + = mergeWithKey f id (const Nil) m1 m2 + +{-------------------------------------------------------------------- + Intersection +--------------------------------------------------------------------} + +-- | /O(n+m)/. The intersection with a combining function. +-- +-- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" + +intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c +intersectionWith f m1 m2 + = intersectionWithKey (\_ x y -> f x y) m1 m2 + +-- | /O(n+m)/. The 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 :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c +intersectionWithKey f m1 m2 + = mergeWithKey' bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 $! f k1 x1 x2) (const Nil) (const Nil) m1 m2 + +{-------------------------------------------------------------------- + MergeWithKey +--------------------------------------------------------------------} + +-- | /O(n+m)/. A high-performance universal combining function. Using +-- 'mergeWithKey', all combining functions can be defined without any loss of +-- efficiency (with exception of 'union', 'difference' and 'intersection', +-- where sharing of some nodes is lost with 'mergeWithKey'). +-- +-- 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 :: (Key -> a -> b -> Maybe c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c) + -> IntMap a -> IntMap b -> IntMap c +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 + {-# INLINE combine #-} +{-# INLINE mergeWithKey #-} + +{-------------------------------------------------------------------- + Min\/Max +--------------------------------------------------------------------} + +-- | /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 :: (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) + _ -> go f t + where + go f' (Bin p m l r) = bin p m (go f' l) r + go f' (Tip k y) = case f' k y of + Just y' -> y' `seq` Tip k y' + Nothing -> Nil + go _ Nil = error "updateMinWithKey Nil" + +-- | /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 :: (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 + _ -> go f t + where + go f' (Bin p m l r) = bin p m l (go f' r) + go f' (Tip k y) = case f' k y of + Just y' -> y' `seq` Tip k y' + Nothing -> Nil + go _ Nil = error "updateMaxWithKey Nil" + +-- | /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) -> IntMap a -> IntMap a +updateMax f = updateMaxWithKey (const f) + +-- | /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) -> IntMap a -> IntMap a +updateMin f = updateMinWithKey (const f) + + +{-------------------------------------------------------------------- + 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) -> 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 + +-- | /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 :: (Key -> a -> b) -> IntMap a -> IntMap b +mapWithKey f t + = case t of + Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r) + Tip k x -> Tip k $! f k x + Nil -> Nil + +-- | /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 -> IntMap b -> (a,IntMap c) +mapAccum f = mapAccumWithKey (\a' _ x -> f a' x) + +-- | /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 -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap 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. Strict in +-- the accumulating argument and the both elements of the +-- result of the function. +mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) +mapAccumL f a t + = case t of + Bin p m l r -> let (a1,l') = mapAccumL f a l + (a2,r') = mapAccumL f a1 r + in (a2 `strictPair` Bin p m l' r') + Tip k x -> let (a',x') = f a k x in x' `seq` (a' `strictPair` Tip k x') + Nil -> (a `strictPair` Nil) + +-- | /O(n)/. The function @'mapAccumR'@ threads an accumulating +-- argument through the map in descending order of keys. +mapAccumRWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) +mapAccumRWithKey f a t + = case t of + Bin p m l r -> let (a1,r') = mapAccumRWithKey f a r + (a2,l') = mapAccumRWithKey f a1 l + in (a2 `strictPair` Bin p m l' r') + Tip k x -> let (a',x') = f a k x in x' `seq` (a' `strictPair` Tip k x') + Nil -> (a `strictPair` Nil) + +-- | /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 :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a +mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] + +{-------------------------------------------------------------------- + Filter +--------------------------------------------------------------------} +-- | /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) -> IntMap a -> IntMap 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 :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b +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 + Nothing -> Nil +mapMaybeWithKey _ Nil = Nil + +-- | /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) -> IntMap a -> (IntMap b, IntMap 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 :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) +mapEitherWithKey f (Bin p m l r) + = bin p m l1 r1 `strictPair` bin p m l2 r2 + where + (l1,l2) = mapEitherWithKey f l + (r1,r2) = mapEitherWithKey f r +mapEitherWithKey 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) +mapEitherWithKey _ Nil = (Nil, Nil) + +{-------------------------------------------------------------------- + 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.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")] +-- > fromSet undefined Data.IntSet.empty == empty + +fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a +fromSet _ IntSet.Nil = Nil +fromSet f (IntSet.Bin p m l r) = Bin p m (fromSet f l) (fromSet f r) +fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1) + where -- This is slightly complicated, as we to convert the dense + -- representation of IntSet into tree representation of IntMap. + -- + -- We are given a nonzero bit mask 'bmask' of 'bits' bits with prefix 'prefix'. + -- 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 + 0 -> Tip prefix $! g prefix + _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of + bits2 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 -> + buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2 + | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 -> + buildTree g prefix bmask bits2 + | otherwise -> + Bin prefix bits2 (buildTree g prefix bmask bits2) (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2) + +{-------------------------------------------------------------------- + Lists +--------------------------------------------------------------------} +-- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs. +-- +-- > 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")] + +fromList :: [(Key,a)] -> IntMap a +fromList xs + = foldlStrict ins empty xs + where + ins t (k,x) = insert k x t + +-- | /O(n*min(n,W))/. Create 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 :: (a -> a -> a) -> [(Key,a)] -> IntMap a +fromListWith f xs + = fromListWithKey (\_ x y -> f x y) xs + +-- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'. +-- +-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")] +-- > fromListWith (++) [] == empty + +fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a +fromListWithKey f xs + = foldlStrict ins empty xs + where + ins t (k,x) = insertWithKey f k x t + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order. +-- +-- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] + +fromAscList :: [(Key,a)] -> IntMap a +fromAscList xs + = fromAscListWithKey (\_ x _ -> x) xs + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order, with a combining function on equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] + +fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a +fromAscListWith f xs + = fromAscListWithKey (\_ x y -> f x y) xs + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order, with a combining function on equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] + +fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a +fromAscListWithKey _ [] = Nil +fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0) + where + -- [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 + | otherwise = z:combineEq x xs + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order and all distinct. +-- /The precondition (input list is strictly ascending) is not checked./ +-- +-- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] + +fromDistinctAscList :: [(Key,a)] -> IntMap a +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 + + 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) + reduce z zs m px tx stk@(Push py ty stk') = + let mxy = branchMask px py + pxy = mask px mxy + in if shorter m mxy + then reduce z zs m pxy (Bin pxy mxy ty tx) stk' + else work z zs (Push px tx stk) + + finish _ t Nada = t + finish px tx (Push py ty stk) = finish p (join py ty px tx) stk + where m = branchMask px py + p = mask px m + +data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada diff --git a/Data/IntSet.hs b/Data/IntSet.hs --- a/Data/IntSet.hs +++ b/Data/IntSet.hs @@ -1,1100 +1,148 @@ -#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 -{-# LANGUAGE Trustworthy #-} -#endif ------------------------------------------------------------------------------ --- | --- Module : Data.IntSet --- Copyright : (c) Daan Leijen 2002 --- License : BSD-style --- Maintainer : libraries@haskell.org --- Stability : provisional --- Portability : portable --- --- An efficient implementation of integer sets. --- --- Since many function names (but not the type name) clash with --- "Prelude" names, this module is usually imported @qualified@, e.g. --- --- > import Data.IntSet (IntSet) --- > import qualified Data.IntSet as IntSet --- --- The implementation is based on /big-endian patricia trees/. This data --- structure performs especially well on binary operations like 'union' --- and 'intersection'. However, my benchmarks show that it is also --- (much) faster on insertions and deletions when compared to a generic --- size-balanced set implementation (see "Data.Set"). --- --- * 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), --- October 1968, pages 514-534. --- --- Many operations have a worst-case complexity of /O(min(n,W))/. --- This means that the operation can become linear in the number of --- elements with a maximum of /W/ -- the number of bits in an 'Int' --- (32 or 64). ------------------------------------------------------------------------------ - --- It is essential that the bit fiddling functions like mask, zero, branchMask --- etc are inlined. If they do not, the memory allocation skyrockets. The GHC --- usually gets it right, but it is disastrous if it does not. Therefore we --- explicitly mark these functions INLINE. - -module Data.IntSet ( - -- * Set type -#if !defined(TESTING) - IntSet -- instance Eq,Show -#else - IntSet(..) -- instance Eq,Show -#endif - - -- * Operators - , (\\) - - -- * Query - , null - , size - , member - , notMember - , isSubsetOf - , isProperSubsetOf - - -- * Construction - , empty - , singleton - , insert - , delete - - -- * Combine - , union - , unions - , difference - , intersection - - -- * Filter - , filter - , partition - , split - , splitMember - - -- * Map - , map - - -- * Folds - , foldr - , foldl - -- ** Strict folds - , foldr' - , foldl' - -- ** Legacy folds - , fold - - -- * Min\/Max - , findMin - , findMax - , deleteMin - , deleteMax - , deleteFindMin - , deleteFindMax - , maxView - , minView - - -- * Conversion - - -- ** List - , elems - , toList - , fromList - - -- ** Ordered list - , toAscList - , fromAscList - , fromDistinctAscList - - -- * Debugging - , showTree - , showTreeWith - -#if defined(TESTING) - -- * Internals - , match -#endif - ) where - - -import Prelude hiding (lookup,filter,foldr,foldl,null,map) -import Data.Bits - -import qualified Data.List as List -import Data.Monoid (Monoid(..)) -import Data.Maybe (fromMaybe) -import Data.Typeable -import Control.DeepSeq (NFData) - -#if __GLASGOW_HASKELL__ -import Text.Read -import Data.Data (Data(..), mkNoRepType) -#endif - -#if __GLASGOW_HASKELL__ >= 503 -import GHC.Exts ( Word(..), Int(..), shiftRL# ) -#elif __GLASGOW_HASKELL__ -import Word -import GlaExts ( Word(..), Int(..), shiftRL# ) -#else -import Data.Word -#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 - -infixl 9 \\{-This comment teaches CPP correct behaviour -} - --- A "Nat" is a natural machine word (an unsigned Int) -type Nat = Word - -natFromInt :: Int -> Nat -natFromInt i = fromIntegral i -{-# INLINE natFromInt #-} - -intFromNat :: Nat -> Int -intFromNat w = fromIntegral w -{-# INLINE intFromNat #-} - -shiftRL :: Nat -> Int -> Nat -#if __GLASGOW_HASKELL__ -{-------------------------------------------------------------------- - GHC: use unboxing to get @shiftRL@ inlined. ---------------------------------------------------------------------} -shiftRL (W# x) (I# i) - = W# (shiftRL# x i) -#else -shiftRL x i = shiftR x i -{-# INLINE shiftRL #-} -#endif - -{-------------------------------------------------------------------- - Operators ---------------------------------------------------------------------} --- | /O(n+m)/. See 'difference'. -(\\) :: IntSet -> IntSet -> IntSet -m1 \\ m2 = difference m1 m2 - -{-------------------------------------------------------------------- - Types ---------------------------------------------------------------------} - --- The order of constructors of IntSet matters when considering performance. --- Currently in GHC 7.0, when type has 3 constructors, they are matched from --- the first to the last -- the best performance is achieved when the --- constructors are ordered by frequency. --- On GHC 7.0, reordering constructors from Nil | Tip | Bin to Bin | Tip | Nil --- improves the containers_benchmark by 11% on x86 and by 9% on x86_64. - --- | A set of integers. -data IntSet = Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !IntSet !IntSet - | Tip {-# UNPACK #-} !Int - | Nil --- Invariant: Nil is never found as a child of Bin. --- Invariant: The Mask is a power of 2. It is the largest bit position at which --- two elements of the set differ. --- Invariant: Prefix is the common high-order bits that all elements share to --- the left of the Mask bit. --- Invariant: In Bin prefix mask left right, left consists of the elements that --- don't have the mask bit set; right is all the elements that do. - - -type Prefix = Int -type Mask = Int - -instance Monoid IntSet where - mempty = empty - mappend = union - mconcat = unions - -#if __GLASGOW_HASKELL__ - -{-------------------------------------------------------------------- - A Data instance ---------------------------------------------------------------------} - --- This instance preserves data abstraction at the cost of inefficiency. --- We omit reflection services for the sake of data abstraction. - -instance Data IntSet where - gfoldl f z is = z fromList `f` (toList is) - toConstr _ = error "toConstr" - gunfold _ _ = error "gunfold" - dataTypeOf _ = mkNoRepType "Data.IntSet.IntSet" - -#endif - -{-------------------------------------------------------------------- - Query ---------------------------------------------------------------------} --- | /O(1)/. Is the set empty? -null :: IntSet -> Bool -null Nil = True -null _ = False - --- | /O(n)/. Cardinality of the set. -size :: IntSet -> Int -size t - = case t of - Bin _ _ l r -> size l + size r - Tip _ -> 1 - Nil -> 0 - --- The 'go' function in the member and lookup causes 10% speedup, but also an --- increased memory allocation. It does not cause speedup with other methods --- like insert and delete, so it is present only in member and lookup. - --- Also mind the 'nomatch' line in member definition, which is not present in --- lookup and not present in IntMap.hs. That condition stops the search if the --- prefix of current vertex is different that the element looked for. The --- member is correct both with and without this condition. With this condition, --- elements not present are rejected sooner, but a little bit more work is done --- for the elements in the set (we are talking about 3-5% slowdown). Any of --- the solutions is better than the other, because we do not know the --- distribution of input data. Current state is historic. - --- | /O(min(n,W))/. Is the value a member of the set? -member :: Int -> IntSet -> Bool -member x = x `seq` go - where - go (Bin p m l r) - | nomatch x p m = False - | zero x m = go l - | otherwise = go r - go (Tip y) = x == y - go Nil = False - --- | /O(min(n,W))/. Is the element not in the set? -notMember :: Int -> IntSet -> Bool -notMember k = not . member k - --- 'lookup' is used by 'intersection' for left-biasing -lookup :: Int -> IntSet -> Maybe Int -lookup k = k `seq` go - where - go (Bin _ m l r) - | zero k m = go l - | otherwise = go r - go (Tip kx) - | k == kx = Just kx - | otherwise = Nothing - go Nil = Nothing - -{-------------------------------------------------------------------- - Construction ---------------------------------------------------------------------} --- | /O(1)/. The empty set. -empty :: IntSet -empty - = Nil - --- | /O(1)/. A set of one element. -singleton :: Int -> IntSet -singleton x - = Tip x - -{-------------------------------------------------------------------- - Insert ---------------------------------------------------------------------} --- | /O(min(n,W))/. Add a value to the set. When the value is already --- an element of the set, it is replaced by the new one, ie. 'insert' --- is left-biased. -insert :: Int -> IntSet -> IntSet -insert x t = x `seq` - case t of - Bin p m l r - | nomatch x p m -> join x (Tip x) p t - | zero x m -> Bin p m (insert x l) r - | otherwise -> Bin p m l (insert x r) - Tip y - | x==y -> Tip x - | otherwise -> join x (Tip x) y t - Nil -> Tip x - --- right-biased insertion, used by 'union' -insertR :: Int -> IntSet -> IntSet -insertR x t = x `seq` - case t of - Bin p m l r - | nomatch x p m -> join x (Tip x) p t - | zero x m -> Bin p m (insert x l) r - | otherwise -> Bin p m l (insert x r) - Tip y - | x==y -> t - | otherwise -> join x (Tip x) y t - Nil -> Tip x - --- | /O(min(n,W))/. Delete a value in the set. Returns the --- original set when the value was not present. -delete :: Int -> IntSet -> IntSet -delete x t = x `seq` - case t of - Bin p m l r - | nomatch x p m -> t - | zero x m -> bin p m (delete x l) r - | otherwise -> bin p m l (delete x r) - Tip y - | x==y -> Nil - | otherwise -> t - Nil -> Nil - - -{-------------------------------------------------------------------- - Union ---------------------------------------------------------------------} --- | The union of a list of sets. -unions :: [IntSet] -> IntSet -unions xs - = foldlStrict union empty xs - - --- | /O(n+m)/. The union of two sets. -union :: IntSet -> IntSet -> IntSet -union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = union1 - | shorter m2 m1 = union2 - | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2) - | otherwise = join p1 t1 p2 t2 - where - union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2 - | zero p2 m1 = Bin p1 m1 (union l1 t2) r1 - | otherwise = Bin p1 m1 l1 (union r1 t2) - - union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2 - | zero p1 m2 = Bin p2 m2 (union t1 l2) r2 - | otherwise = Bin p2 m2 l2 (union t1 r2) - -union (Tip x) t = insert x t -union t (Tip x) = insertR x t -- right bias -union Nil t = t -union t Nil = t - - -{-------------------------------------------------------------------- - Difference ---------------------------------------------------------------------} --- | /O(n+m)/. Difference between two sets. -difference :: IntSet -> IntSet -> IntSet -difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = difference1 - | shorter m2 m1 = difference2 - | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2) - | otherwise = t1 - where - difference1 | nomatch p2 p1 m1 = t1 - | zero p2 m1 = bin p1 m1 (difference l1 t2) r1 - | otherwise = bin p1 m1 l1 (difference r1 t2) - - difference2 | nomatch p1 p2 m2 = t1 - | zero p1 m2 = difference t1 l2 - | otherwise = difference t1 r2 - -difference t1@(Tip x) t2 - | member x t2 = Nil - | otherwise = t1 - -difference Nil _ = Nil -difference t (Tip x) = delete x t -difference t Nil = t - - - -{-------------------------------------------------------------------- - Intersection ---------------------------------------------------------------------} --- | /O(n+m)/. The intersection of two sets. -intersection :: IntSet -> IntSet -> IntSet -intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = intersection1 - | shorter m2 m1 = intersection2 - | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2) - | otherwise = Nil - where - intersection1 | nomatch p2 p1 m1 = Nil - | zero p2 m1 = intersection l1 t2 - | otherwise = intersection r1 t2 - - intersection2 | nomatch p1 p2 m2 = Nil - | zero p1 m2 = intersection t1 l2 - | otherwise = intersection t1 r2 - -intersection t1@(Tip x) t2 - | member x t2 = t1 - | otherwise = Nil -intersection t (Tip x) - = case lookup x t of - Just y -> Tip y - Nothing -> Nil -intersection Nil _ = Nil -intersection _ Nil = Nil - - - -{-------------------------------------------------------------------- - Subset ---------------------------------------------------------------------} --- | /O(n+m)/. Is this a proper subset? (ie. a subset but not equal). -isProperSubsetOf :: IntSet -> IntSet -> Bool -isProperSubsetOf t1 t2 - = case subsetCmp t1 t2 of - LT -> True - _ -> False - -subsetCmp :: IntSet -> IntSet -> Ordering -subsetCmp t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - | shorter m1 m2 = GT - | shorter m2 m1 = case subsetCmpLt of - GT -> GT - _ -> LT - | p1 == p2 = subsetCmpEq - | otherwise = GT -- disjoint - where - subsetCmpLt | nomatch p1 p2 m2 = GT - | zero p1 m2 = subsetCmp t1 l2 - | otherwise = subsetCmp t1 r2 - subsetCmpEq = case (subsetCmp l1 l2, subsetCmp r1 r2) of - (GT,_ ) -> GT - (_ ,GT) -> GT - (EQ,EQ) -> EQ - _ -> LT - -subsetCmp (Bin _ _ _ _) _ = GT -subsetCmp (Tip x) (Tip y) - | x==y = EQ - | otherwise = GT -- disjoint -subsetCmp (Tip x) t - | member x t = LT - | otherwise = GT -- disjoint -subsetCmp Nil Nil = EQ -subsetCmp Nil _ = LT - --- | /O(n+m)/. Is this a subset? --- @(s1 `isSubsetOf` s2)@ tells whether @s1@ is a subset of @s2@. - -isSubsetOf :: IntSet -> IntSet -> Bool -isSubsetOf t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - | shorter m1 m2 = False - | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubsetOf t1 l2 - else isSubsetOf t1 r2) - | otherwise = (p1==p2) && isSubsetOf l1 l2 && isSubsetOf r1 r2 -isSubsetOf (Bin _ _ _ _) _ = False -isSubsetOf (Tip x) t = member x t -isSubsetOf Nil _ = True - - -{-------------------------------------------------------------------- - Filter ---------------------------------------------------------------------} --- | /O(n)/. Filter all elements that satisfy some predicate. -filter :: (Int -> Bool) -> IntSet -> IntSet -filter predicate t - = case t of - Bin p m l r - -> bin p m (filter predicate l) (filter predicate r) - Tip x - | predicate x -> t - | otherwise -> Nil - Nil -> Nil - --- | /O(n)/. partition the set according to some predicate. -partition :: (Int -> Bool) -> IntSet -> (IntSet,IntSet) -partition predicate t - = case t of - Bin p m l r - -> let (l1,l2) = partition predicate l - (r1,r2) = partition predicate r - in (bin p m l1 r1, bin p m l2 r2) - Tip x - | predicate x -> (t,Nil) - | otherwise -> (Nil,t) - Nil -> (Nil,Nil) - - --- | /O(min(n,W))/. 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 3 (fromList [1..5]) == (fromList [1,2], fromList [4,5]) -split :: Int -> IntSet -> (IntSet,IntSet) -split x t - = case t of - Bin _ m l r - | m < 0 -> if x >= 0 then let (lt,gt) = split' x l in (union r lt, gt) - else let (lt,gt) = split' x r in (lt, union gt l) - -- handle negative numbers. - | otherwise -> split' x t - Tip y - | x>y -> (t,Nil) - | x (Nil,t) - | otherwise -> (Nil,Nil) - Nil -> (Nil, Nil) - -split' :: Int -> IntSet -> (IntSet,IntSet) -split' x t - = case t of - Bin p m l r - | match x p m -> if zero x m then let (lt,gt) = split' x l in (lt,union gt r) - else let (lt,gt) = split' x r in (union l lt,gt) - | otherwise -> if x < p then (Nil, t) - else (t, Nil) - Tip y - | x>y -> (t,Nil) - | x (Nil,t) - | otherwise -> (Nil,Nil) - Nil -> (Nil,Nil) - --- | /O(min(n,W))/. Performs a 'split' but also returns whether the pivot --- element was found in the original set. -splitMember :: Int -> IntSet -> (IntSet,Bool,IntSet) -splitMember x t - = case t of - Bin _ m l r - | m < 0 -> if x >= 0 then let (lt,found,gt) = splitMember' x l in (union r lt, found, gt) - else let (lt,found,gt) = splitMember' x r in (lt, found, union gt l) - -- handle negative numbers. - | otherwise -> splitMember' x t - Tip y - | x>y -> (t,False,Nil) - | x (Nil,False,t) - | otherwise -> (Nil,True,Nil) - Nil -> (Nil,False,Nil) - -splitMember' :: Int -> IntSet -> (IntSet,Bool,IntSet) -splitMember' x t - = case t of - Bin p m l r - | match x p m -> if zero x m then let (lt,found,gt) = splitMember x l in (lt,found,union gt r) - else let (lt,found,gt) = splitMember x r in (union l lt,found,gt) - | otherwise -> if x < p then (Nil, False, t) - else (t, False, Nil) - Tip y - | x>y -> (t,False,Nil) - | x (Nil,False,t) - | otherwise -> (Nil,True,Nil) - Nil -> (Nil,False,Nil) - -{---------------------------------------------------------------------- - Min/Max -----------------------------------------------------------------------} - --- | /O(min(n,W))/. Retrieves the maximal key of the set, and the set --- stripped of that element, or 'Nothing' if passed an empty set. -maxView :: IntSet -> Maybe (Int, IntSet) -maxView t - = case t of - Bin p m l r | m < 0 -> let (result,t') = maxViewUnsigned l in Just (result, bin p m t' r) - Bin p m l r -> let (result,t') = maxViewUnsigned r in Just (result, bin p m l t') - Tip y -> Just (y,Nil) - Nil -> Nothing - -maxViewUnsigned :: IntSet -> (Int, IntSet) -maxViewUnsigned t - = case t of - Bin p m l r -> let (result,t') = maxViewUnsigned r in (result, bin p m l t') - Tip y -> (y, Nil) - Nil -> error "maxViewUnsigned Nil" - --- | /O(min(n,W))/. Retrieves the minimal key of the set, and the set --- stripped of that element, or 'Nothing' if passed an empty set. -minView :: IntSet -> Maybe (Int, IntSet) -minView t - = case t of - Bin p m l r | m < 0 -> let (result,t') = minViewUnsigned r in Just (result, bin p m l t') - Bin p m l r -> let (result,t') = minViewUnsigned l in Just (result, bin p m t' r) - Tip y -> Just (y, Nil) - Nil -> Nothing - -minViewUnsigned :: IntSet -> (Int, IntSet) -minViewUnsigned t - = case t of - Bin p m l r -> let (result,t') = minViewUnsigned l in (result, bin p m t' r) - Tip y -> (y, Nil) - Nil -> error "minViewUnsigned Nil" - --- | /O(min(n,W))/. Delete and find the minimal element. --- --- > deleteFindMin set = (findMin set, deleteMin set) -deleteFindMin :: IntSet -> (Int, IntSet) -deleteFindMin = fromMaybe (error "deleteFindMin: empty set has no minimal element") . minView - --- | /O(min(n,W))/. Delete and find the maximal element. --- --- > deleteFindMax set = (findMax set, deleteMax set) -deleteFindMax :: IntSet -> (Int, IntSet) -deleteFindMax = fromMaybe (error "deleteFindMax: empty set has no maximal element") . maxView - - --- | /O(min(n,W))/. The minimal element of the set. -findMin :: IntSet -> Int -findMin Nil = error "findMin: empty set has no minimal element" -findMin (Tip x) = x -findMin (Bin _ m l r) - | m < 0 = find r - | otherwise = find l - where find (Tip x) = x - find (Bin _ _ l' _) = find l' - find Nil = error "findMin Nil" - --- | /O(min(n,W))/. The maximal element of a set. -findMax :: IntSet -> Int -findMax Nil = error "findMax: empty set has no maximal element" -findMax (Tip x) = x -findMax (Bin _ m l r) - | m < 0 = find l - | otherwise = find r - where find (Tip x) = x - find (Bin _ _ _ r') = find r' - find Nil = error "findMax Nil" - - --- | /O(min(n,W))/. Delete the minimal element. -deleteMin :: IntSet -> IntSet -deleteMin = maybe (error "deleteMin: empty set has no minimal element") snd . minView - --- | /O(min(n,W))/. Delete the maximal element. -deleteMax :: IntSet -> IntSet -deleteMax = maybe (error "deleteMax: empty set has no maximal element") snd . maxView - -{---------------------------------------------------------------------- - Map -----------------------------------------------------------------------} - --- | /O(n*min(n,W))/. --- @'map' f s@ is the set obtained by applying @f@ to each element of @s@. --- --- It's worth noting that the size of the result may be smaller if, --- for some @(x,y)@, @x \/= y && f x == f y@ - -map :: (Int->Int) -> IntSet -> IntSet -map f = fromList . List.map f . toList - -{-------------------------------------------------------------------- - Fold ---------------------------------------------------------------------} --- | /O(n)/. Fold the elements in the set using the given right-associative --- binary operator. This function is an equivalent of 'foldr' and is present --- for compatibility only. --- --- /Please note that fold will be deprecated in the future and removed./ -fold :: (Int -> b -> b) -> b -> IntSet -> b -fold = foldr -{-# INLINE fold #-} - --- | /O(n)/. Fold the elements in the set using the given right-associative --- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'toAscList'@. --- --- For example, --- --- > toAscList set = foldr (:) [] set -foldr :: (Int -> b -> b) -> b -> IntSet -> b -foldr f z t = - case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before - _ -> go z t - where - go z' Nil = z' - go z' (Tip x) = f x z' - go z' (Bin _ _ l r) = go (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' :: (Int -> b -> b) -> b -> IntSet -> b -foldr' f z t = - case t of Bin 0 m l r | m < 0 -> go (go z l) r -- put negative numbers before - _ -> go z t - where - STRICT_1_OF_2(go) - go z' Nil = z' - go z' (Tip x) = f x z' - go z' (Bin _ _ l r) = go (go z' r) l -{-# INLINE foldr' #-} - --- | /O(n)/. Fold the elements in the set using the given left-associative --- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'toAscList'@. --- --- For example, --- --- > toDescList set = foldl (flip (:)) [] set -foldl :: (a -> Int -> a) -> a -> IntSet -> a -foldl f z t = - case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before - _ -> go z t - where - STRICT_1_OF_2(go) - go z' Nil = z' - go z' (Tip x) = f z' x - go z' (Bin _ _ l r) = go (go z' l) 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 -> Int -> a) -> a -> IntSet -> a -foldl' f z t = - case t of Bin 0 m l r | m < 0 -> go (go z r) l -- put negative numbers before - _ -> go z t - where - STRICT_1_OF_2(go) - go z' Nil = z' - go z' (Tip x) = f z' x - go z' (Bin _ _ l r) = go (go z' l) r -{-# INLINE foldl' #-} - -{-------------------------------------------------------------------- - List variations ---------------------------------------------------------------------} --- | /O(n)/. The elements of a set. (For sets, this is equivalent to toList) -elems :: IntSet -> [Int] -elems s - = toList s - -{-------------------------------------------------------------------- - Lists ---------------------------------------------------------------------} --- | /O(n)/. Convert the set to a list of elements. -toList :: IntSet -> [Int] -toList t - = fold (:) [] t - --- | /O(n)/. Convert the set to an ascending list of elements. -toAscList :: IntSet -> [Int] -toAscList t = toList t - --- | /O(n*min(n,W))/. Create a set from a list of integers. -fromList :: [Int] -> IntSet -fromList xs - = foldlStrict ins empty xs - where - ins t x = insert x t - --- | /O(n)/. Build a set from an ascending list of elements. --- /The precondition (input list is ascending) is not checked./ -fromAscList :: [Int] -> IntSet -fromAscList [] = Nil -fromAscList (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0) - where - combineEq x' [] = [x'] - combineEq x' (x:xs) - | x==x' = combineEq x' xs - | otherwise = x' : combineEq x xs - --- | /O(n)/. Build a set from an ascending list of distinct elements. --- /The precondition (input list is strictly ascending) is not checked./ -fromDistinctAscList :: [Int] -> IntSet -fromDistinctAscList [] = Nil -fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada - where - work x [] stk = finish x (Tip x) stk - work x (z:zs) stk = reduce z zs (branchMask z x) x (Tip x) stk - - 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 - pxy = mask px mxy - in if shorter m mxy - then reduce z zs m pxy (Bin pxy mxy ty tx) stk' - else work z zs (Push px tx stk) - - finish _ t Nada = t - finish px tx (Push py ty stk) = finish p (join py ty px tx) stk - where m = branchMask px py - p = mask px m - -data Stack = Push {-# UNPACK #-} !Prefix !IntSet !Stack | Nada - - -{-------------------------------------------------------------------- - Eq ---------------------------------------------------------------------} -instance Eq IntSet where - t1 == t2 = equal t1 t2 - t1 /= t2 = nequal t1 t2 - -equal :: IntSet -> IntSet -> Bool -equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2) -equal (Tip x) (Tip y) - = (x==y) -equal Nil Nil = True -equal _ _ = False - -nequal :: IntSet -> IntSet -> Bool -nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2) -nequal (Tip x) (Tip y) - = (x/=y) -nequal Nil Nil = False -nequal _ _ = True - -{-------------------------------------------------------------------- - Ord ---------------------------------------------------------------------} - -instance Ord IntSet where - compare s1 s2 = compare (toAscList s1) (toAscList s2) - -- tentative implementation. See if more efficient exists. - -{-------------------------------------------------------------------- - Show ---------------------------------------------------------------------} -instance Show IntSet where - showsPrec p xs = showParen (p > 10) $ - showString "fromList " . shows (toList xs) - -{- -XXX unused code -showSet :: [Int] -> ShowS -showSet [] - = showString "{}" -showSet (x:xs) - = showChar '{' . shows x . showTail xs - where - showTail [] = showChar '}' - showTail (x':xs') = showChar ',' . shows x' . showTail xs' --} - -{-------------------------------------------------------------------- - Read ---------------------------------------------------------------------} -instance Read IntSet 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 - -{-------------------------------------------------------------------- - Typeable ---------------------------------------------------------------------} - -#include "Typeable.h" -INSTANCE_TYPEABLE0(IntSet,intSetTc,"IntSet") - -{-------------------------------------------------------------------- - NFData ---------------------------------------------------------------------} - --- The IntSet constructors consist only of strict fields of Ints and --- IntSets, thus the default NFData instance which evaluates to whnf --- should suffice -instance NFData IntSet - -{-------------------------------------------------------------------- - Debugging ---------------------------------------------------------------------} --- | /O(n)/. Show the tree that implements the set. The tree is shown --- in a compressed, hanging format. -showTree :: IntSet -> String -showTree s - = showTreeWith True False s - - -{- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows - the tree that implements the set. 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. --} -showTreeWith :: Bool -> Bool -> IntSet -> String -showTreeWith hang wide t - | hang = (showsTreeHang wide [] t) "" - | otherwise = (showsTree wide [] [] t) "" - -showsTree :: Bool -> [String] -> [String] -> IntSet -> ShowS -showsTree wide lbars rbars t - = case t of - Bin p m l r - -> showsTree wide (withBar rbars) (withEmpty rbars) r . - showWide wide rbars . - showsBars lbars . showString (showBin p m) . showString "\n" . - showWide wide lbars . - showsTree wide (withEmpty lbars) (withBar lbars) l - Tip x - -> showsBars lbars . showString " " . shows x . showString "\n" - Nil -> showsBars lbars . showString "|\n" - -showsTreeHang :: Bool -> [String] -> IntSet -> ShowS -showsTreeHang wide bars t - = case t of - Bin p m l r - -> showsBars bars . showString (showBin p m) . showString "\n" . - showWide wide bars . - showsTreeHang wide (withBar bars) l . - showWide wide bars . - showsTreeHang wide (withEmpty bars) r - Tip x - -> showsBars bars . showString " " . shows x . showString "\n" - Nil -> showsBars bars . showString "|\n" - -showBin :: Prefix -> Mask -> String -showBin _ _ - = "*" -- ++ show (p,m) - -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 - - -{-------------------------------------------------------------------- - Helpers ---------------------------------------------------------------------} -{-------------------------------------------------------------------- - Join ---------------------------------------------------------------------} -join :: Prefix -> IntSet -> Prefix -> IntSet -> IntSet -join p1 t1 p2 t2 - | zero p1 m = Bin p m t1 t2 - | otherwise = Bin p m t2 t1 - where - m = branchMask p1 p2 - p = mask p1 m -{-# INLINE join #-} - -{-------------------------------------------------------------------- - @bin@ assures that we never have empty trees within a tree. ---------------------------------------------------------------------} -bin :: Prefix -> Mask -> IntSet -> IntSet -> IntSet -bin _ _ l Nil = l -bin _ _ Nil r = r -bin p m l r = Bin p m l r -{-# INLINE bin #-} - - -{-------------------------------------------------------------------- - Endian independent bit twiddling ---------------------------------------------------------------------} -zero :: Int -> Mask -> Bool -zero i m - = (natFromInt i) .&. (natFromInt m) == 0 -{-# INLINE zero #-} - -nomatch,match :: Int -> Prefix -> Mask -> Bool -nomatch i p m - = (mask i m) /= p -{-# INLINE nomatch #-} - -match i p m - = (mask i m) == p -{-# INLINE match #-} - --- Suppose a is largest such that 2^a divides 2*m. --- Then mask i m is i with the low a bits zeroed out. -mask :: Int -> Mask -> Prefix -mask i m - = maskW (natFromInt i) (natFromInt m) -{-# INLINE mask #-} - -{-------------------------------------------------------------------- - Big endian operations ---------------------------------------------------------------------} -maskW :: Nat -> Nat -> Prefix -maskW i m - = intFromNat (i .&. (complement (m-1) `xor` m)) -{-# INLINE maskW #-} - -shorter :: Mask -> Mask -> Bool -shorter m1 m2 - = (natFromInt m1) > (natFromInt m2) -{-# INLINE shorter #-} - -branchMask :: Prefix -> Prefix -> Mask -branchMask p1 p2 - = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2)) -{-# INLINE branchMask #-} - -{---------------------------------------------------------------------- - Finding the highest bit (mask) in a word [x] can be done efficiently in - three ways: - * convert to a floating point value and the mantissa tells us the - [log2(x)] that corresponds with the highest bit position. The mantissa - is retrieved either via the standard C function [frexp] or by some bit - twiddling on IEEE compatible numbers (float). Note that one needs to - use at least [double] precision for an accurate mantissa of 32 bit - numbers. - * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit). - * use processor specific assembler instruction (asm). - - The most portable way would be [bit], but is it efficient enough? - I have measured the cycle counts of the different methods on an AMD - Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction: - - highestBitMask: method cycles - -------------- - frexp 200 - float 33 - bit 11 - asm 12 - - highestBit: method cycles - -------------- - frexp 195 - float 33 - bit 11 - asm 11 - - Wow, the bit twiddling is on today's RISC like machines even faster - than a single CISC instruction (BSR)! -----------------------------------------------------------------------} - -{---------------------------------------------------------------------- - [highestBitMask] returns a word where only the highest bit is set. - It is found by first setting all bits in lower positions than the - highest bit and than taking an exclusive or with the original value. - Allthough the function may look expensive, GHC compiles this into - excellent C code that subsequently compiled into highly efficient - machine code. The algorithm is derived from Jorg Arndt's FXT library. -----------------------------------------------------------------------} -highestBitMask :: Nat -> Nat -highestBitMask x0 - = case (x0 .|. shiftRL x0 1) of - x1 -> case (x1 .|. shiftRL x1 2) of - x2 -> case (x2 .|. shiftRL x2 4) of - x3 -> case (x3 .|. shiftRL x3 8) of - x4 -> case (x4 .|. shiftRL x4 16) of - x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms - x6 -> (x6 `xor` (shiftRL x6 1)) -{-# INLINE highestBitMask #-} - - -{-------------------------------------------------------------------- - Utilities ---------------------------------------------------------------------} -foldlStrict :: (a -> b -> a) -> a -> [b] -> a -foldlStrict f = go - where - go z [] = z - go z (x:xs) = let z' = f z x in z' `seq` go z' xs -{-# INLINE foldlStrict #-} +{-# LANGUAGE CPP #-} +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Safe #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.IntSet +-- Copyright : (c) Daan Leijen 2002 +-- (c) Joachim Breitner 2011 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- An efficient implementation of integer sets. +-- +-- These modules are intended to be imported qualified, to avoid name +-- clashes with Prelude functions, e.g. +-- +-- > import Data.IntSet (IntSet) +-- > import qualified Data.IntSet as IntSet +-- +-- The implementation is based on /big-endian patricia trees/. This data +-- structure performs especially well on binary operations like 'union' +-- and 'intersection'. However, my benchmarks show that it is also +-- (much) faster on insertions and deletions when compared to a generic +-- size-balanced set implementation (see "Data.Set"). +-- +-- * 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), +-- October 1968, pages 514-534. +-- +-- Additionally, this implementation places bitmaps in the leaves of the tree. +-- Their size is the natural size of a machine word (32 or 64 bits) and greatly +-- reduce memory footprint and execution times for dense sets, e.g. sets where +-- it is likely that many values lie close to each other. The asymptotics are +-- not affected by this optimization. +-- +-- Many operations have a worst-case complexity of /O(min(n,W))/. +-- This means that the operation can become linear in the number of +-- elements with a maximum of /W/ -- the number of bits in an 'Int' +-- (32 or 64). +----------------------------------------------------------------------------- + +module Data.IntSet ( + -- * Strictness properties + -- $strictness + + -- * Set type +#if !defined(TESTING) + IntSet -- instance Eq,Show +#else + IntSet(..) -- instance Eq,Show +#endif + + -- * Operators + , (\\) + + -- * Query + , IS.null + , size + , member + , notMember + , lookupLT + , lookupGT + , lookupLE + , lookupGE + , isSubsetOf + , isProperSubsetOf + + -- * Construction + , empty + , singleton + , insert + , delete + + -- * Combine + , union + , unions + , difference + , intersection + + -- * Filter + , IS.filter + , partition + , split + , splitMember + + -- * Map + , IS.map + + -- * Folds + , IS.foldr + , IS.foldl + -- ** Strict folds + , foldr' + , foldl' + -- ** Legacy folds + , fold + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , maxView + , minView + + -- * Conversion + + -- ** List + , elems + , toList + , fromList + + -- ** Ordered list + , toAscList + , toDescList + , fromAscList + , fromDistinctAscList + + -- * Debugging + , showTree + , showTreeWith + +#if defined(TESTING) + -- * Internals + , match +#endif + ) where + +import Data.IntSet.Base as IS + +-- $strictness +-- +-- This module satisfies the following strictness property: +-- +-- * Key arguments are evaluated to WHNF +-- +-- Here are some examples that illustrate the property: +-- +-- > delete undefined s == undefined diff --git a/Data/IntSet/Base.hs b/Data/IntSet/Base.hs new file mode 100644 --- /dev/null +++ b/Data/IntSet/Base.hs @@ -0,0 +1,1485 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE MagicHash, BangPatterns, DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.IntSet.Base +-- Copyright : (c) Daan Leijen 2002 +-- (c) Joachim Breitner 2011 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- An efficient implementation of integer sets. +-- +-- These modules are intended to be imported qualified, to avoid name +-- clashes with Prelude functions, e.g. +-- +-- > import Data.IntSet (IntSet) +-- > import qualified Data.IntSet as IntSet +-- +-- The implementation is based on /big-endian patricia trees/. This data +-- structure performs especially well on binary operations like 'union' +-- and 'intersection'. However, my benchmarks show that it is also +-- (much) faster on insertions and deletions when compared to a generic +-- size-balanced set implementation (see "Data.Set"). +-- +-- * 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), +-- October 1968, pages 514-534. +-- +-- Additionally, this implementation places bitmaps in the leaves of the tree. +-- Their size is the natural size of a machine word (32 or 64 bits) and greatly +-- reduce memory footprint and execution times for dense sets, e.g. sets where +-- it is likely that many values lie close to each other. The asymptotics are +-- not affected by this optimization. +-- +-- Many operations have a worst-case complexity of /O(min(n,W))/. +-- This means that the operation can become linear in the number of +-- elements with a maximum of /W/ -- the number of bits in an 'Int' +-- (32 or 64). +----------------------------------------------------------------------------- + +-- [Note: INLINE bit fiddling] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- It is essential that the bit fiddling functions like mask, zero, branchMask +-- etc are inlined. If they do not, the memory allocation skyrockets. The GHC +-- usually gets it right, but it is disastrous if it does not. Therefore we +-- explicitly mark these functions INLINE. + + +-- [Note: Local 'go' functions and capturing] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- Care must be taken when using 'go' function which captures an argument. +-- Sometimes (for example when the argument is passed to a data constructor, +-- as in insert), GHC heap-allocates more than necessary. Therefore C-- code +-- must be checked for increased allocation when creating and modifying such +-- functions. + + +-- [Note: Order of constructors] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- The order of constructors of IntSet matters when considering performance. +-- Currently in GHC 7.0, when type has 3 constructors, they are matched from +-- the first to the last -- the best performance is achieved when the +-- constructors are ordered by frequency. +-- On GHC 7.0, reordering constructors from Nil | Tip | Bin to Bin | Tip | Nil +-- improves the benchmark by circa 10%. + +module Data.IntSet.Base ( + -- * Set type + IntSet(..) -- instance Eq,Show + + -- * Operators + , (\\) + + -- * Query + , null + , size + , member + , notMember + , lookupLT + , lookupGT + , lookupLE + , lookupGE + , isSubsetOf + , isProperSubsetOf + + -- * Construction + , empty + , singleton + , insert + , delete + + -- * Combine + , union + , unions + , difference + , intersection + + -- * Filter + , filter + , partition + , split + , splitMember + + -- * Map + , map + + -- * Folds + , foldr + , foldl + -- ** Strict folds + , foldr' + , foldl' + -- ** Legacy folds + , fold + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , maxView + , minView + + -- * Conversion + + -- ** List + , elems + , toList + , fromList + + -- ** Ordered list + , toAscList + , toDescList + , fromAscList + , fromDistinctAscList + + -- * Debugging + , showTree + , showTreeWith + + -- * Internals + , match + , suffixBitMask + , prefixBitMask + , bitmapOf + ) where + + +import Prelude hiding (filter,foldr,foldl,null,map) +import Data.Bits + +import qualified Data.List as List +import Data.Monoid (Monoid(..)) +import Data.Maybe (fromMaybe) +import Data.Typeable +import Control.DeepSeq (NFData) + +#if __GLASGOW_HASKELL__ +import Text.Read +import Data.Data (Data(..), mkNoRepType) +#endif + +#if __GLASGOW_HASKELL__ +import GHC.Exts ( Word(..), Int(..), build ) +import GHC.Prim ( uncheckedShiftL#, uncheckedShiftRL#, indexInt8OffAddr# ) +#else +import Data.Word +#endif + +-- On GHC, include MachDeps.h to get WORD_SIZE_IN_BITS macro. +#if defined(__GLASGOW_HASKELL__) +#include "MachDeps.h" +#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 + +infixl 9 \\{-This comment teaches CPP correct behaviour -} + +-- A "Nat" is a natural machine word (an unsigned Int) +type Nat = Word + +natFromInt :: Int -> Nat +natFromInt i = fromIntegral i +{-# INLINE natFromInt #-} + +intFromNat :: Nat -> Int +intFromNat w = fromIntegral w +{-# INLINE intFromNat #-} + +-- Right and left logical shifts. +shiftRL, shiftLL :: Nat -> Int -> Nat +#if __GLASGOW_HASKELL__ +{-------------------------------------------------------------------- + GHC: use unboxing to get @shiftRL@ and @shiftLL@ inlined. +--------------------------------------------------------------------} +shiftRL (W# x) (I# i) = W# (uncheckedShiftRL# x i) +shiftLL (W# x) (I# i) = W# (uncheckedShiftL# x i) +#else +shiftRL x i = shiftR x i +shiftLL x i = shiftL x i +#endif +{-# INLINE shiftRL #-} +{-# INLINE shiftLL #-} + +{-------------------------------------------------------------------- + Operators +--------------------------------------------------------------------} +-- | /O(n+m)/. See 'difference'. +(\\) :: IntSet -> IntSet -> IntSet +m1 \\ m2 = difference m1 m2 + +{-------------------------------------------------------------------- + Types +--------------------------------------------------------------------} + +-- | A set of integers. + +-- See Note: Order of constructors +data IntSet = Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !IntSet !IntSet +-- Invariant: Nil is never found as a child of Bin. +-- Invariant: The Mask is a power of 2. It is the largest bit position at which +-- two elements of the set differ. +-- Invariant: Prefix is the common high-order bits that all elements share to +-- the left of the Mask bit. +-- Invariant: In Bin prefix mask left right, left consists of the elements that +-- don't have the mask bit set; right is all the elements that do. + | Tip {-# UNPACK #-} !Prefix {-# UNPACK #-} !BitMap +-- Invariant: The Prefix is zero for all but the last 5 (on 32 bit arches) or 6 +-- bits (on 64 bit arches). The values of the map represented by a tip +-- are the prefix plus the indices of the set bits in the bit map. + | Nil + +-- A number stored in a set is stored as +-- * Prefix (all but last 5-6 bits) and +-- * BitMap (last 5-6 bits stored as a bitmask) +-- Last 5-6 bits are called a Suffix. + +type Prefix = Int +type Mask = Int +type BitMap = Word + +instance Monoid IntSet where + mempty = empty + mappend = union + mconcat = unions + +#if __GLASGOW_HASKELL__ + +{-------------------------------------------------------------------- + A Data instance +--------------------------------------------------------------------} + +-- This instance preserves data abstraction at the cost of inefficiency. +-- We omit reflection services for the sake of data abstraction. + +instance Data IntSet where + gfoldl f z is = z fromList `f` (toList is) + toConstr _ = error "toConstr" + gunfold _ _ = error "gunfold" + dataTypeOf _ = mkNoRepType "Data.IntSet.IntSet" + +#endif + +{-------------------------------------------------------------------- + Query +--------------------------------------------------------------------} +-- | /O(1)/. Is the set empty? +null :: IntSet -> Bool +null Nil = True +null _ = False +{-# INLINE null #-} + +-- | /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 + +-- | /O(min(n,W))/. Is the value a member of the set? + +-- See Note: Local 'go' functions and capturing] +member :: Int -> IntSet -> Bool +member x = x `seq` go + where + go (Bin p m l r) + | nomatch x p m = False + | zero x m = go l + | otherwise = go r + go (Tip y bm) = prefixOf x == y && bitmapOf x .&. bm /= 0 + go Nil = False + +-- | /O(min(n,W))/. Is the element not in the set? +notMember :: Int -> IntSet -> Bool +notMember k = not . member k + +-- | /O(log n)/. Find largest element smaller than the given one. +-- +-- > lookupLT 3 (fromList [3, 5]) == Nothing +-- > lookupLT 5 (fromList [3, 5]) == Just 3 + +-- See Note: Local 'go' functions and capturing. +lookupLT :: Int -> IntSet -> Maybe Int +lookupLT x t = x `seq` case t of + Bin _ m l r | m < 0 -> if x >= 0 then go r l else go Nil r + _ -> go Nil t + where + go def (Bin p m l r) | nomatch x p m = if x < p then unsafeFindMax def else unsafeFindMax r + | zero x m = go def l + | otherwise = go l r + go def (Tip kx bm) | prefixOf x > kx = Just $ kx + highestBitSet bm + | prefixOf x == kx && maskLT /= 0 = Just $ kx + highestBitSet maskLT + | otherwise = unsafeFindMax def + where maskLT = (bitmapOf x - 1) .&. bm + go def Nil = unsafeFindMax def + + +-- | /O(log n)/. Find smallest element greater than the given one. +-- +-- > lookupGT 4 (fromList [3, 5]) == Just 5 +-- > lookupGT 5 (fromList [3, 5]) == Nothing + +-- See Note: Local 'go' functions and capturing. +lookupGT :: Int -> IntSet -> Maybe Int +lookupGT x t = x `seq` case t of + Bin _ m l r | m < 0 -> if x >= 0 then go Nil l else go l r + _ -> go Nil t + where + go def (Bin p m l r) | nomatch x p m = if x < p then unsafeFindMin l else unsafeFindMin def + | zero x m = go r l + | otherwise = go def r + go def (Tip kx bm) | prefixOf x < kx = Just $ kx + lowestBitSet bm + | prefixOf x == kx && maskGT /= 0 = Just $ kx + lowestBitSet maskGT + | otherwise = unsafeFindMin def + where maskGT = (- ((bitmapOf x) `shiftLL` 1)) .&. bm + go def Nil = unsafeFindMin def + + +-- | /O(log n)/. Find largest element smaller or equal to the given one. +-- +-- > lookupLE 2 (fromList [3, 5]) == Nothing +-- > lookupLE 4 (fromList [3, 5]) == Just 3 +-- > lookupLE 5 (fromList [3, 5]) == Just 5 + +-- See Note: Local 'go' functions and capturing. +lookupLE :: Int -> IntSet -> Maybe Int +lookupLE x t = x `seq` case t of + Bin _ m l r | m < 0 -> if x >= 0 then go r l else go Nil r + _ -> go Nil t + where + go def (Bin p m l r) | nomatch x p m = if x < p then unsafeFindMax def else unsafeFindMax r + | zero x m = go def l + | otherwise = go l r + go def (Tip kx bm) | prefixOf x > kx = Just $ kx + highestBitSet bm + | prefixOf x == kx && maskLE /= 0 = Just $ kx + highestBitSet maskLE + | otherwise = unsafeFindMax def + where maskLE = (((bitmapOf x) `shiftLL` 1) - 1) .&. bm + go def Nil = unsafeFindMax def + + +-- | /O(log n)/. Find smallest element greater or equal to the given one. +-- +-- > lookupGE 3 (fromList [3, 5]) == Just 3 +-- > lookupGE 4 (fromList [3, 5]) == Just 5 +-- > lookupGE 6 (fromList [3, 5]) == Nothing + +-- See Note: Local 'go' functions and capturing. +lookupGE :: Int -> IntSet -> Maybe Int +lookupGE x t = x `seq` case t of + Bin _ m l r | m < 0 -> if x >= 0 then go Nil l else go l r + _ -> go Nil t + where + go def (Bin p m l r) | nomatch x p m = if x < p then unsafeFindMin l else unsafeFindMin def + | zero x m = go r l + | otherwise = go def r + go def (Tip kx bm) | prefixOf x < kx = Just $ kx + lowestBitSet bm + | prefixOf x == kx && maskGE /= 0 = Just $ kx + lowestBitSet maskGE + | otherwise = unsafeFindMin def + where maskGE = (- (bitmapOf x)) .&. bm + go def Nil = unsafeFindMin def + + + +-- Helper function for lookupGE and lookupGT. It assumes that if a Bin node is +-- given, it has m > 0. +unsafeFindMin :: IntSet -> Maybe Int +unsafeFindMin Nil = Nothing +unsafeFindMin (Tip kx bm) = Just $ kx + lowestBitSet bm +unsafeFindMin (Bin _ _ l _) = unsafeFindMin l + +-- Helper function for lookupLE and lookupLT. It assumes that if a Bin node is +-- given, it has m > 0. +unsafeFindMax :: IntSet -> Maybe Int +unsafeFindMax Nil = Nothing +unsafeFindMax (Tip kx bm) = Just $ kx + highestBitSet bm +unsafeFindMax (Bin _ _ _ r) = unsafeFindMax r + +{-------------------------------------------------------------------- + Construction +--------------------------------------------------------------------} +-- | /O(1)/. The empty set. +empty :: IntSet +empty + = Nil +{-# INLINE empty #-} + +-- | /O(1)/. A set of one element. +singleton :: Int -> IntSet +singleton x + = Tip (prefixOf x) (bitmapOf x) +{-# INLINE singleton #-} + +{-------------------------------------------------------------------- + Insert +--------------------------------------------------------------------} +-- | /O(min(n,W))/. Add a value to the set. There is no left- or right bias for +-- IntSets. +insert :: Int -> IntSet -> IntSet +insert x = x `seq` 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 -> join 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 -> join kx (Tip kx bm) kx' t + 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 :: Int -> IntSet -> IntSet +delete x = x `seq` 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 + + +{-------------------------------------------------------------------- + Union +--------------------------------------------------------------------} +-- | The union of a list of sets. +unions :: [IntSet] -> IntSet +unions xs + = foldlStrict union empty xs + + +-- | /O(n+m)/. The union of two sets. +union :: IntSet -> IntSet -> IntSet +union t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) + | shorter m1 m2 = union1 + | shorter m2 m1 = union2 + | p1 == p2 = Bin p1 m1 (union l1 l2) (union r1 r2) + | otherwise = join p1 t1 p2 t2 + where + union1 | nomatch p2 p1 m1 = join p1 t1 p2 t2 + | zero p2 m1 = Bin p1 m1 (union l1 t2) r1 + | otherwise = Bin p1 m1 l1 (union r1 t2) + + union2 | nomatch p1 p2 m2 = join p1 t1 p2 t2 + | zero p1 m2 = Bin p2 m2 (union t1 l2) r2 + | otherwise = Bin p2 m2 l2 (union t1 r2) + +union t@(Bin _ _ _ _) (Tip kx bm) = insertBM kx bm t +union t@(Bin _ _ _ _) Nil = t +union (Tip kx bm) t = insertBM kx bm t +union Nil t = t + + +{-------------------------------------------------------------------- + Difference +--------------------------------------------------------------------} +-- | /O(n+m)/. Difference between two sets. +difference :: IntSet -> IntSet -> IntSet +difference t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) + | shorter m1 m2 = difference1 + | shorter m2 m1 = difference2 + | p1 == p2 = bin p1 m1 (difference l1 l2) (difference r1 r2) + | otherwise = t1 + where + difference1 | nomatch p2 p1 m1 = t1 + | zero p2 m1 = bin p1 m1 (difference l1 t2) r1 + | otherwise = bin p1 m1 l1 (difference r1 t2) + + difference2 | nomatch p1 p2 m2 = t1 + | zero p1 m2 = difference t1 l2 + | otherwise = difference t1 r2 + +difference t@(Bin _ _ _ _) (Tip kx bm) = deleteBM kx bm t +difference t@(Bin _ _ _ _) Nil = t + +difference t1@(Tip kx bm) t2 = differenceTip t2 + where differenceTip (Bin p2 m2 l2 r2) | nomatch kx p2 m2 = t1 + | zero kx m2 = differenceTip l2 + | otherwise = differenceTip r2 + differenceTip (Tip kx2 bm2) | kx == kx2 = tip kx (bm .&. complement bm2) + | otherwise = t1 + differenceTip Nil = t1 + +difference Nil _ = Nil + + + +{-------------------------------------------------------------------- + Intersection +--------------------------------------------------------------------} +-- | /O(n+m)/. The intersection of two sets. +intersection :: IntSet -> IntSet -> IntSet +intersection t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) + | shorter m1 m2 = intersection1 + | shorter m2 m1 = intersection2 + | p1 == p2 = bin p1 m1 (intersection l1 l2) (intersection r1 r2) + | otherwise = Nil + where + intersection1 | nomatch p2 p1 m1 = Nil + | zero p2 m1 = intersection l1 t2 + | otherwise = intersection r1 t2 + + intersection2 | nomatch p1 p2 m2 = Nil + | zero p1 m2 = intersection t1 l2 + | otherwise = intersection t1 r2 + +intersection t1@(Bin _ _ _ _) (Tip kx2 bm2) = intersectBM t1 + where intersectBM (Bin p1 m1 l1 r1) | nomatch kx2 p1 m1 = Nil + | zero kx2 m1 = intersectBM l1 + | otherwise = intersectBM r1 + intersectBM (Tip kx1 bm1) | kx1 == kx2 = tip kx1 (bm1 .&. bm2) + | otherwise = Nil + intersectBM Nil = Nil + +intersection (Bin _ _ _ _) Nil = Nil + +intersection (Tip kx1 bm1) t2 = intersectBM t2 + where intersectBM (Bin p2 m2 l2 r2) | nomatch kx1 p2 m2 = Nil + | zero kx1 m2 = intersectBM l2 + | otherwise = intersectBM r2 + intersectBM (Tip kx2 bm2) | kx1 == kx2 = tip kx1 (bm1 .&. bm2) + | otherwise = Nil + intersectBM Nil = Nil + +intersection Nil _ = Nil + +{-------------------------------------------------------------------- + Subset +--------------------------------------------------------------------} +-- | /O(n+m)/. Is this a proper subset? (ie. a subset but not equal). +isProperSubsetOf :: IntSet -> IntSet -> Bool +isProperSubsetOf t1 t2 + = case subsetCmp t1 t2 of + LT -> True + _ -> False + +subsetCmp :: IntSet -> IntSet -> Ordering +subsetCmp t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) + | shorter m1 m2 = GT + | shorter m2 m1 = case subsetCmpLt of + GT -> GT + _ -> LT + | p1 == p2 = subsetCmpEq + | otherwise = GT -- disjoint + where + subsetCmpLt | nomatch p1 p2 m2 = GT + | zero p1 m2 = subsetCmp t1 l2 + | otherwise = subsetCmp t1 r2 + subsetCmpEq = case (subsetCmp l1 l2, subsetCmp r1 r2) of + (GT,_ ) -> GT + (_ ,GT) -> GT + (EQ,EQ) -> EQ + _ -> LT + +subsetCmp (Bin _ _ _ _) _ = GT +subsetCmp (Tip kx1 bm1) (Tip kx2 bm2) + | kx1 /= kx2 = GT -- disjoint + | bm1 == bm2 = EQ + | bm1 .&. complement bm2 == 0 = LT + | otherwise = GT +subsetCmp t1@(Tip kx _) (Bin p m l r) + | nomatch kx p m = GT + | zero kx m = case subsetCmp t1 l of GT -> GT ; _ -> LT + | otherwise = case subsetCmp t1 r of GT -> GT ; _ -> LT +subsetCmp (Tip _ _) Nil = GT -- disjoint +subsetCmp Nil Nil = EQ +subsetCmp Nil _ = LT + +-- | /O(n+m)/. Is this a subset? +-- @(s1 `isSubsetOf` s2)@ tells whether @s1@ is a subset of @s2@. + +isSubsetOf :: IntSet -> IntSet -> Bool +isSubsetOf t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) + | shorter m1 m2 = False + | shorter m2 m1 = match p1 p2 m2 && (if zero p1 m2 then isSubsetOf t1 l2 + else isSubsetOf t1 r2) + | otherwise = (p1==p2) && isSubsetOf l1 l2 && isSubsetOf r1 r2 +isSubsetOf (Bin _ _ _ _) _ = False +isSubsetOf (Tip kx1 bm1) (Tip kx2 bm2) = kx1 == kx2 && bm1 .&. complement bm2 == 0 +isSubsetOf t1@(Tip kx _) (Bin p m l r) + | nomatch kx p m = False + | zero kx m = isSubsetOf t1 l + | otherwise = isSubsetOf t1 r +isSubsetOf (Tip _ _) Nil = False +isSubsetOf Nil _ = True + + +{-------------------------------------------------------------------- + Filter +--------------------------------------------------------------------} +-- | /O(n)/. Filter all elements that satisfy some predicate. +filter :: (Int -> Bool) -> IntSet -> IntSet +filter predicate t + = case t of + Bin p m l r + -> bin p m (filter predicate l) (filter predicate r) + Tip kx bm + -> tip kx (foldl'Bits 0 (bitPred kx) 0 bm) + Nil -> Nil + where bitPred kx bm bi | predicate (kx + bi) = bm .|. bitmapOfSuffix bi + | otherwise = bm + {-# INLINE bitPred #-} + +-- | /O(n)/. partition the set according to some predicate. +partition :: (Int -> Bool) -> IntSet -> (IntSet,IntSet) +partition predicate t + = case t of + Bin p m l r + -> let (l1,l2) = partition predicate l + (r1,r2) = partition predicate r + in (bin p m l1 r1, bin p m l2 r2) + Tip kx bm + -> let bm1 = foldl'Bits 0 (bitPred kx) 0 bm + in (tip kx bm1, tip kx (bm `xor` bm1)) + Nil -> (Nil,Nil) + where bitPred kx bm bi | predicate (kx + bi) = bm .|. bitmapOfSuffix bi + | otherwise = bm + {-# INLINE bitPred #-} + + +-- | /O(min(n,W))/. 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 3 (fromList [1..5]) == (fromList [1,2], fromList [4,5]) +split :: Int -> IntSet -> (IntSet,IntSet) +split x t = + case t of Bin _ m l r | m < 0 -> if x >= 0 then case go x l of (lt, gt) -> (union lt r, gt) + else case go x r of (lt, gt) -> (lt, union gt l) + _ -> go x t + where + go x' t'@(Bin p m l r) | match x' p m = if zero x' m then case go x' l of (lt, gt) -> (lt, union gt r) + else case go x' r of (lt, gt) -> (union lt l, gt) + | otherwise = if x' < p then (Nil, t') + else (t', Nil) + go x' t'@(Tip kx' bm) | kx' > x' = (Nil, t') + -- equivalent to kx' > prefixOf x' + | kx' < prefixOf x' = (t', Nil) + | otherwise = (tip kx' (bm .&. lowerBitmap), tip kx' (bm .&. higherBitmap)) + where lowerBitmap = bitmapOf x' - 1 + higherBitmap = complement (lowerBitmap + bitmapOf x') + go _ Nil = (Nil, Nil) + +-- | /O(min(n,W))/. Performs a 'split' but also returns whether the pivot +-- element was found in the original set. +splitMember :: Int -> IntSet -> (IntSet,Bool,IntSet) +splitMember x t = + case t of Bin _ m l r | m < 0 -> if x >= 0 then case go x l of (lt, fnd, gt) -> (union lt r, fnd, gt) + else case go x r of (lt, fnd, gt) -> (lt, fnd, union gt l) + _ -> go x t + where + go x' t'@(Bin p m l r) | match x' p m = if zero x' m then case go x' l of (lt, fnd, gt) -> (lt, fnd, union gt r) + else case go x' r of (lt, fnd, gt) -> (union lt l, fnd, gt) + | otherwise = if x' < p then (Nil, False, t') + else (t', False, Nil) + go x' t'@(Tip kx' bm) | kx' > x' = (Nil, False, t') + -- equivalent to kx' > prefixOf x' + | kx' < prefixOf x' = (t', False, Nil) + | otherwise = (tip kx' (bm .&. lowerBitmap), (bm .&. bitmapOfx') /= 0, tip kx' (bm .&. higherBitmap)) + where bitmapOfx' = bitmapOf x' + lowerBitmap = bitmapOfx' - 1 + higherBitmap = complement (lowerBitmap + bitmapOfx') + go _ Nil = (Nil, False, Nil) + + +{---------------------------------------------------------------------- + Min/Max +----------------------------------------------------------------------} + +-- | /O(min(n,W))/. Retrieves the maximal key of the set, and the set +-- stripped of that element, or 'Nothing' if passed an empty set. +maxView :: IntSet -> Maybe (Int, IntSet) +maxView 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) + _ -> Just (go t) + where + go (Bin p m l r) = case go r of (result, r') -> (result, bin p m l r') + go (Tip kx bm) = case highestBitSet bm of bi -> (kx + bi, tip kx (bm .&. complement (bitmapOfSuffix bi))) + go Nil = error "maxView Nil" + +-- | /O(min(n,W))/. Retrieves the minimal key of the set, and the set +-- stripped of that element, or 'Nothing' if passed an empty set. +minView :: IntSet -> Maybe (Int, IntSet) +minView 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') + _ -> Just (go t) + where + go (Bin p m l r) = case go l of (result, l') -> (result, bin p m l' r) + go (Tip kx bm) = case lowestBitSet bm of bi -> (kx + bi, tip kx (bm .&. complement (bitmapOfSuffix bi))) + go Nil = error "minView Nil" + +-- | /O(min(n,W))/. Delete and find the minimal element. +-- +-- > deleteFindMin set = (findMin set, deleteMin set) +deleteFindMin :: IntSet -> (Int, IntSet) +deleteFindMin = fromMaybe (error "deleteFindMin: empty set has no minimal element") . minView + +-- | /O(min(n,W))/. Delete and find the maximal element. +-- +-- > deleteFindMax set = (findMax set, deleteMax set) +deleteFindMax :: IntSet -> (Int, IntSet) +deleteFindMax = fromMaybe (error "deleteFindMax: empty set has no maximal element") . maxView + + +-- | /O(min(n,W))/. The minimal element of the set. +findMin :: IntSet -> Int +findMin Nil = error "findMin: empty set has no minimal element" +findMin (Tip kx bm) = kx + lowestBitSet bm +findMin (Bin _ m l r) + | m < 0 = find r + | otherwise = find l + where find (Tip kx bm) = kx + lowestBitSet bm + find (Bin _ _ l' _) = find l' + find Nil = error "findMin Nil" + +-- | /O(min(n,W))/. The maximal element of a set. +findMax :: IntSet -> Int +findMax Nil = error "findMax: empty set has no maximal element" +findMax (Tip kx bm) = kx + highestBitSet bm +findMax (Bin _ m l r) + | m < 0 = find l + | otherwise = find r + where find (Tip kx bm) = kx + highestBitSet bm + find (Bin _ _ _ r') = find r' + find Nil = error "findMax Nil" + + +-- | /O(min(n,W))/. Delete the minimal element. +deleteMin :: IntSet -> IntSet +deleteMin = maybe Nil snd . minView + +-- | /O(min(n,W))/. Delete the maximal element. +deleteMax :: IntSet -> IntSet +deleteMax = maybe Nil snd . maxView + +{---------------------------------------------------------------------- + Map +----------------------------------------------------------------------} + +-- | /O(n*min(n,W))/. +-- @'map' f s@ is the set obtained by applying @f@ to each element of @s@. +-- +-- It's worth noting that the size of the result may be smaller if, +-- for some @(x,y)@, @x \/= y && f x == f y@ + +map :: (Int->Int) -> IntSet -> IntSet +map f = fromList . List.map f . toList + +{-------------------------------------------------------------------- + Fold +--------------------------------------------------------------------} +-- | /O(n)/. Fold the elements in the set using the given right-associative +-- binary operator. This function is an equivalent of 'foldr' and is present +-- for compatibility only. +-- +-- /Please note that fold will be deprecated in the future and removed./ +fold :: (Int -> b -> b) -> b -> IntSet -> b +fold = foldr +{-# INLINE fold #-} + +-- | /O(n)/. Fold the elements in the set using the given right-associative +-- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'toAscList'@. +-- +-- For example, +-- +-- > toAscList set = foldr (:) [] set +foldr :: (Int -> b -> b) -> b -> IntSet -> b +foldr f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t + where + go z' Nil = z' + go z' (Tip kx bm) = foldrBits kx f z' bm + go z' (Bin _ _ l r) = go (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' :: (Int -> b -> b) -> b -> IntSet -> b +foldr' f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t + where + STRICT_1_OF_2(go) + 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' #-} + +-- | /O(n)/. Fold the elements in the set using the given left-associative +-- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'toAscList'@. +-- +-- For example, +-- +-- > toDescList set = foldl (flip (:)) [] set +foldl :: (a -> Int -> a) -> a -> IntSet -> a +foldl f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t + where + STRICT_1_OF_2(go) + go z' Nil = z' + go z' (Tip kx bm) = foldlBits kx f z' bm + go z' (Bin _ _ l r) = go (go z' l) 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 -> Int -> a) -> a -> IntSet -> a +foldl' f z = \t -> -- Use lambda t to be inlinable with two arguments only. + case t of Bin _ m l r | m < 0 -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t + where + STRICT_1_OF_2(go) + 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' #-} + +{-------------------------------------------------------------------- + List variations +--------------------------------------------------------------------} +-- | /O(n)/. An alias of 'toAscList'. The elements of a set in ascending order. +-- Subject to list fusion. +elems :: IntSet -> [Int] +elems + = toAscList + +{-------------------------------------------------------------------- + Lists +--------------------------------------------------------------------} +-- | /O(n)/. Convert the set to a list of elements. Subject to list fusion. +toList :: IntSet -> [Int] +toList + = toAscList + +-- | /O(n)/. Convert the set to an ascending list of elements. Subject to list +-- fusion. +toAscList :: IntSet -> [Int] +toAscList = foldr (:) [] + +-- | /O(n)/. Convert the set to a descending list of elements. Subject to list +-- fusion. +toDescList :: IntSet -> [Int] +toDescList = foldl (flip (:)) [] + +-- List fusion for the list generating functions. +#if __GLASGOW_HASKELL__ +-- The foldrFB and foldlFB are foldr and foldl equivalents, used for list fusion. +-- They are important to convert unfused to{Asc,Desc}List back, see mapFB in prelude. +foldrFB :: (Int -> b -> b) -> b -> IntSet -> b +foldrFB = foldr +{-# INLINE[0] foldrFB #-} +foldlFB :: (a -> Int -> a) -> a -> IntSet -> a +foldlFB = foldl +{-# INLINE[0] foldlFB #-} + +-- Inline elems and toList, so that we need to fuse only toAscList. +{-# INLINE elems #-} +{-# INLINE toList #-} + +-- The fusion is enabled up to phase 2 included. If it does not succeed, +-- convert in phase 1 the expanded to{Asc,Desc}List calls back to +-- 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 to{Asc,Desc}List -- it was forbidden to inline it +-- before phase 0, otherwise the fusion rules would not fire at all. +{-# NOINLINE[0] toAscList #-} +{-# NOINLINE[0] toDescList #-} +{-# RULES "IntSet.toAscList" [~1] forall s . toAscList s = build (\c n -> foldrFB c n s) #-} +{-# RULES "IntSet.toAscListBack" [1] foldrFB (:) [] = toAscList #-} +{-# RULES "IntSet.toDescList" [~1] forall s . toDescList s = build (\c n -> foldlFB (\xs x -> c x xs) n s) #-} +{-# RULES "IntSet.toDescListBack" [1] foldlFB (\xs x -> x : xs) [] = toDescList #-} +#endif + + +-- | /O(n*min(n,W))/. Create a set from a list of integers. +fromList :: [Int] -> IntSet +fromList xs + = foldlStrict ins empty xs + where + ins t x = insert x t + +-- | /O(n)/. Build a set from an ascending list of elements. +-- /The precondition (input list is ascending) is not checked./ +fromAscList :: [Int] -> IntSet +fromAscList [] = Nil +fromAscList (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0) + where + combineEq x' [] = [x'] + combineEq x' (x:xs) + | x==x' = combineEq x' xs + | otherwise = x' : combineEq x xs + +-- | /O(n)/. Build a set from an ascending list of distinct elements. +-- /The precondition (input list is strictly ascending) is not checked./ +fromDistinctAscList :: [Int] -> IntSet +fromDistinctAscList [] = Nil +fromDistinctAscList (z0 : zs0) = work (prefixOf z0) (bitmapOf z0) zs0 Nada + where + -- 'work' accumulates all values that go into one tip, before passing this Tip + -- to 'reduce' + work kx bm [] stk = finish kx (Tip kx bm) stk + work kx bm (z:zs) stk | kx == prefixOf z = work kx (bm .|. bitmapOf z) zs stk + work kx bm (z:zs) stk = reduce z zs (branchMask z kx) kx (Tip kx bm) stk + + reduce z zs _ px tx Nada = work (prefixOf z) (bitmapOf z) zs (Push px tx Nada) + reduce z zs m px tx stk@(Push py ty stk') = + let mxy = branchMask px py + pxy = mask px mxy + in if shorter m mxy + then reduce z zs m pxy (Bin pxy mxy ty tx) stk' + else work (prefixOf z) (bitmapOf z) zs (Push px tx stk) + + finish _ t Nada = t + finish px tx (Push py ty stk) = finish p (join py ty px tx) stk + where m = branchMask px py + p = mask px m + +data Stack = Push {-# UNPACK #-} !Prefix !IntSet !Stack | Nada + + +{-------------------------------------------------------------------- + Eq +--------------------------------------------------------------------} +instance Eq IntSet where + t1 == t2 = equal t1 t2 + t1 /= t2 = nequal t1 t2 + +equal :: IntSet -> IntSet -> Bool +equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) + = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2) +equal (Tip kx1 bm1) (Tip kx2 bm2) + = kx1 == kx2 && bm1 == bm2 +equal Nil Nil = True +equal _ _ = False + +nequal :: IntSet -> IntSet -> Bool +nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) + = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2) +nequal (Tip kx1 bm1) (Tip kx2 bm2) + = kx1 /= kx2 || bm1 /= bm2 +nequal Nil Nil = False +nequal _ _ = True + +{-------------------------------------------------------------------- + Ord +--------------------------------------------------------------------} + +instance Ord IntSet where + compare s1 s2 = compare (toAscList s1) (toAscList s2) + -- tentative implementation. See if more efficient exists. + +{-------------------------------------------------------------------- + Show +--------------------------------------------------------------------} +instance Show IntSet where + showsPrec p xs = showParen (p > 10) $ + showString "fromList " . shows (toList xs) + +{-------------------------------------------------------------------- + Read +--------------------------------------------------------------------} +instance Read IntSet 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 + +{-------------------------------------------------------------------- + Typeable +--------------------------------------------------------------------} + +#include "Typeable.h" +INSTANCE_TYPEABLE0(IntSet,intSetTc,"IntSet") + +{-------------------------------------------------------------------- + NFData +--------------------------------------------------------------------} + +-- The IntSet constructors consist only of strict fields of Ints and +-- IntSets, thus the default NFData instance which evaluates to whnf +-- should suffice +instance NFData IntSet + +{-------------------------------------------------------------------- + Debugging +--------------------------------------------------------------------} +-- | /O(n)/. Show the tree that implements the set. The tree is shown +-- in a compressed, hanging format. +showTree :: IntSet -> String +showTree s + = showTreeWith True False s + + +{- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows + the tree that implements the set. 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. +-} +showTreeWith :: Bool -> Bool -> IntSet -> String +showTreeWith hang wide t + | hang = (showsTreeHang wide [] t) "" + | otherwise = (showsTree wide [] [] t) "" + +showsTree :: Bool -> [String] -> [String] -> IntSet -> ShowS +showsTree wide lbars rbars t + = case t of + Bin p m l r + -> showsTree wide (withBar rbars) (withEmpty rbars) r . + showWide wide rbars . + showsBars lbars . showString (showBin p m) . showString "\n" . + showWide wide lbars . + showsTree wide (withEmpty lbars) (withBar lbars) l + Tip kx bm + -> showsBars lbars . showString " " . shows kx . showString " + " . + showsBitMap bm . showString "\n" + Nil -> showsBars lbars . showString "|\n" + +showsTreeHang :: Bool -> [String] -> IntSet -> ShowS +showsTreeHang wide bars t + = case t of + Bin p m l r + -> showsBars bars . showString (showBin p m) . showString "\n" . + showWide wide bars . + showsTreeHang wide (withBar bars) l . + showWide wide bars . + showsTreeHang wide (withEmpty bars) r + Tip kx bm + -> showsBars bars . showString " " . shows kx . showString " + " . + showsBitMap bm . showString "\n" + Nil -> showsBars bars . showString "|\n" + +showBin :: Prefix -> Mask -> String +showBin _ _ + = "*" -- ++ show (p,m) + +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 + +showsBitMap :: Word -> ShowS +showsBitMap = showString . showBitMap + +showBitMap :: Word -> String +showBitMap w = show $ foldrBits 0 (:) [] w + +node :: String +node = "+--" + +withBar, withEmpty :: [String] -> [String] +withBar bars = "| ":bars +withEmpty bars = " ":bars + + +{-------------------------------------------------------------------- + Helpers +--------------------------------------------------------------------} +{-------------------------------------------------------------------- + Join +--------------------------------------------------------------------} +join :: Prefix -> IntSet -> Prefix -> IntSet -> IntSet +join p1 t1 p2 t2 + | zero p1 m = Bin p m t1 t2 + | otherwise = Bin p m t2 t1 + where + m = branchMask p1 p2 + p = mask p1 m +{-# INLINE join #-} + +{-------------------------------------------------------------------- + @bin@ assures that we never have empty trees within a tree. +--------------------------------------------------------------------} +bin :: Prefix -> Mask -> IntSet -> IntSet -> IntSet +bin _ _ l Nil = l +bin _ _ Nil r = r +bin p m l r = Bin p m l r +{-# INLINE bin #-} + +{-------------------------------------------------------------------- + @tip@ assures that we never have empty bitmaps within a tree. +--------------------------------------------------------------------} +tip :: Prefix -> BitMap -> IntSet +tip _ 0 = Nil +tip kx bm = Tip kx bm +{-# INLINE tip #-} + + +{---------------------------------------------------------------------- + Functions that generate Prefix and BitMap of a Key or a Suffix. +----------------------------------------------------------------------} + +suffixBitMask :: Int +suffixBitMask = bitSize (undefined::Word) - 1 +{-# INLINE suffixBitMask #-} + +prefixBitMask :: Int +prefixBitMask = complement suffixBitMask +{-# INLINE prefixBitMask #-} + +prefixOf :: Int -> Prefix +prefixOf x = x .&. prefixBitMask +{-# INLINE prefixOf #-} + +suffixOf :: Int -> Int +suffixOf x = x .&. suffixBitMask +{-# INLINE suffixOf #-} + +bitmapOfSuffix :: Int -> BitMap +bitmapOfSuffix s = 1 `shiftLL` s +{-# INLINE bitmapOfSuffix #-} + +bitmapOf :: Int -> BitMap +bitmapOf x = bitmapOfSuffix (suffixOf x) +{-# INLINE bitmapOf #-} + + +{-------------------------------------------------------------------- + Endian independent bit twiddling +--------------------------------------------------------------------} +zero :: Int -> Mask -> Bool +zero i m + = (natFromInt i) .&. (natFromInt m) == 0 +{-# INLINE zero #-} + +nomatch,match :: Int -> Prefix -> Mask -> Bool +nomatch i p m + = (mask i m) /= p +{-# INLINE nomatch #-} + +match i p m + = (mask i m) == p +{-# INLINE match #-} + +-- Suppose a is largest such that 2^a divides 2*m. +-- Then mask i m is i with the low a bits zeroed out. +mask :: Int -> Mask -> Prefix +mask i m + = maskW (natFromInt i) (natFromInt m) +{-# INLINE mask #-} + +{-------------------------------------------------------------------- + Big endian operations +--------------------------------------------------------------------} +maskW :: Nat -> Nat -> Prefix +maskW i m + = intFromNat (i .&. (complement (m-1) `xor` m)) +{-# INLINE maskW #-} + +shorter :: Mask -> Mask -> Bool +shorter m1 m2 + = (natFromInt m1) > (natFromInt m2) +{-# INLINE shorter #-} + +branchMask :: Prefix -> Prefix -> Mask +branchMask p1 p2 + = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2)) +{-# INLINE branchMask #-} + +{---------------------------------------------------------------------- + Finding the highest bit (mask) in a word [x] can be done efficiently in + three ways: + * convert to a floating point value and the mantissa tells us the + [log2(x)] that corresponds with the highest bit position. The mantissa + is retrieved either via the standard C function [frexp] or by some bit + twiddling on IEEE compatible numbers (float). Note that one needs to + use at least [double] precision for an accurate mantissa of 32 bit + numbers. + * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit). + * use processor specific assembler instruction (asm). + + The most portable way would be [bit], but is it efficient enough? + I have measured the cycle counts of the different methods on an AMD + Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction: + + highestBitMask: method cycles + -------------- + frexp 200 + float 33 + bit 11 + asm 12 + + highestBit: method cycles + -------------- + frexp 195 + float 33 + bit 11 + asm 11 + + Wow, the bit twiddling is on today's RISC like machines even faster + than a single CISC instruction (BSR)! +----------------------------------------------------------------------} + +{---------------------------------------------------------------------- + [highestBitMask] returns a word where only the highest bit is set. + It is found by first setting all bits in lower positions than the + highest bit and than taking an exclusive or with the original value. + Allthough the function may look expensive, GHC compiles this into + excellent C code that subsequently compiled into highly efficient + machine code. The algorithm is derived from Jorg Arndt's FXT library. +----------------------------------------------------------------------} +highestBitMask :: Nat -> Nat +highestBitMask x0 + = case (x0 .|. shiftRL x0 1) of + x1 -> case (x1 .|. shiftRL x1 2) of + x2 -> case (x2 .|. shiftRL x2 4) of + x3 -> case (x3 .|. shiftRL x3 8) of + x4 -> case (x4 .|. shiftRL x4 16) of +#if !(defined(__GLASGOW_HASKELL__) && WORD_SIZE_IN_BITS==32) + x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms +#endif + x6 -> (x6 `xor` (shiftRL x6 1)) +{-# INLINE highestBitMask #-} + +{---------------------------------------------------------------------- + To get best performance, we provide fast implementations of + lowestBitSet, highestBitSet and fold[lr][l]Bits for GHC. + If the intel bsf and bsr instructions ever become GHC primops, + this code should be reimplemented using these. + + Performance of this code is crucial for folds, toList, filter, partition. + + The signatures of methods in question are placed after this comment. +----------------------------------------------------------------------} + +lowestBitSet :: Nat -> Int +highestBitSet :: Nat -> Int +foldlBits :: Int -> (a -> Int -> a) -> a -> Nat -> a +foldl'Bits :: Int -> (a -> Int -> a) -> a -> Nat -> a +foldrBits :: Int -> (Int -> a -> a) -> a -> Nat -> a +foldr'Bits :: Int -> (Int -> a -> a) -> a -> Nat -> a + +{-# INLINE lowestBitSet #-} +{-# INLINE highestBitSet #-} +{-# INLINE foldlBits #-} +{-# INLINE foldl'Bits #-} +{-# INLINE foldrBits #-} +{-# INLINE foldr'Bits #-} + +#if defined(__GLASGOW_HASKELL__) && (WORD_SIZE_IN_BITS==32 || WORD_SIZE_IN_BITS==64) +{---------------------------------------------------------------------- + For lowestBitSet we use wordsize-dependant implementation based on + multiplication and DeBrujn indeces, which was proposed by Edward Kmett + + + The core of this implementation is fast indexOfTheOnlyBit, + which is given a Nat with exactly one bit set, and returns + its index. + + Lot of effort was put in these implementations, please benchmark carefully + before changing this code. +----------------------------------------------------------------------} + +indexOfTheOnlyBit :: Nat -> Int +{-# INLINE indexOfTheOnlyBit #-} +indexOfTheOnlyBit bitmask = + I# (lsbArray `indexInt8OffAddr#` unboxInt (intFromNat ((bitmask * magic) `shiftRL` offset))) + where unboxInt (I# i) = i +#if WORD_SIZE_IN_BITS==32 + magic = 0x077CB531 + offset = 27 + !lsbArray = "\0\1\28\2\29\14\24\3\30\22\20\15\25\17\4\8\31\27\13\23\21\19\16\7\26\12\18\6\11\5\10\9"# +#else + magic = 0x07EDD5E59A4E28C2 + offset = 58 + !lsbArray = "\63\0\58\1\59\47\53\2\60\39\48\27\54\33\42\3\61\51\37\40\49\18\28\20\55\30\34\11\43\14\22\4\62\57\46\52\38\26\32\41\50\36\17\19\29\10\13\21\56\45\25\31\35\16\9\12\44\24\15\8\23\7\6\5"# +#endif +-- The lsbArray gets inlined to every call site of indexOfTheOnlyBit. +-- That cannot be easily avoided, as GHC forbids top-level Addr# literal. +-- One could go around that by supplying getLsbArray :: () -> Addr# marked +-- as NOINLINE. But the code size of calling it and processing the result +-- is 48B on 32-bit and 56B on 64-bit architectures -- so the 32B and 64B array +-- is actually improvement on 32-bit and only a 8B size increase on 64-bit. + +lowestBitMask :: Nat -> Nat +lowestBitMask x = x .&. negate x +{-# INLINE lowestBitMask #-} + +-- Reverse the order of bits in the Nat. +revNat :: Nat -> Nat +#if WORD_SIZE_IN_BITS==32 +revNat x1 = case ((x1 `shiftRL` 1) .&. 0x55555555) .|. ((x1 .&. 0x55555555) `shiftLL` 1) of + x2 -> case ((x2 `shiftRL` 2) .&. 0x33333333) .|. ((x2 .&. 0x33333333) `shiftLL` 2) of + x3 -> case ((x3 `shiftRL` 4) .&. 0x0F0F0F0F) .|. ((x3 .&. 0x0F0F0F0F) `shiftLL` 4) of + x4 -> case ((x4 `shiftRL` 8) .&. 0x00FF00FF) .|. ((x4 .&. 0x00FF00FF) `shiftLL` 8) of + x5 -> ( x5 `shiftRL` 16 ) .|. ( x5 `shiftLL` 16); +#else +revNat x1 = case ((x1 `shiftRL` 1) .&. 0x5555555555555555) .|. ((x1 .&. 0x5555555555555555) `shiftLL` 1) of + x2 -> case ((x2 `shiftRL` 2) .&. 0x3333333333333333) .|. ((x2 .&. 0x3333333333333333) `shiftLL` 2) of + x3 -> case ((x3 `shiftRL` 4) .&. 0x0F0F0F0F0F0F0F0F) .|. ((x3 .&. 0x0F0F0F0F0F0F0F0F) `shiftLL` 4) of + x4 -> case ((x4 `shiftRL` 8) .&. 0x00FF00FF00FF00FF) .|. ((x4 .&. 0x00FF00FF00FF00FF) `shiftLL` 8) of + x5 -> case ((x5 `shiftRL` 16) .&. 0x0000FFFF0000FFFF) .|. ((x5 .&. 0x0000FFFF0000FFFF) `shiftLL` 16) of + x6 -> ( x6 `shiftRL` 32 ) .|. ( x6 `shiftLL` 32); +#endif + +lowestBitSet x = indexOfTheOnlyBit (lowestBitMask x) + +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)) + +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)) + +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) + +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) + +#else +{---------------------------------------------------------------------- + In general case we use logarithmic implementation of + lowestBitSet and highestBitSet, which works up to bit sizes of 64. + + Folds are linear scans. +----------------------------------------------------------------------} + +lowestBitSet n0 = + let (n1,b1) = if n0 .&. 0xFFFFFFFF /= 0 then (n0,0) else (n0 `shiftRL` 32, 32) + (n2,b2) = if n1 .&. 0xFFFF /= 0 then (n1,b1) else (n1 `shiftRL` 16, 16+b1) + (n3,b3) = if n2 .&. 0xFF /= 0 then (n2,b2) else (n2 `shiftRL` 8, 8+b2) + (n4,b4) = if n3 .&. 0xF /= 0 then (n3,b3) else (n3 `shiftRL` 4, 4+b3) + (n5,b5) = if n4 .&. 0x3 /= 0 then (n4,b4) else (n4 `shiftRL` 2, 2+b4) + b6 = if n5 .&. 0x1 /= 0 then b5 else 1+b5 + in b6 + +highestBitSet n0 = + let (n1,b1) = if n0 .&. 0xFFFFFFFF00000000 /= 0 then (n0 `shiftRL` 32, 32) else (n0,0) + (n2,b2) = if n1 .&. 0xFFFF0000 /= 0 then (n1 `shiftRL` 16, 16+b1) else (n1,b1) + (n3,b3) = if n2 .&. 0xFF00 /= 0 then (n2 `shiftRL` 8, 8+b2) else (n2,b2) + (n4,b4) = if n3 .&. 0xF0 /= 0 then (n3 `shiftRL` 4, 4+b3) else (n3,b3) + (n5,b5) = if n4 .&. 0xC /= 0 then (n4 `shiftRL` 2, 2+b4) else (n4,b4) + b6 = if n5 .&. 0x2 /= 0 then 1+b5 else b5 + in b6 + +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 + 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 + 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 + 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 + go bi n | n `testBit` 0 = f bi $! go (bi + 1) (n `shiftRL` 1) + | otherwise = go (bi + 1) (n `shiftRL` 1) + +#endif + +{---------------------------------------------------------------------- + [bitcount] as posted by David F. Place to haskell-cafe on April 11, 2006, + based on the code on + http://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetKernighan, + where the following source is given: + Published in 1988, the C Programming Language 2nd Ed. (by Brian W. + Kernighan and Dennis M. Ritchie) mentions this in exercise 2-9. On April + 19, 2006 Don Knuth pointed out to me that this method "was first published + by Peter Wegner in CACM 3 (1960), 322. (Also discovered independently by + Derrick Lehmer and published in 1964 in a book edited by Beckenbach.)" +----------------------------------------------------------------------} +bitcount :: Int -> Word -> Int +bitcount a0 x0 = go a0 x0 + where go a 0 = a + go a x = go (a + 1) (x .&. (x-1)) +{-# INLINE bitcount #-} + + +{-------------------------------------------------------------------- + Utilities +--------------------------------------------------------------------} +foldlStrict :: (a -> b -> a) -> a -> [b] -> a +foldlStrict f = go + where + go z [] = z + go z (x:xs) = let z' = f z x in z' `seq` go z' xs +{-# INLINE foldlStrict #-} diff --git a/Data/Map.hs b/Data/Map.hs --- a/Data/Map.hs +++ b/Data/Map.hs @@ -1,2650 +1,104 @@ -{-# LANGUAGE NoBangPatterns #-} -#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 -{-# LANGUAGE Safe #-} -#endif ------------------------------------------------------------------------------ --- | --- Module : Data.Map --- 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 . ------------------------------------------------------------------------------ - --- 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). --- --- 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 that can be INLINE are not recursive -- a 'go' function doing --- the real work is provided. - -module Data.Map ( - -- * 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 - - -- * Construction - , empty - , singleton - - -- ** Insertion - , insert - , insertWith - , insertWith' - , insertWithKey - , insertWithKey' - , insertLookupWithKey - , 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 - - -- * Traversal - -- ** Map - , map - , mapWithKey - , mapAccum - , mapAccumWithKey - , mapAccumRWithKey - , mapKeys - , mapKeysWith - , mapKeysMonotonic - - -- * Folds - , foldr - , foldl - , foldrWithKey - , foldlWithKey - -- ** Strict folds - , foldr' - , foldl' - , foldrWithKey' - , foldlWithKey' - -- ** Legacy folds - , fold - , foldWithKey - - -- * Conversion - , elems - , keys - , keysSet - , assocs - - -- ** Lists - , toList - , fromList - , fromListWith - , fromListWithKey - - -- ** Ordered lists - , toAscList - , toDescList - , fromAscList - , fromAscListWith - , fromAscListWithKey - , fromDistinctAscList - - -- * Filter - , filter - , filterWithKey - , partition - , partitionWithKey - - , mapMaybe - , mapMaybeWithKey - , mapEither - , mapEitherWithKey - - , split - , splitLookup - - -- * 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 - , join - , merge -#endif - - ) where - -import Prelude hiding (lookup,map,filter,foldr,foldl,null) -import qualified Data.Set as Set -import qualified Data.List as List -import Data.Monoid (Monoid(..)) -import Control.Applicative (Applicative(..), (<$>)) -import Data.Traversable (Traversable(traverse)) -import qualified Data.Foldable as Foldable -import Data.Typeable -import Control.DeepSeq (NFData(rnf)) - -#if __GLASGOW_HASKELL__ -import Text.Read -import Data.Data -#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_1_OF_3(fn) fn arg _ _ | arg `seq` False = undefined -#define STRICT_2_OF_3(fn) fn _ arg _ | arg `seq` False = undefined -#define STRICT_2_OF_4(fn) fn _ arg _ _ | arg `seq` False = undefined - -{-------------------------------------------------------------------- - 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 -{-# INLINE (!) #-} - --- | 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@. -data Map k a = Tip - | Bin {-# UNPACK #-} !Size !k a !(Map k a) !(Map k a) - -type Size = Int - -instance (Ord k) => Monoid (Map k v) where - mempty = empty - mappend = union - mconcat = unions - -#if __GLASGOW_HASKELL__ - -{-------------------------------------------------------------------- - A Data instance ---------------------------------------------------------------------} - --- This instance preserves data abstraction at the cost of inefficiency. --- We omit 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 _ = error "toConstr" - gunfold _ _ = error "gunfold" - dataTypeOf _ = mkNoRepType "Data.Map.Map" - dataCast2 f = gcast2 f - -#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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE null #-} -#endif - --- | /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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE size #-} -#endif - - --- | /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 - -lookupAssoc :: Ord k => k -> Map k a -> Maybe (k,a) -lookupAssoc = 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 (kx,x) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINEABLE lookupAssoc #-} -#else -{-# INLINE lookupAssoc #-} -#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 k m = case lookup k m of - Nothing -> False - Just _ -> True -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINEABLE 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 -{-# INLINE notMember #-} - --- | /O(log n)/. Find the value at a key. --- Calls 'error' when the element can not be found. --- Consider using 'lookup' when elements may not be present. -find :: Ord k => k -> Map k a -> a -find k m = case lookup k m of - Nothing -> error "Map.find: element not in the map" - Just x -> 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 def k m = case lookup k m of - Nothing -> def - Just x -> x -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE findWithDefault #-} -#else -{-# INLINE findWithDefault #-} -#endif - -{-------------------------------------------------------------------- - Construction ---------------------------------------------------------------------} --- | /O(1)/. The empty map. --- --- > empty == fromList [] --- > size empty == 0 - -empty :: Map k a -empty = Tip - --- | /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 - -{-------------------------------------------------------------------- - 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' - -insert :: Ord k => k -> a -> Map k a -> Map k a -insert = go - where - 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 -{-# INLINEABLE 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') -{-# INLINE insertWith #-} - --- | Same as 'insertWith', but the combining function is applied strictly. --- This is often the most desirable behavior. --- --- For example, to update a counter: --- --- > insertWith' (+) k 1 m --- -insertWith' :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWith' f = insertWithKey' (\_ x' y' -> f x' y') -{-# INLINE insertWith' #-} - --- | /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" - -insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWithKey = go - where - 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 -{-# INLINEABLE insertWithKey #-} -#else -{-# INLINE insertWithKey #-} -#endif - --- | Same as 'insertWithKey', but the combining function is applied strictly. -insertWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWithKey' = go - where - STRICT_2_OF_4(go) - go _ kx x Tip = x `seq` 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 -{-# INLINEABLE 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")]) - -insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a - -> (Maybe a, Map k a) -insertLookupWithKey = go - where - 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 -{-# INLINEABLE insertLookupWithKey #-} -#else -{-# INLINE insertLookupWithKey #-} -#endif - --- | /O(log n)/. A strict version of 'insertLookupWithKey'. -insertLookupWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a - -> (Maybe a, Map k a) -insertLookupWithKey' = go - where - STRICT_2_OF_4(go) - go _ kx x Tip = x `seq` (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 -{-# INLINEABLE insertLookupWithKey' #-} -#else -{-# INLINE insertLookupWithKey' #-} -#endif - -{-------------------------------------------------------------------- - Deletion - [delete] is the inlined version of [deleteWith (\k x -> Nothing)] ---------------------------------------------------------------------} --- | /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 - -delete :: Ord k => k -> Map k a -> Map k a -delete = go - where - 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 -{-# INLINEABLE 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) -{-# INLINE adjust #-} - --- | /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')) -{-# INLINE adjustWithKey #-} - --- | /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) -{-# INLINE update #-} - --- | /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" - -updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a -updateWithKey = go - where - 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 -{-# INLINEABLE 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") - -updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) -updateLookupWithKey = go - where - 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 -{-# INLINEABLE 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")] - -alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a -alter = go - where - 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 -{-# INLINEABLE alter #-} -#else -{-# INLINE alter #-} -#endif - -{-------------------------------------------------------------------- - Indexing ---------------------------------------------------------------------} --- | /O(log n)/. Return the /index/ of a key. The index is a number from --- /0/ up to, but not including, the 'size' of the map. Calls 'error' when --- the key is not a 'member' of the map. --- --- > findIndex 2 (fromList [(5,"a"), (3,"b")]) Error: element is not in the map --- > findIndex 3 (fromList [(5,"a"), (3,"b")]) == 0 --- > findIndex 5 (fromList [(5,"a"), (3,"b")]) == 1 --- > findIndex 6 (fromList [(5,"a"), (3,"b")]) Error: element is not in the map - -findIndex :: Ord k => k -> Map k a -> Int -findIndex k t - = case lookupIndex k t of - Nothing -> error "Map.findIndex: element is not in the map" - Just idx -> idx -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE findIndex #-} -#endif - --- | /O(log n)/. Lookup the /index/ of a key. 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 - -lookupIndex :: Ord k => k -> Map k a -> Maybe Int -lookupIndex k = lkp k 0 - where - STRICT_1_OF_3(lkp) - STRICT_2_OF_3(lkp) - lkp _ _ Tip = Nothing - lkp key idx (Bin _ kx _ l r) - = case compare key kx of - LT -> lkp key idx l - GT -> lkp key (idx + size l + 1) r - EQ -> let idx' = idx + size l in idx' `seq` Just idx' -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE lookupIndex #-} -#endif - --- | /O(log n)/. Retrieve an element by /index/. Calls 'error' when an --- invalid index is used. --- --- > elemAt 0 (fromList [(5,"a"), (3,"b")]) == (3,"b") --- > elemAt 1 (fromList [(5,"a"), (3,"b")]) == (5, "a") --- > elemAt 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range - -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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE elemAt #-} -#endif - --- | /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' -> Bin sx kx x' l r - Nothing -> glue l r - where - sizeL = size l -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE updateAt #-} -#endif - --- | /O(log n)/. Delete the element at /index/. --- Defined as (@'deleteAt' i map = 'updateAt' (\k x -> 'Nothing') i map@). --- --- > deleteAt 0 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" --- > deleteAt 1 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" --- > deleteAt 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range --- > deleteAt (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range - -deleteAt :: Int -> Map k a -> Map k a -deleteAt i m - = updateAt (\_ _ -> Nothing) i m -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE deleteAt #-} -#endif - - -{-------------------------------------------------------------------- - 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" -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE findMin #-} -#endif - --- | /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" -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE findMax #-} -#endif - --- | /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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE deleteMin #-} -#endif - --- | /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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE deleteMax #-} -#endif - --- | /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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE updateMin #-} -#endif - --- | /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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE updateMax #-} -#endif - - --- | /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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE updateMinWithKey #-} -#endif - --- | /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) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE updateMaxWithKey #-} -#endif - --- | /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) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE minViewWithKey #-} -#endif - --- | /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) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE maxViewWithKey #-} -#endif - --- | /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) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE minView #-} -#endif - --- | /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 --- --- > 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) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE maxView #-} -#endif - --- 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. --- Hedge-union is more efficient on (bigset \``union`\` smallset). --- --- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")] - -union :: Ord k => Map k a -> Map k a -> Map k a -union Tip t2 = t2 -union t1 Tip = t1 -union (Bin _ k x Tip Tip) t = insert k x t -union t (Bin _ k x Tip Tip) = insertWith (\_ y->y) k x t -union t1 t2 = hedgeUnionL NothingS NothingS t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE union #-} -#endif - --- left-biased hedge union -hedgeUnionL :: Ord a - => MaybeS a -> MaybeS a -> Map a b -> Map a b - -> Map a b -hedgeUnionL _ _ t1 Tip - = t1 -hedgeUnionL blo bhi Tip (Bin _ kx x l r) - = join kx x (filterGt blo l) (filterLt bhi r) -hedgeUnionL blo bhi (Bin _ kx x l r) t2 - = join kx x (hedgeUnionL blo bmi l (trim blo bmi t2)) - (hedgeUnionL bmi bhi r (trim bmi bhi t2)) - where - bmi = JustS kx -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE hedgeUnionL #-} -#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 -{-# INLINE unionWith #-} - --- | /O(n+m)/. --- Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm. --- Hedge-union is more efficient on (bigset \``union`\` smallset). --- --- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value --- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] - -unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a -unionWithKey _ Tip t2 = t2 -unionWithKey _ t1 Tip = t1 -unionWithKey f t1 t2 = hedgeUnionWithKey f NothingS NothingS t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE unionWithKey #-} -#endif - -hedgeUnionWithKey :: Ord a - => (a -> b -> b -> b) - -> MaybeS a -> MaybeS a - -> Map a b -> Map a b - -> Map a b -hedgeUnionWithKey _ _ _ t1 Tip - = t1 -hedgeUnionWithKey _ blo bhi Tip (Bin _ kx x l r) - = join kx x (filterGt blo l) (filterLt bhi r) -hedgeUnionWithKey f blo bhi (Bin _ kx x l r) t2 - = join kx newx (hedgeUnionWithKey f blo bmi l lt) - (hedgeUnionWithKey f bmi bhi r gt) - where - bmi = JustS kx - lt = trim blo bmi t2 - (found,gt) = trimLookupLo kx bhi t2 - newx = case found of - Nothing -> x - Just (_,y) -> f kx x y -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE hedgeUnionWithKey #-} -#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 - = join 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 -{-# INLINE differenceWith #-} - --- | /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 _ Tip _ = Tip -differenceWithKey _ t1 Tip = t1 -differenceWithKey f t1 t2 = hedgeDiffWithKey f NothingS NothingS t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE differenceWithKey #-} -#endif - -hedgeDiffWithKey :: Ord a - => (a -> b -> c -> Maybe b) - -> MaybeS a -> MaybeS a - -> Map a b -> Map a c - -> Map a b -hedgeDiffWithKey _ _ _ Tip _ - = Tip -hedgeDiffWithKey _ blo bhi (Bin _ kx x l r) Tip - = join kx x (filterGt blo l) (filterLt bhi r) -hedgeDiffWithKey f blo bhi t (Bin _ kx x l r) - = case found of - Nothing -> merge tl tr - Just (ky,y) -> - case f ky y x of - Nothing -> merge tl tr - Just z -> join ky z tl tr - where - bmi = JustS kx - lt = trim blo bmi t - (found,gt) = trimLookupLo kx bhi t - tl = hedgeDiffWithKey f blo bmi lt l - tr = hedgeDiffWithKey f bmi bhi gt r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE hedgeDiffWithKey #-} -#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@). --- --- > 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 m1 m2 - = intersectionWithKey (\_ x _ -> x) m1 m2 -{-# INLINE intersection #-} - --- | /O(n+m)/. 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 m1 m2 - = intersectionWithKey (\_ x y -> f x y) m1 m2 -{-# INLINE intersectionWith #-} - --- | /O(n+m)/. Intersection with a combining function. --- Intersection is more efficient on (bigset \``intersection`\` smallset). --- --- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar --- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A" - - -intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c -intersectionWithKey _ Tip _ = Tip -intersectionWithKey _ _ Tip = Tip -intersectionWithKey f t1@(Bin s1 k1 x1 l1 r1) t2@(Bin s2 k2 x2 l2 r2) = - if s1 >= s2 then - let (lt,found,gt) = splitLookupWithKey k2 t1 - tl = intersectionWithKey f lt l2 - tr = intersectionWithKey f gt r2 - in case found of - Just (k,x) -> join k (f k x x2) tl tr - Nothing -> merge tl tr - else let (lt,found,gt) = splitLookup k1 t2 - tl = intersectionWithKey f l1 lt - tr = intersectionWithKey f r1 gt - in case found of - Just x -> join k1 (f k1 x1 x) tl tr - Nothing -> merge tl tr -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE intersectionWithKey #-} -#endif - - - -{-------------------------------------------------------------------- - 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 :: Ord k => (a -> Bool) -> Map k a -> Map k a -filter p m - = filterWithKey (\_ x -> p x) m -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE filter #-} -#endif - --- | /O(n)/. Filter all keys\/values that satisfy the predicate. --- --- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -filterWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> Map k a -filterWithKey _ Tip = Tip -filterWithKey p (Bin _ kx x l r) - | p kx x = join kx x (filterWithKey p l) (filterWithKey p r) - | otherwise = merge (filterWithKey p l) (filterWithKey p r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE filterWithKey #-} -#endif - --- | /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 :: Ord k => (a -> Bool) -> Map k a -> (Map k a,Map k a) -partition p m - = partitionWithKey (\_ x -> p x) m -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE partition #-} -#endif - --- | /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 :: Ord k => (k -> a -> Bool) -> Map k a -> (Map k a,Map k a) -partitionWithKey _ Tip = (Tip,Tip) -partitionWithKey p (Bin _ kx x l r) - | p kx x = (join kx x l1 r1,merge l2 r2) - | otherwise = (merge l1 r1,join kx x l2 r2) - where - (l1,l2) = partitionWithKey p l - (r1,r2) = partitionWithKey p r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE partitionWithKey #-} -#endif - --- | /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 :: Ord k => (a -> Maybe b) -> Map k a -> Map k b -mapMaybe f = mapMaybeWithKey (\_ x -> f x) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapMaybe #-} -#endif - --- | /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 :: Ord k => (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 -> join kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r) - Nothing -> merge (mapMaybeWithKey f l) (mapMaybeWithKey f r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapMaybeWithKey #-} -#endif - --- | /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 :: Ord k => (a -> Either b c) -> Map k a -> (Map k b, Map k c) -mapEither f m - = mapEitherWithKey (\_ x -> f x) m -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapEither #-} -#endif - --- | /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 :: Ord k => - (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) -mapEitherWithKey _ Tip = (Tip, Tip) -mapEitherWithKey f (Bin _ kx x l r) = case f kx x of - Left y -> (join kx y l1 r1, merge l2 r2) - Right z -> (merge l1 r1, join kx z l2 r2) - where - (l1,l2) = mapEitherWithKey f l - (r1,r2) = mapEitherWithKey f r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapEitherWithKey #-} -#endif - -{-------------------------------------------------------------------- - 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 = mapWithKey (\_ x -> f x) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE map #-} -#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) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapWithKey #-} -#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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapAccum #-} -#endif - --- | /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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapAccumWithKey #-} -#endif - --- | /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') -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapAccumL #-} -#endif - --- | /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') -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapAccumRWithKey #-} -#endif - --- | /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 smallest of --- these keys is retained. --- --- > mapKeys (+ 1) (fromList [(5,"a"), (3,"b")]) == fromList [(4, "b"), (6, "a")] --- > mapKeys (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "c" --- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c" - -mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a -mapKeys = mapKeysWith (\x _ -> x) -#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 . List.map fFirst . toList - where fFirst (x,y) = (f x, y) -#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) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapKeysMonotonic #-} -#endif - -{-------------------------------------------------------------------- - Folds ---------------------------------------------------------------------} - --- | /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. --- --- /Please note that fold will be deprecated in the future and removed./ -fold :: (a -> b -> b) -> b -> Map k a -> b -fold = foldr -{-# INLINE fold #-} - --- | /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 = go - 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 = go - 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 = go - 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 = go - 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. This function is an equivalent of 'foldrWithKey' and is present --- for compatibility only. --- --- /Please note that foldWithKey will be deprecated in the future and removed./ -foldWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b -foldWithKey = foldrWithKey -{-# INLINE foldWithKey #-} - --- | /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 = go - 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 = go - 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 = go - 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 = go - 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' #-} - -{-------------------------------------------------------------------- - List variations ---------------------------------------------------------------------} --- | /O(n)/. --- Return all elements of the map in the ascending order of their keys. --- --- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"] --- > elems empty == [] - -elems :: Map k a -> [a] -elems m - = [x | (_,x) <- assocs m] -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE elems #-} -#endif - --- | /O(n)/. Return all keys of the map in ascending order. --- --- > keys (fromList [(5,"a"), (3,"b")]) == [3,5] --- > keys empty == [] - -keys :: Map k a -> [k] -keys m - = [k | (k,_) <- assocs m] -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE keys #-} -#endif - --- | /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 m = Set.fromDistinctAscList (keys m) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE keysSet #-} -#endif - --- | /O(n)/. Return all key\/value pairs in the map in ascending key order. --- --- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] --- > assocs empty == [] - -assocs :: Map k a -> [(k,a)] -assocs m - = toList m -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE assocs #-} -#endif - -{-------------------------------------------------------------------- - 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. --- --- > 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")] - -fromList :: Ord k => [(k,a)] -> Map k a -fromList xs - = foldlStrict ins empty xs - where - ins t (k,x) = insert k x t -#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 to a list of key\/value pairs. --- --- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] --- > toList empty == [] - -toList :: Map k a -> [(k,a)] -toList t = toAscList t -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE toList #-} -#endif - --- | /O(n)/. Convert to an ascending list. --- --- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] - -toAscList :: Map k a -> [(k,a)] -toAscList t = foldrWithKey (\k x xs -> (k,x):xs) [] t -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE toAscList #-} -#endif - --- | /O(n)/. Convert to a descending list. -toDescList :: Map k a -> [(k,a)] -toDescList t = foldlWithKey (\xs k x -> (k,x):xs) [] t -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE 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 - -fromDistinctAscList :: [(k,a)] -> Map k a -fromDistinctAscList xs - = build const (length xs) xs - where - -- 1) use continuations so that we use heap space instead of stack space. - -- 2) special case for n==5 to build bushier trees. - build c 0 xs' = c Tip xs' - build c 5 xs' = case xs' of - ((k1,x1):(k2,x2):(k3,x3):(k4,x4):(k5,x5):xx) - -> c (bin k4 x4 (bin k2 x2 (singleton k1 x1) (singleton k3 x3)) (singleton k5 x5)) xx - _ -> error "fromDistinctAscList build" - build c n xs' = seq nr $ build (buildR nr c) nl xs' - where - nl = n `div` 2 - nr = n - nl - 1 - - buildR n c l ((k,x):ys) = build (buildB l k x c) n ys - buildR _ _ _ [] = error "fromDistinctAscList buildR []" - buildB l k x c r zs = c (bin k x l r) zs -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromDistinctAscList #-} -#endif - - -{-------------------------------------------------------------------- - 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 - -trimLookupLo :: Ord k => k -> MaybeS k -> Map k a -> (Maybe (k,a), Map k a) -trimLookupLo _ _ Tip = (Nothing, Tip) -trimLookupLo lo hi t@(Bin _ kx x l r) - = case compare lo kx of - LT -> case compare' kx hi of - LT -> (lookupAssoc lo t, t) - _ -> trimLookupLo lo hi l - GT -> trimLookupLo lo hi r - EQ -> (Just (kx,x),trim (JustS lo) hi r) - where compare' _ NothingS = LT - compare' kx' (JustS hi') = compare kx' hi' -#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 -> join 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 -> join 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 k t = k `seq` - case t of - Tip -> (Tip, Tip) - Bin _ kx x l r -> case compare k kx of - LT -> let (lt,gt) = split k l in (lt,join kx x gt r) - GT -> let (lt,gt) = split k r in (join 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 in (lt,z,join kx x gt r) - GT -> let (lt,z,gt) = splitLookup k r in (join kx x l lt,z,gt) - EQ -> (l,Just x,r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE splitLookup #-} -#endif - --- | /O(log n)/. -splitLookupWithKey :: Ord k => k -> Map k a -> (Map k a,Maybe (k,a),Map k a) -splitLookupWithKey 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) = splitLookupWithKey k l in (lt,z,join kx x gt r) - GT -> let (lt,z,gt) = splitLookupWithKey k r in (join kx x l lt,z,gt) - EQ -> (l,Just (kx, x),r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE splitLookupWithKey #-} -#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. - [join 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 [join], [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. ---------------------------------------------------------------------} - -{-------------------------------------------------------------------- - Join ---------------------------------------------------------------------} -join :: Ord k => k -> a -> Map k a -> Map k a -> Map k a -join kx x Tip r = insertMin kx x r -join kx x l Tip = insertMax kx x l -join kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz) - | delta*sizeL < sizeR = balanceL kz z (join kx x l lz) rz - | delta*sizeR < sizeL = balanceR ky y ly (join kx x ry r) - | otherwise = bin kx x l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE join #-} -#endif - - --- 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) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertMax #-} -#endif - -insertMin kx x t - = case t of - Tip -> singleton kx x - Bin _ ky y l r - -> balanceL ky y (insertMin kx x l) r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertMin #-} -#endif - -{-------------------------------------------------------------------- - [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 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE merge #-} -#endif - -{-------------------------------------------------------------------- - [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' -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE glue #-} -#endif - - --- | /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) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE deleteFindMin #-} -#endif - --- | /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) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE deleteFindMax #-} -#endif - - -{-------------------------------------------------------------------- - [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 _ Tip = pure Tip - traverse f (Bin s k v l r) - = flip (Bin s k) <$> traverse f l <*> f v <*> traverse f r - -instance Foldable.Foldable (Map k) where - fold Tip = mempty - fold (Bin _ _ v l r) = Foldable.fold l `mappend` v `mappend` Foldable.fold r - foldr = foldr - foldl = foldl - foldMap _ Tip = mempty - foldMap f (Bin _ _ v l r) = Foldable.foldMap f l `mappend` f v `mappend` Foldable.foldMap f r - -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 ---------------------------------------------------------------------} - -#include "Typeable.h" -INSTANCE_TYPEABLE2(Map,mapTc,"Map") - -{-------------------------------------------------------------------- - Assertions ---------------------------------------------------------------------} --- | /O(n)/. Test if the internal map structure is valid. --- --- > valid (fromAscList [(3,"b"), (5,"a")]) == True --- > valid (fromAscList [(5,"a"), (3,"b")]) == False - -valid :: Ord k => Map k a -> Bool -valid t - = balanced t && ordered t && validsize t - -ordered :: Ord a => Map a b -> Bool -ordered t - = bounded (const True) (const True) t - where - bounded lo hi t' - = case t' of - Tip -> True - Bin _ kx _ l r -> (lo kx) && (hi kx) && bounded lo (kx) hi r - --- | Exported only for "Debug.QuickCheck" -balanced :: Map k a -> Bool -balanced t - = case t of - Tip -> True - Bin _ _ _ l r -> (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) && - balanced l && balanced r - -validsize :: Map a b -> Bool -validsize t - = (realsize t == Just (size t)) - where - realsize t' - = case t' of - Tip -> Just 0 - Bin sz _ _ l r -> case (realsize l,realsize r) of - (Just n,Just m) | n+m+1 == sz -> Just sz - _ -> Nothing - -{-------------------------------------------------------------------- - Utilities ---------------------------------------------------------------------} -foldlStrict :: (a -> b -> a) -> a -> [b] -> a -foldlStrict f = go - where - go z [] = z - go z (x:xs) = let z' = f z x in z' `seq` go z' xs -{-# INLINE foldlStrict #-} +{-# LANGUAGE CPP #-} +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Safe #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.Map +-- 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). +-- +-- This module re-exports the value lazy 'Data.Map.Lazy' API, plus +-- several value strict functions from 'Data.Map.Strict'. +-- +-- These modules are intended to be imported qualified, to avoid name +-- clashes with Prelude functions, e.g. +-- +-- > 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 (). +----------------------------------------------------------------------------- + +module Data.Map + ( module Data.Map.Lazy + , insertWith' + , insertWithKey' + , insertLookupWithKey' + , fold + , foldWithKey + ) where + +import Data.Map.Lazy +import qualified Data.Map.Lazy as L +import qualified Data.Map.Strict as S + +-- | /Deprecated./ As of version 0.5, replaced by 'S.insertWith'. +-- +-- /O(log n)/. Same as 'insertWith', but the combining function is +-- applied strictly. This is often the most desirable behavior. +-- +-- For example, to update a counter: +-- +-- > insertWith' (+) k 1 m +-- +insertWith' :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWith' = S.insertWith +{-# INLINABLE insertWith' #-} + +-- | /Deprecated./ As of version 0.5, replaced by 'S.insertWithKey'. +-- +-- /O(log n)/. Same as 'insertWithKey', but the combining function is +-- applied strictly. +insertWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWithKey' = S.insertWithKey +{-# INLINABLE insertWithKey' #-} + +-- | /Deprecated./ As of version 0.5, replaced by +-- 'S.insertLookupWithKey'. +-- +-- /O(log n)/. A strict version of 'insertLookupWithKey'. +insertLookupWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a + -> (Maybe a, Map k a) +insertLookupWithKey' = S.insertLookupWithKey +{-# INLINABLE insertLookupWithKey' #-} + +-- | /Deprecated./ As of version 0.5, replaced by 'L.foldr'. +-- +-- /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. +fold :: (a -> b -> b) -> b -> Map k a -> b +fold = L.foldr +{-# INLINE fold #-} + +-- | /Deprecated./ As of version 0.4, replaced by 'L.foldrWithKey'. +-- +-- /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. +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 new file mode 100644 --- /dev/null +++ b/Data/Map/Base.hs @@ -0,0 +1,2722 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +----------------------------------------------------------------------------- +-- | +-- 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 IntMap, 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 + -- ** 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 + + -- * 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 + , join + , merge + , glue + , trim + , trimLookupLo + , foldlStrict + , MaybeS(..) + , filterGt + , filterLt + ) where + +import Prelude hiding (lookup,map,filter,foldr,foldl,null) +import qualified Data.Set.Base as Set +import Data.StrictPair +import Data.Monoid (Monoid(..)) +import Control.Applicative (Applicative(..), (<$>)) +import Data.Traversable (Traversable(traverse)) +import qualified Data.Foldable as Foldable +import Data.Typeable +import Control.DeepSeq (NFData(rnf)) + +#if __GLASGOW_HASKELL__ +import GHC.Exts ( build ) +import Text.Read +import Data.Data +#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_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 + +{-------------------------------------------------------------------- + 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 + +instance (Ord k) => Monoid (Map k v) where + mempty = empty + mappend = union + mconcat = unions + +#if __GLASGOW_HASKELL__ + +{-------------------------------------------------------------------- + A Data instance +--------------------------------------------------------------------} + +-- This instance preserves data abstraction at the cost of inefficiency. +-- We omit 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 _ = error "toConstr" + gunfold _ _ = error "gunfold" + dataTypeOf _ = mkNoRepType "Data.Map.Map" + dataCast2 f = gcast2 f + +#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. 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. 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 /index/. Calls 'error' when an +-- invalid index is used. +-- +-- > elemAt 0 (fromList [(5,"a"), (3,"b")]) == (3,"b") +-- > elemAt 1 (fromList [(5,"a"), (3,"b")]) == (5, "a") +-- > elemAt 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range + +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/. 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' -> Bin sx kx x' l r + Nothing -> glue l r + where + sizeL = size l + +-- | /O(log n)/. Delete the element at /index/. +-- Defined as (@'deleteAt' i map = 'updateAt' (\k x -> 'Nothing') i map@). +-- +-- > deleteAt 0 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +-- > deleteAt 1 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > deleteAt 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range +-- > deleteAt (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range + +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 +-- +-- > 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. +-- Hedge-union is more efficient on (bigset \``union`\` smallset). +-- +-- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")] + +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) = join 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 = join 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. +-- Hedge-union is more efficient on (bigset \``union`\` smallset). +-- +-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value +-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] + +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 = join 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@). +-- +-- > 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 join 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. +-- +-- > 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. +-- Intersection is more efficient on (bigset \``intersection`\` smallset). +-- +-- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar +-- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A" + + +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 $ join 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) -> join 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' -> join 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 = join 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 _ Tip = (Tip,Tip) +partitionWithKey p (Bin _ kx x l r) + | p kx x = (join kx x l1 r1,merge l2 r2) + | otherwise = (merge l1 r1,join kx x l2 r2) + where + (l1,l2) = partitionWithKey p l + (r1,r2) = partitionWithKey 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 -> join 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 _ Tip = (Tip, Tip) +mapEitherWithKey f (Bin _ kx x l r) = case f kx x of + Left y -> (join kx y l1 r1, merge l2 r2) + Right z -> (merge l1 r1, join kx z l2 r2) + where + (l1,l2) = mapEitherWithKey f l + (r1,r2) = mapEitherWithKey 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) + +-- | /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) + +-- | /O(n)/. +-- @'traverseWithKey' f s == '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 +{-# INLINE traverseWithKey #-} +traverseWithKey :: Applicative t => (k -> a -> t b) -> Map k a -> t (Map k b) +traverseWithKey f = go + where + go Tip = pure Tip + go (Bin s k v l r) + = flip (Bin s k) <$> go l <*> f k v <*> go r + +-- | /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' #-} + +{-------------------------------------------------------------------- + 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 +--------------------------------------------------------------------} +-- | /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. +-- +-- > 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")] + +fromList :: Ord k => [(k,a)] -> Map k a +fromList xs + = foldlStrict ins empty xs + where + ins t (k,x) = insert k x t +#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 + +fromDistinctAscList :: [(k,a)] -> Map k a +fromDistinctAscList xs + = create const (length xs) xs + where + -- 1) use continuations so that we use heap space instead of stack space. + -- 2) special case for n==5 to create bushier trees. + create c 0 xs' = c Tip xs' + create c 5 xs' = case xs' of + ((k1,x1):(k2,x2):(k3,x3):(k4,x4):(k5,x5):xx) + -> c (bin k4 x4 (bin k2 x2 (singleton k1 x1) (singleton k3 x3)) (singleton k5 x5)) xx + _ -> error "fromDistinctAscList create" + create c n xs' = seq nr $ create (createR nr c) nl xs' + where nl = n `div` 2 + nr = n - nl - 1 + + createR n c l ((k,x):ys) = create (createB l k x c) n ys + createR _ _ _ [] = error "fromDistinctAscList createR []" + createB l k x c r zs = c (bin k x 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 lk NothingS t = greater lk t + where greater :: Ord k => k -> Map k a -> (Maybe a, Map k a) + greater lo t'@(Bin _ kx x l r) = case compare lo kx of LT -> lookup lo l `strictPair` t' + EQ -> (Just x, r) + GT -> greater lo r + greater _ Tip = (Nothing, Tip) +trimLookupLo lk (JustS hk) t = middle lk hk t + where middle :: Ord k => k -> k -> Map k a -> (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 `strictPair` t' + | otherwise -> middle lo hi l + EQ -> Just x `strictPair` 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 -> join 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 -> join 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 k t = k `seq` + case t of + Tip -> (Tip, Tip) + Bin _ kx x l r -> case compare k kx of + LT -> let (lt,gt) = split k l in (lt,join kx x gt r) + GT -> let (lt,gt) = split k r in (join 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 in (lt,z,join kx x gt r) + GT -> let (lt,z,gt) = splitLookup k r in (join kx x l 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. + [join 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 [join], [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. +--------------------------------------------------------------------} + +{-------------------------------------------------------------------- + Join +--------------------------------------------------------------------} +join :: k -> a -> Map k a -> Map k a -> Map k a +join kx x Tip r = insertMin kx x r +join kx x l Tip = insertMax kx x l +join kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz) + | delta*sizeL < sizeR = balanceL kz z (join kx x l lz) rz + | delta*sizeR < sizeL = balanceR ky y ly (join kx x ry r) + | otherwise = bin kx x l r + + +-- insertMin and insertMax don't perform potentially expensive comparisons. +insertMax,insertMin :: 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) + +instance Foldable.Foldable (Map k) where + fold Tip = mempty + fold (Bin _ _ v l r) = Foldable.fold l `mappend` v `mappend` Foldable.fold r + foldr = foldr + foldl = foldl + foldMap _ Tip = mempty + foldMap f (Bin _ _ v l r) = Foldable.foldMap f l `mappend` f v `mappend` Foldable.foldMap f r + +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 +--------------------------------------------------------------------} + +#include "Typeable.h" +INSTANCE_TYPEABLE2(Map,mapTc,"Map") + +{-------------------------------------------------------------------- + Assertions +--------------------------------------------------------------------} +-- | /O(n)/. Test if the internal map structure is valid. +-- +-- > valid (fromAscList [(3,"b"), (5,"a")]) == True +-- > valid (fromAscList [(5,"a"), (3,"b")]) == False + +valid :: Ord k => Map k a -> Bool +valid t + = balanced t && ordered t && validsize t + +ordered :: Ord a => Map a b -> Bool +ordered t + = bounded (const True) (const True) t + where + bounded lo hi t' + = case t' of + Tip -> True + Bin _ kx _ l r -> (lo kx) && (hi kx) && bounded lo (kx) hi r + +-- | Exported only for "Debug.QuickCheck" +balanced :: Map k a -> Bool +balanced t + = case t of + Tip -> True + Bin _ _ _ l r -> (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) && + balanced l && balanced r + +validsize :: Map a b -> Bool +validsize t + = (realsize t == Just (size t)) + where + realsize t' + = case t' of + Tip -> Just 0 + Bin sz _ _ l r -> case (realsize l,realsize r) of + (Just n,Just m) | n+m+1 == sz -> Just sz + _ -> Nothing + +{-------------------------------------------------------------------- + Utilities +--------------------------------------------------------------------} +foldlStrict :: (a -> b -> a) -> a -> [b] -> a +foldlStrict f = go + where + go z [] = z + go z (x:xs) = let z' = f z x in z' `seq` go z' xs +{-# INLINE foldlStrict #-} diff --git a/Data/Map/Lazy.hs b/Data/Map/Lazy.hs new file mode 100644 --- /dev/null +++ b/Data/Map/Lazy.hs @@ -0,0 +1,227 @@ +{-# LANGUAGE CPP #-} +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Safe #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.Map.Lazy +-- 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 the keys, but lazy in the values. +-- If you need value-strict maps, use 'Data.Map.Strict' instead. +-- The 'Map' type itself 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.Lazy 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 (). +----------------------------------------------------------------------------- + +module Data.Map.Lazy ( + -- * Strictness properties + -- $strictness + + -- * Map type +#if !defined(TESTING) + Map -- instance Eq,Show,Read +#else + Map(..) -- instance Eq,Show,Read +#endif + + -- * Operators + , (!), (\\) + + -- * Query + , M.null + , size + , member + , notMember + , M.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 + , M.map + , mapWithKey + , traverseWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeys + , mapKeysWith + , mapKeysMonotonic + + -- * Folds + , M.foldr + , M.foldl + , foldrWithKey + , foldlWithKey + -- ** 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 + , M.filter + , filterWithKey + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + + -- * 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 + , join + , merge +#endif + + ) where + +import Data.Map.Base as M + +-- $strictness +-- +-- This module satisfies the following strictness property: +-- +-- * Key arguments are evaluated to WHNF +-- +-- Here are some examples that illustrate the property: +-- +-- > insertWith (\ new old -> old) undefined v m == undefined +-- > insertWith (\ new old -> old) k undefined m == OK +-- > delete undefined m == undefined diff --git a/Data/Map/Strict.hs b/Data/Map/Strict.hs new file mode 100644 --- /dev/null +++ b/Data/Map/Strict.hs @@ -0,0 +1,1139 @@ +{-# LANGUAGE CPP #-} +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Safe #-} +#endif +----------------------------------------------------------------------------- +-- | +-- 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'. +-- +-- 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 + -- ** 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 + + -- * 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 + , join + , 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.StrictPair + +-- Use macros to define strictness of functions. STRICT_x_OF_y +-- denotes an y-ary function strict in the x-th parameter. Similarly +-- STRICT_x_y_OF_z denotes an z-ary function strict in the x-th and +-- y-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_2_OF_3(fn) fn _ arg _ | arg `seq` False = undefined +#define STRICT_1_2_OF_3(fn) fn arg1 arg2 _ | arg1 `seq` arg2 `seq` False = undefined +#define STRICT_2_3_OF_4(fn) fn _ arg1 arg2 _ | arg1 `seq` arg2 `seq` False = undefined + +-- $strictness +-- +-- This module satisfies the following strictness properties: +-- +-- 1. Key and value arguments are evaluated to WHNF; +-- +-- 2. Keys and values are evaluated to WHNF before they are stored in +-- the map. +-- +-- Here are some examples that illustrate the first property: +-- +-- > insertWith (\ new old -> old) k undefined m == undefined +-- > 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 = def `seq` 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_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_3_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 = go + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a) + STRICT_2_3_OF_4(go) + go _ kx x Tip = Nothing `strictPair` 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 `strictPair` balanceL ky y l' r + GT -> let (found, r') = go f kx x r + in found `strictPair` balanceR ky y l r' + EQ -> let x' = f kx x y + in x' `seq` (Just y `strictPair` 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 = 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 `strictPair` balanceR kx x l' r + GT -> let (found,r') = go f k r + in found `strictPair` balanceL kx x l r' + EQ -> case f kx x of + Just x' -> x' `seq` (Just x' `strictPair` 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. +-- Hedge-union is more efficient on (bigset \``union`\` smallset). +-- +-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value +-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] + +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. +-- +-- > 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. +-- Intersection is more efficient on (bigset \``intersection`\` smallset). +-- +-- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar +-- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A" + + +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 $ join 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) -> join 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` join 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` join 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 _ Tip = (Tip, Tip) +mapEitherWithKey f (Bin _ kx x l r) = case f kx x of + Left y -> y `seq` (join kx y l1 r1 `strictPair` merge l2 r2) + Right z -> z `seq` (merge l1 r1 `strictPair` join kx z l2 r2) + where + (l1,l2) = mapEitherWithKey f l + (r1,r2) = mapEitherWithKey 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) + +-- | /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) + +-- | /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. +-- +-- > 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")] + +fromList :: Ord k => [(k,a)] -> Map k a +fromList xs + = foldlStrict ins empty xs + where + ins t (k,x) = insert k x t +#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 + +fromDistinctAscList :: [(k,a)] -> Map k a +fromDistinctAscList xs + = create const (length xs) xs + where + -- 1) use continuations so that we use heap space instead of stack space. + -- 2) special case for n==5 to create bushier trees. + create c 0 xs' = c Tip xs' + create c 5 xs' = case xs' of + ((k1,x1):(k2,x2):(k3,x3):(k4,x4):(k5,x5):xx) + -> x1 `seq` x2 `seq` x3 `seq` x4 `seq` x5 `seq` + c (bin k4 x4 (bin k2 x2 (singleton k1 x1) (singleton k3 x3)) + (singleton k5 x5)) xx + _ -> error "fromDistinctAscList create" + create c n xs' = seq nr $ create (createR nr c) nl xs' + where nl = n `div` 2 + nr = n - nl - 1 + + createR n c l ((k,x):ys) = x `seq` create (createB l k x c) n ys + createR _ _ _ [] = error "fromDistinctAscList createR []" + createB l k x c r zs = x `seq` c (bin k x l r) zs diff --git a/Data/Sequence.hs b/Data/Sequence.hs --- a/Data/Sequence.hs +++ b/Data/Sequence.hs @@ -1,3 +1,7 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif #if __GLASGOW_HASKELL__ >= 703 {-# LANGUAGE Trustworthy #-} #endif @@ -35,7 +39,11 @@ ----------------------------------------------------------------------------- module Data.Sequence ( +#if !defined(TESTING) Seq, +#else + Seq(..), Elem(..), FingerTree(..), Node(..), Digit(..), +#endif -- * Construction empty, -- :: Seq a singleton, -- :: a -> Seq a @@ -121,7 +129,10 @@ 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 - valid, + Sized(..), + deep, + node2, + node3, #endif ) where @@ -130,8 +141,9 @@ 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 (foldl', sortBy) +import qualified Data.List import Control.Applicative (Applicative(..), (<$>), WrappedMonad(..), liftA, liftA2, liftA3) +import Control.DeepSeq (NFData(rnf)) import Control.Monad (MonadPlus(..), ap) import Data.Monoid (Monoid(..)) import Data.Functor (Functor(..)) @@ -146,11 +158,6 @@ import Data.Data #endif -#if TESTING -import qualified Data.List (zipWith) -import Test.QuickCheck hiding ((><)) -#endif - infixr 5 `consTree` infixl 5 `snocTree` @@ -183,6 +190,9 @@ 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 = singleton xs >>= f = foldl' add empty xs @@ -307,9 +317,12 @@ 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 m `seq` rnf sf + {-# INLINE deep #-} -{-# SPECIALIZE INLINE deep :: Digit (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> FingerTree (Elem a) #-} -{-# SPECIALIZE INLINE deep :: Digit (Node a) -> FingerTree (Node (Node a)) -> Digit (Node a) -> FingerTree (Node a) #-} 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 @@ -383,6 +396,12 @@ 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 @@ -429,19 +448,19 @@ 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 #-} -{-# SPECIALIZE node2 :: Elem a -> Elem a -> Node (Elem a) #-} -{-# SPECIALIZE node2 :: Node a -> Node a -> Node (Node a) #-} node2 :: Sized a => a -> a -> Node a node2 a b = Node2 (size a + size b) a b {-# INLINE node3 #-} -{-# SPECIALIZE node3 :: Elem a -> Elem a -> Elem a -> Node (Elem a) #-} -{-# SPECIALIZE node3 :: Node a -> Node a -> Node a -> Node (Node a) #-} node3 :: Sized a => a -> a -> a -> Node a node3 a b c = Node3 (size a + size b + size c) a b c @@ -452,6 +471,9 @@ -- Elements newtype Elem a = Elem { getElem :: a } +#if TESTING + deriving Show +#endif instance Sized (Elem a) where size _ = 1 @@ -466,10 +488,8 @@ instance Traversable Elem where traverse f (Elem x) = Elem <$> f x -#ifdef TESTING -instance (Show a) => Show (Elem a) where - showsPrec p (Elem x) = showsPrec p x -#endif +instance NFData a => NFData (Elem a) where + rnf (Elem x) = rnf x ------------------------------------------------------- -- Applicative construction @@ -1771,85 +1791,3 @@ mergePQ cmp q1@(PQueue x1 ts1) q2@(PQueue x2 ts2) | cmp x1 x2 == GT = PQueue x2 (q1 :& ts2) | otherwise = PQueue x1 (q2 :& ts1) - -#if TESTING - ------------------------------------------------------------------------- --- QuickCheck ------------------------------------------------------------------------- - -instance Arbitrary a => Arbitrary (Seq a) where - arbitrary = Seq <$> arbitrary - shrink (Seq x) = map Seq (shrink x) - -instance Arbitrary a => Arbitrary (Elem a) where - arbitrary = Elem <$> arbitrary - -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 1 = Single <$> arbitrary - arb n = deep <$> arbitrary <*> arb (n `div` 2) <*> arbitrary - - shrink (Deep _ (One a) Empty (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 = [] - -instance (Arbitrary a, Sized a) => Arbitrary (Node a) where - arbitrary = oneof [ - node2 <$> arbitrary <*> arbitrary, - node3 <$> arbitrary <*> arbitrary <*> arbitrary] - - shrink (Node2 _ a b) = - [node2 a' b | a' <- shrink a] ++ - [node2 a b' | b' <- shrink b] - shrink (Node3 _ a b c) = - [node2 a b, node2 a c, node2 b c] ++ - [node3 a' b c | a' <- shrink a] ++ - [node3 a b' c | b' <- shrink b] ++ - [node3 a b c' | c' <- shrink c] - -instance Arbitrary a => Arbitrary (Digit a) where - arbitrary = oneof [ - One <$> arbitrary, - Two <$> arbitrary <*> arbitrary, - Three <$> arbitrary <*> arbitrary <*> arbitrary, - Four <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary] - - shrink (One a) = map One (shrink a) - shrink (Two a b) = [One a, One b] - shrink (Three a b c) = [Two a b, Two a c, Two b c] - shrink (Four a b c d) = [Three a b c, Three a b d, Three a c d, Three b c d] - ------------------------------------------------------------------------- --- Valid trees ------------------------------------------------------------------------- - -class Valid a where - valid :: a -> Bool - -instance Valid (Elem a) where - valid _ = True - -instance Valid (Seq a) where - valid (Seq xs) = valid xs - -instance (Sized a, Valid a) => Valid (FingerTree a) where - valid Empty = 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 - -instance (Sized a, Valid a) => Valid (Node a) where - valid node = size node == sum (fmap size node) && all valid node - -instance Valid a => Valid (Digit a) where - valid = all valid - -#endif diff --git a/Data/Set.hs b/Data/Set.hs --- a/Data/Set.hs +++ b/Data/Set.hs @@ -1,1271 +1,144 @@ -#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 -{-# LANGUAGE Safe #-} -#endif ------------------------------------------------------------------------------ --- | --- Module : Data.Set --- Copyright : (c) Daan Leijen 2002 --- License : BSD-style --- Maintainer : libraries@haskell.org --- Stability : provisional --- Portability : portable --- --- An efficient implementation of sets. --- --- Since many function names (but not the type name) clash with --- "Prelude" names, this module is usually imported @qualified@, e.g. --- --- > import Data.Set (Set) --- > import qualified Data.Set as Set --- --- The implementation of 'Set' 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'. Of course, left-biasing can only be observed --- when equality is an equivalence relation instead of structural --- equality. ------------------------------------------------------------------------------ - --- 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). --- --- 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 that can be INLINE are not recursive -- a 'go' function doing --- the real work is provided. - -module Data.Set ( - -- * Set type -#if !defined(TESTING) - Set -- instance Eq,Ord,Show,Read,Data,Typeable -#else - Set(..) -#endif - - -- * Operators - , (\\) - - -- * Query - , null - , size - , member - , notMember - , isSubsetOf - , isProperSubsetOf - - -- * Construction - , empty - , singleton - , insert - , delete - - -- * Combine - , union - , unions - , difference - , intersection - - -- * Filter - , filter - , partition - , split - , splitMember - - -- * Map - , map - , mapMonotonic - - -- * Folds - , foldr - , foldl - -- ** Strict folds - , foldr' - , foldl' - -- ** Legacy folds - , fold - - -- * Min\/Max - , findMin - , findMax - , deleteMin - , deleteMax - , deleteFindMin - , deleteFindMax - , maxView - , minView - - -- * Conversion - - -- ** List - , elems - , toList - , fromList - - -- ** Ordered list - , toAscList - , fromAscList - , fromDistinctAscList - - -- * Debugging - , showTree - , showTreeWith - , valid - -#if defined(TESTING) - -- Internals (for testing) - , bin - , balanced - , join - , merge -#endif - ) where - -import Prelude hiding (filter,foldl,foldr,null,map) -import qualified Data.List as List -import Data.Monoid (Monoid(..)) -import qualified Data.Foldable as Foldable -import Data.Typeable -import Control.DeepSeq (NFData(rnf)) - -{- --- just for testing -import QuickCheck -import List (nub,sort) -import qualified List --} - -#if __GLASGOW_HASKELL__ -import Text.Read -import Data.Data -#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 - -{-------------------------------------------------------------------- - Operators ---------------------------------------------------------------------} -infixl 9 \\ -- - --- | /O(n+m)/. See 'difference'. -(\\) :: Ord a => Set a -> Set a -> Set a -m1 \\ m2 = difference m1 m2 -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE (\\) #-} -#endif - -{-------------------------------------------------------------------- - Sets are size balanced trees ---------------------------------------------------------------------} --- | A set of values @a@. -data Set a = Tip - | Bin {-# UNPACK #-} !Size !a !(Set a) !(Set a) - -type Size = Int - -instance Ord a => Monoid (Set a) where - mempty = empty - mappend = union - mconcat = unions - -instance Foldable.Foldable Set where - fold Tip = mempty - fold (Bin _ k l r) = Foldable.fold l `mappend` k `mappend` Foldable.fold r - foldr = foldr - foldl = foldl - foldMap _ Tip = mempty - foldMap f (Bin _ k l r) = Foldable.foldMap f l `mappend` f k `mappend` Foldable.foldMap f r - -#if __GLASGOW_HASKELL__ - -{-------------------------------------------------------------------- - A Data instance ---------------------------------------------------------------------} - --- This instance preserves data abstraction at the cost of inefficiency. --- We omit reflection services for the sake of data abstraction. - -instance (Data a, Ord a) => Data (Set a) where - gfoldl f z set = z fromList `f` (toList set) - toConstr _ = error "toConstr" - gunfold _ _ = error "gunfold" - dataTypeOf _ = mkNoRepType "Data.Set.Set" - dataCast1 f = gcast1 f - -#endif - -{-------------------------------------------------------------------- - Query ---------------------------------------------------------------------} --- | /O(1)/. Is this the empty set? -null :: Set a -> Bool -null Tip = True -null (Bin {}) = False -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE null #-} -#endif - --- | /O(1)/. The number of elements in the set. -size :: Set a -> Int -size Tip = 0 -size (Bin sz _ _ _) = sz -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE size #-} -#endif - --- | /O(log n)/. Is the element in the set? -member :: Ord a => a -> Set a -> Bool -member = go - where - STRICT_1_OF_2(go) - 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 -{-# INLINABLE member #-} -#else -{-# INLINE member #-} -#endif - --- | /O(log n)/. Is the element not in the set? -notMember :: Ord a => a -> Set a -> Bool -notMember a t = not $ member a t -{-# INLINE notMember #-} - -{-------------------------------------------------------------------- - Construction ---------------------------------------------------------------------} --- | /O(1)/. The empty set. -empty :: Set a -empty = Tip - --- | /O(1)/. Create a singleton set. -singleton :: a -> Set a -singleton x = Bin 1 x Tip Tip - -{-------------------------------------------------------------------- - Insertion, Deletion ---------------------------------------------------------------------} --- | /O(log n)/. Insert an element in a set. --- If the set already contains an element equal to the given value, --- it is replaced with the new value. -insert :: Ord a => a -> Set a -> Set a -insert = go - where - 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 -{-# INLINEABLE insert #-} -#else -{-# INLINE insert #-} -#endif - --- Insert an element to the set only if it is not in the set. Used by --- `union`. -insertR :: Ord a => a -> Set a -> Set a -insertR = go - where - 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) - EQ -> t -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINEABLE insertR #-} -#else -{-# INLINE insertR #-} -#endif - --- | /O(log n)/. Delete an element from a set. -delete :: Ord a => a -> Set a -> Set a -delete = go - where - 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) - EQ -> glue l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINEABLE delete #-} -#else -{-# INLINE delete #-} -#endif - -{-------------------------------------------------------------------- - Subset ---------------------------------------------------------------------} --- | /O(n+m)/. Is this a proper subset? (ie. a subset but not equal). -isProperSubsetOf :: Ord a => Set a -> Set a -> Bool -isProperSubsetOf s1 s2 - = (size s1 < size s2) && (isSubsetOf s1 s2) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE isProperSubsetOf #-} -#endif - - --- | /O(n+m)/. Is this a subset? --- @(s1 `isSubsetOf` s2)@ tells whether @s1@ is a subset of @s2@. -isSubsetOf :: Ord a => Set a -> Set a -> Bool -isSubsetOf t1 t2 - = (size t1 <= size t2) && (isSubsetOfX t1 t2) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE isSubsetOf #-} -#endif - -isSubsetOfX :: Ord a => Set a -> Set a -> Bool -isSubsetOfX Tip _ = True -isSubsetOfX _ Tip = False -isSubsetOfX (Bin _ x l r) t - = found && isSubsetOfX l lt && isSubsetOfX r gt - where - (lt,found,gt) = splitMember x t -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE isSubsetOfX #-} -#endif - - -{-------------------------------------------------------------------- - Minimal, Maximal ---------------------------------------------------------------------} --- | /O(log n)/. The minimal element of a set. -findMin :: Set a -> a -findMin (Bin _ x Tip _) = x -findMin (Bin _ _ l _) = findMin l -findMin Tip = error "Set.findMin: empty set has no minimal element" -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE findMin #-} -#endif - --- | /O(log n)/. The maximal element of a set. -findMax :: Set a -> a -findMax (Bin _ x _ Tip) = x -findMax (Bin _ _ _ r) = findMax r -findMax Tip = error "Set.findMax: empty set has no maximal element" -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE findMax #-} -#endif - --- | /O(log n)/. Delete the minimal element. -deleteMin :: Set a -> Set a -deleteMin (Bin _ _ Tip r) = r -deleteMin (Bin _ x l r) = balanceR x (deleteMin l) r -deleteMin Tip = Tip -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE deleteMin #-} -#endif - --- | /O(log n)/. Delete the maximal element. -deleteMax :: Set a -> Set a -deleteMax (Bin _ _ l Tip) = l -deleteMax (Bin _ x l r) = balanceL x l (deleteMax r) -deleteMax Tip = Tip -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE deleteMax #-} -#endif - -{-------------------------------------------------------------------- - Union. ---------------------------------------------------------------------} --- | 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 -{-# INLINABLE unions #-} -#endif - --- | /O(n+m)/. The union of two sets, preferring the first set when --- equal elements are encountered. --- The implementation uses the efficient /hedge-union/ algorithm. --- Hedge-union is more efficient on (bigset `union` smallset). -union :: Ord a => Set a -> Set a -> Set a -union Tip t2 = t2 -union t1 Tip = t1 -union (Bin _ x Tip Tip) t = insert x t -union t (Bin _ x Tip Tip) = insertR x t -union t1 t2 = hedgeUnion NothingS NothingS t1 t2 -#if __GLASGOW_HASKELL__ >= 700 -{-# 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) - = join x (filterGt blo l) (filterLt bhi r) -hedgeUnion blo bhi (Bin _ x l r) t2 - = join 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/. -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 -{-# 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 - = join 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. --- Elements of the result come from the first set, so for example --- --- > import qualified Data.Set as S --- > data AB = A | B deriving Show --- > instance Ord AB where compare _ _ = EQ --- > instance Eq AB where _ == _ = True --- > main = print (S.singleton A `S.intersection` S.singleton B, --- > S.singleton B `S.intersection` S.singleton A) --- --- prints @(fromList [A],fromList [B])@. -intersection :: Ord a => Set a -> Set a -> Set a -intersection Tip _ = Tip -intersection _ Tip = Tip -intersection t1@(Bin s1 x1 l1 r1) t2@(Bin s2 x2 l2 r2) = - if s1 >= s2 then - let (lt,found,gt) = splitLookup x2 t1 - tl = intersection lt l2 - tr = intersection gt r2 - in case found of - Just x -> join x tl tr - Nothing -> merge tl tr - else let (lt,found,gt) = splitMember x1 t2 - tl = intersection l1 lt - tr = intersection r1 gt - in if found then join x1 tl tr - else merge tl tr -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE intersection #-} -#endif - -{-------------------------------------------------------------------- - Filter and partition ---------------------------------------------------------------------} --- | /O(n)/. Filter all elements that satisfy the predicate. -filter :: Ord a => (a -> Bool) -> Set a -> Set a -filter _ Tip = Tip -filter p (Bin _ x l r) - | p x = join x (filter p l) (filter p r) - | otherwise = merge (filter p l) (filter p r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE filter #-} -#endif - --- | /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. --- See also 'split'. -partition :: Ord a => (a -> Bool) -> Set a -> (Set a,Set a) -partition _ Tip = (Tip, Tip) -partition p (Bin _ x l r) = case (partition p l, partition p r) of - ((l1, l2), (r1, r2)) - | p x -> (join x l1 r1, merge l2 r2) - | otherwise -> (merge l1 r1, join x l2 r2) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE partition #-} -#endif - -{---------------------------------------------------------------------- - Map -----------------------------------------------------------------------} - --- | /O(n*log n)/. --- @'map' f s@ is the set obtained by applying @f@ to each element of @s@. --- --- It's worth noting that the size of the result may be smaller if, --- for some @(x,y)@, @x \/= y && f x == f y@ - -map :: (Ord a, Ord b) => (a->b) -> Set a -> Set b -map f = fromList . List.map f . toList -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE map #-} -#endif - --- | /O(n)/. The --- --- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is monotonic. --- /The precondition is not checked./ --- Semi-formally, we have: --- --- > and [x < y ==> f x < f y | x <- ls, y <- ls] --- > ==> mapMonotonic f s == map f s --- > where ls = toList s - -mapMonotonic :: (a->b) -> Set a -> Set b -mapMonotonic _ Tip = Tip -mapMonotonic f (Bin sz x l r) = Bin sz (f x) (mapMonotonic f l) (mapMonotonic f r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE mapMonotonic #-} -#endif - -{-------------------------------------------------------------------- - Fold ---------------------------------------------------------------------} --- | /O(n)/. Fold the elements in the set using the given right-associative --- binary operator. This function is an equivalent of 'foldr' and is present --- for compatibility only. --- --- /Please note that fold will be deprecated in the future and removed./ -fold :: (a -> b -> b) -> b -> Set a -> b -fold = foldr -{-# INLINE fold #-} - --- | /O(n)/. Fold the elements in the set using the given right-associative --- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'toAscList'@. --- --- For example, --- --- > toAscList set = foldr (:) [] set -foldr :: (a -> b -> b) -> b -> Set a -> b -foldr f = go - 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 -> Set a -> b -foldr' f = go - 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 elements in the set using the given left-associative --- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'toAscList'@. --- --- For example, --- --- > toDescList set = foldl (flip (:)) [] set -foldl :: (a -> b -> a) -> a -> Set b -> a -foldl f = go - 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 -> Set b -> a -foldl' f = go - 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' #-} - -{-------------------------------------------------------------------- - List variations ---------------------------------------------------------------------} --- | /O(n)/. The elements of a set. -elems :: Set a -> [a] -elems = toList -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE elems #-} -#endif - -{-------------------------------------------------------------------- - Lists ---------------------------------------------------------------------} --- | /O(n)/. Convert the set to a list of elements. -toList :: Set a -> [a] -toList = toAscList -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE toList #-} -#endif - --- | /O(n)/. Convert the set to an ascending list of elements. -toAscList :: Set a -> [a] -toAscList = foldr (:) [] -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE toAscList #-} -#endif - --- | /O(n*log n)/. Create a set from a list of elements. -fromList :: Ord a => [a] -> Set a -fromList = foldlStrict ins empty - where - ins t x = insert x t -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromList #-} -#endif - -{-------------------------------------------------------------------- - Building trees from ascending/descending lists can be done in linear time. - - Note that if [xs] is ascending that: - fromAscList xs == fromList xs ---------------------------------------------------------------------} --- | /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 -{-# INLINABLE fromAscList #-} -#endif - - --- | /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./ -fromDistinctAscList :: [a] -> Set a -fromDistinctAscList xs - = build const (length xs) xs - where - -- 1) use continutations so that we use heap space instead of stack space. - -- 2) special case for n==5 to build bushier trees. - build c 0 xs' = c Tip xs' - build c 5 xs' = case xs' of - (x1:x2:x3:x4:x5:xx) - -> c (bin x4 (bin x2 (singleton x1) (singleton x3)) (singleton x5)) xx - _ -> error "fromDistinctAscList build 5" - build c n xs' = seq nr $ build (buildR nr c) nl xs' - where - nl = n `div` 2 - nr = n - nl - 1 - - buildR n c l (x:ys) = build (buildB l x c) n ys - buildR _ _ _ [] = error "fromDistinctAscList buildR []" - buildB l x c r zs = c (bin x l r) zs -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE fromDistinctAscList #-} -#endif - -{-------------------------------------------------------------------- - Eq converts the set 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 a => Eq (Set a) where - t1 == t2 = (size t1 == size t2) && (toAscList t1 == toAscList t2) - -{-------------------------------------------------------------------- - Ord ---------------------------------------------------------------------} - -instance Ord a => Ord (Set a) where - compare s1 s2 = compare (toAscList s1) (toAscList s2) - -{-------------------------------------------------------------------- - Show ---------------------------------------------------------------------} -instance Show a => Show (Set a) where - showsPrec p xs = showParen (p > 10) $ - showString "fromList " . shows (toList xs) - -{-------------------------------------------------------------------- - Read ---------------------------------------------------------------------} -instance (Read a, Ord a) => Read (Set 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 - -{-------------------------------------------------------------------- - Typeable/Data ---------------------------------------------------------------------} - -#include "Typeable.h" -INSTANCE_TYPEABLE1(Set,setTc,"Set") - -{-------------------------------------------------------------------- - NFData ---------------------------------------------------------------------} - -instance NFData a => NFData (Set a) where - rnf Tip = () - 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 -> join 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 -> join 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 _ Tip = (Tip,Tip) -split x (Bin _ y l r) - = case compare x y of - LT -> let (lt,gt) = split x l in (lt,join y gt r) - GT -> let (lt,gt) = split x r in (join y l lt,gt) - EQ -> (l,r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE split #-} -#endif - --- | /O(log n)/. Performs a 'split' but also returns whether the pivot --- element was found in the original set. -splitMember :: Ord a => a -> Set a -> (Set a,Bool,Set a) -splitMember x t = let (l,m,r) = splitLookup x t in - (l,maybe False (const True) m,r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE splitMember #-} -#endif - --- | /O(log n)/. Performs a 'split' but also returns the pivot --- element that was found in the original set. -splitLookup :: Ord a => a -> Set a -> (Set a,Maybe a,Set a) -splitLookup _ Tip = (Tip,Nothing,Tip) -splitLookup x (Bin _ y l r) - = case compare x y of - LT -> let (lt,found,gt) = splitLookup x l in (lt,found,join y gt r) - GT -> let (lt,found,gt) = splitLookup x r in (join y l lt,found,gt) - EQ -> (l,Just y,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] < [x] and all values - in [r] > [x], and that [l] and [r] are valid trees. - - In order of sophistication: - [Bin sz x l r] The type constructor. - [bin x l r] Maintains the correct size, assumes that both [l] - and [r] are balanced with respect to each other. - [balance 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. - [join 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 [join], [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. ---------------------------------------------------------------------} - -{-------------------------------------------------------------------- - Join ---------------------------------------------------------------------} -join :: a -> Set a -> Set a -> Set a -join x Tip r = insertMin x r -join x l Tip = insertMax x l -join x l@(Bin sizeL y ly ry) r@(Bin sizeR z lz rz) - | delta*sizeL < sizeR = balanceL z (join x l lz) rz - | delta*sizeR < sizeL = balanceR y ly (join x ry r) - | otherwise = bin x l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE join #-} -#endif - - --- insertMin and insertMax don't perform potentially expensive comparisons. -insertMax,insertMin :: a -> Set a -> Set a -insertMax x t - = case t of - Tip -> singleton x - Bin _ y l r - -> balanceR y l (insertMax x r) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertMax #-} -#endif - -insertMin x t - = case t of - Tip -> singleton x - Bin _ y l r - -> balanceL y (insertMin x l) r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE insertMin #-} -#endif - -{-------------------------------------------------------------------- - [merge l r]: merges two trees. ---------------------------------------------------------------------} -merge :: Set a -> Set a -> Set a -merge Tip r = r -merge l Tip = l -merge l@(Bin sizeL x lx rx) r@(Bin sizeR y ly ry) - | delta*sizeL < sizeR = balanceL y (merge l ly) ry - | delta*sizeR < sizeL = balanceR x lx (merge rx r) - | otherwise = glue l r -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE merge #-} -#endif - -{-------------------------------------------------------------------- - [glue l r]: glues two trees together. - Assumes that [l] and [r] are already balanced with respect to each other. ---------------------------------------------------------------------} -glue :: Set a -> Set a -> Set a -glue Tip r = r -glue l Tip = l -glue l r - | size l > size r = let (m,l') = deleteFindMax l in balanceR m l' r - | otherwise = let (m,r') = deleteFindMin r in balanceL m l r' -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE glue #-} -#endif - - --- | /O(log n)/. Delete and find the minimal element. --- --- > deleteFindMin set = (findMin set, deleteMin set) - -deleteFindMin :: Set a -> (a,Set a) -deleteFindMin t - = case t of - Bin _ x Tip r -> (x,r) - Bin _ x l r -> let (xm,l') = deleteFindMin l in (xm,balanceR x l' r) - Tip -> (error "Set.deleteFindMin: can not return the minimal element of an empty set", Tip) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE deleteFindMin #-} -#endif - --- | /O(log n)/. Delete and find the maximal element. --- --- > deleteFindMax set = (findMax set, deleteMax set) -deleteFindMax :: Set a -> (a,Set a) -deleteFindMax t - = case t of - Bin _ x l Tip -> (x,l) - Bin _ x l r -> let (xm,r') = deleteFindMax r in (xm,balanceL x l r') - Tip -> (error "Set.deleteFindMax: can not return the maximal element of an empty set", Tip) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE deleteFindMax #-} -#endif - --- | /O(log n)/. Retrieves the minimal key of the set, and the set --- stripped of that element, or 'Nothing' if passed an empty set. -minView :: Set a -> Maybe (a, Set a) -minView Tip = Nothing -minView x = Just (deleteFindMin x) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE minView #-} -#endif - --- | /O(log n)/. Retrieves the maximal key of the set, and the set --- stripped of that element, or 'Nothing' if passed an empty set. -maxView :: Set a -> Maybe (a, Set a) -maxView Tip = Nothing -maxView x = Just (deleteFindMax x) -#if __GLASGOW_HASKELL__ >= 700 -{-# INLINABLE maxView #-} -#endif - -{-------------------------------------------------------------------- - [balance x l 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 correspondes 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 errorneous: - - 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 :: a -> Set a -> Set a -> Set a --- balance x l r --- | sizeL + sizeR <= 1 = Bin sizeX x l r --- | sizeR > delta*sizeL = rotateL x l r --- | sizeL > delta*sizeR = rotateR x l r --- | otherwise = Bin sizeX x l r --- where --- sizeL = size l --- sizeR = size r --- sizeX = sizeL + sizeR + 1 --- --- rotateL :: a -> Set a -> Set a -> Set a --- rotateL x l r@(Bin _ _ ly ry) | size ly < ratio*size ry = singleL x l r --- | otherwise = doubleL x l r --- rotateR :: a -> Set a -> Set a -> Set a --- rotateR x l@(Bin _ _ ly ry) r | size ry < ratio*size ly = singleR x l r --- | otherwise = doubleR x l r --- --- singleL, singleR :: a -> Set a -> Set a -> Set a --- singleL x1 t1 (Bin _ x2 t2 t3) = bin x2 (bin x1 t1 t2) t3 --- singleR x1 (Bin _ x2 t1 t2) t3 = bin x2 t1 (bin x1 t2 t3) --- --- doubleL, doubleR :: a -> Set a -> Set a -> Set a --- doubleL x1 t1 (Bin _ x2 (Bin _ x3 t2 t3) t4) = bin x3 (bin x1 t1 t2) (bin x2 t3 t4) --- doubleR x1 (Bin _ x2 t1 (Bin _ x3 t2 t3)) t4 = bin x3 (bin x2 t1 t2) (bin x1 t3 t4) --- --- It is only written in such a way that every node is pattern-matched only once. --- --- Only balanceL and balanceR are needed at the moment, so balance is not here anymore. --- In case it is needed, it can be found in Data.Map. - --- 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 :: a -> Set a -> Set a -> Set a -balanceL x l r = case r of - Tip -> case l of - Tip -> Bin 1 x Tip Tip - (Bin _ _ Tip Tip) -> Bin 2 x l Tip - (Bin _ lx Tip (Bin _ lrx _ _)) -> Bin 3 lrx (Bin 1 lx Tip Tip) (Bin 1 x Tip Tip) - (Bin _ lx ll@(Bin _ _ _ _) Tip) -> Bin 3 lx ll (Bin 1 x Tip Tip) - (Bin ls lx ll@(Bin lls _ _ _) lr@(Bin lrs lrx lrl lrr)) - | lrs < ratio*lls -> Bin (1+ls) lx ll (Bin (1+lrs) x lr Tip) - | otherwise -> Bin (1+ls) lrx (Bin (1+lls+size lrl) lx ll lrl) (Bin (1+size lrr) x lrr Tip) - - (Bin rs _ _ _) -> case l of - Tip -> Bin (1+rs) x Tip r - - (Bin ls lx ll lr) - | ls > delta*rs -> case (ll, lr) of - (Bin lls _ _ _, Bin lrs lrx lrl lrr) - | lrs < ratio*lls -> Bin (1+ls+rs) lx ll (Bin (1+rs+lrs) x lr r) - | otherwise -> Bin (1+ls+rs) lrx (Bin (1+lls+size lrl) lx ll lrl) (Bin (1+rs+size lrr) x lrr r) - (_, _) -> error "Failure in Data.Map.balanceL" - | otherwise -> Bin (1+ls+rs) 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 :: a -> Set a -> Set a -> Set a -balanceR x l r = case l of - Tip -> case r of - Tip -> Bin 1 x Tip Tip - (Bin _ _ Tip Tip) -> Bin 2 x Tip r - (Bin _ rx Tip rr@(Bin _ _ _ _)) -> Bin 3 rx (Bin 1 x Tip Tip) rr - (Bin _ rx (Bin _ rlx _ _) Tip) -> Bin 3 rlx (Bin 1 x Tip Tip) (Bin 1 rx Tip Tip) - (Bin rs rx rl@(Bin rls rlx rll rlr) rr@(Bin rrs _ _ _)) - | rls < ratio*rrs -> Bin (1+rs) rx (Bin (1+rls) x Tip rl) rr - | otherwise -> Bin (1+rs) rlx (Bin (1+size rll) x Tip rll) (Bin (1+rrs+size rlr) rx rlr rr) - - (Bin ls _ _ _) -> case r of - Tip -> Bin (1+ls) x l Tip - - (Bin rs rx rl rr) - | rs > delta*ls -> case (rl, rr) of - (Bin rls rlx rll rlr, Bin rrs _ _ _) - | rls < ratio*rrs -> Bin (1+ls+rs) rx (Bin (1+ls+rls) x l rl) rr - | otherwise -> Bin (1+ls+rs) rlx (Bin (1+ls+size rll) x l rll) (Bin (1+rrs+size rlr) rx rlr rr) - (_, _) -> error "Failure in Data.Map.balanceR" - | otherwise -> Bin (1+ls+rs) x l r -{-# NOINLINE balanceR #-} - -{-------------------------------------------------------------------- - The bin constructor maintains the size of the tree ---------------------------------------------------------------------} -bin :: a -> Set a -> Set a -> Set a -bin x l r - = Bin (size l + size r + 1) x l r -{-# INLINE bin #-} - - -{-------------------------------------------------------------------- - Utilities ---------------------------------------------------------------------} -foldlStrict :: (a -> b -> a) -> a -> [b] -> a -foldlStrict f = go - where - go z [] = z - go z (x:xs) = let z' = f z x in z' `seq` go z' xs -{-# INLINE foldlStrict #-} - -{-------------------------------------------------------------------- - Debugging ---------------------------------------------------------------------} --- | /O(n)/. Show the tree that implements the set. The tree is shown --- in a compressed, hanging format. -showTree :: Show a => Set a -> String -showTree s - = showTreeWith True False s - - -{- | /O(n)/. The expression (@showTreeWith hang wide map@) shows - the tree that implements the set. 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. - -> Set> putStrLn $ showTreeWith True False $ fromDistinctAscList [1..5] -> 4 -> +--2 -> | +--1 -> | +--3 -> +--5 -> -> Set> putStrLn $ showTreeWith True True $ fromDistinctAscList [1..5] -> 4 -> | -> +--2 -> | | -> | +--1 -> | | -> | +--3 -> | -> +--5 -> -> Set> putStrLn $ showTreeWith False True $ fromDistinctAscList [1..5] -> +--5 -> | -> 4 -> | -> | +--3 -> | | -> +--2 -> | -> +--1 - --} -showTreeWith :: Show a => Bool -> Bool -> Set a -> String -showTreeWith hang wide t - | hang = (showsTreeHang wide [] t) "" - | otherwise = (showsTree wide [] [] t) "" - -showsTree :: Show a => Bool -> [String] -> [String] -> Set a -> ShowS -showsTree wide lbars rbars t - = case t of - Tip -> showsBars lbars . showString "|\n" - Bin _ x Tip Tip - -> showsBars lbars . shows x . showString "\n" - Bin _ x l r - -> showsTree wide (withBar rbars) (withEmpty rbars) r . - showWide wide rbars . - showsBars lbars . shows x . showString "\n" . - showWide wide lbars . - showsTree wide (withEmpty lbars) (withBar lbars) l - -showsTreeHang :: Show a => Bool -> [String] -> Set a -> ShowS -showsTreeHang wide bars t - = case t of - Tip -> showsBars bars . showString "|\n" - Bin _ x Tip Tip - -> showsBars bars . shows x . showString "\n" - Bin _ x l r - -> showsBars bars . shows x . showString "\n" . - showWide wide bars . - showsTreeHang wide (withBar bars) l . - showWide wide bars . - showsTreeHang 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 - -{-------------------------------------------------------------------- - Assertions ---------------------------------------------------------------------} --- | /O(n)/. Test if the internal set structure is valid. -valid :: Ord a => Set a -> Bool -valid t - = balanced t && ordered t && validsize t - -ordered :: Ord a => Set a -> Bool -ordered t - = bounded (const True) (const True) t - where - bounded lo hi t' - = case t' of - Tip -> True - Bin _ x l r -> (lo x) && (hi x) && bounded lo (x) hi r - -balanced :: Set a -> Bool -balanced t - = case t of - Tip -> True - Bin _ _ l r -> (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) && - balanced l && balanced r - -validsize :: Set a -> Bool -validsize t - = (realsize t == Just (size t)) - where - realsize t' - = case t' of - Tip -> Just 0 - Bin sz _ l r -> case (realsize l,realsize r) of - (Just n,Just m) | n+m+1 == sz -> Just sz - _ -> Nothing +{-# LANGUAGE CPP #-} +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Safe #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.Set +-- Copyright : (c) Daan Leijen 2002 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- An efficient implementation of sets. +-- +-- These modules are intended to be imported qualified, to avoid name +-- clashes with Prelude functions, e.g. +-- +-- > import Data.Set (Set) +-- > import qualified Data.Set as Set +-- +-- The implementation of 'Set' 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'. Of course, left-biasing can only be observed +-- when equality is an equivalence relation instead of structural +-- equality. +----------------------------------------------------------------------------- + +module Data.Set ( + -- * Strictness properties + -- $strictness + + -- * Set type +#if !defined(TESTING) + Set -- instance Eq,Ord,Show,Read,Data,Typeable +#else + Set(..) +#endif + + -- * Operators + , (\\) + + -- * Query + , S.null + , size + , member + , notMember + , lookupLT + , lookupGT + , lookupLE + , lookupGE + , isSubsetOf + , isProperSubsetOf + + -- * Construction + , empty + , singleton + , insert + , delete + + -- * Combine + , union + , unions + , difference + , intersection + + -- * Filter + , S.filter + , partition + , split + , splitMember + + -- * Map + , S.map + , mapMonotonic + + -- * Folds + , S.foldr + , S.foldl + -- ** Strict folds + , foldr' + , foldl' + -- ** Legacy folds + , fold + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , maxView + , minView + + -- * Conversion + + -- ** List + , elems + , toList + , fromList + + -- ** Ordered list + , toAscList + , toDescList + , fromAscList + , fromDistinctAscList + + -- * Debugging + , showTree + , showTreeWith + , valid + +#if defined(TESTING) + -- Internals (for testing) + , bin + , balanced + , join + , merge +#endif + ) where + +import Data.Set.Base as S + +-- $strictness +-- +-- This module satisfies the following strictness property: +-- +-- * Key arguments are evaluated to WHNF +-- +-- Here are some examples that illustrate the property: +-- +-- > delete undefined s == undefined diff --git a/Data/Set/Base.hs b/Data/Set/Base.hs new file mode 100644 --- /dev/null +++ b/Data/Set/Base.hs @@ -0,0 +1,1364 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.Set.Base +-- Copyright : (c) Daan Leijen 2002 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- An efficient implementation of sets. +-- +-- These modules are intended to be imported qualified, to avoid name +-- clashes with Prelude functions, e.g. +-- +-- > import Data.Set (Set) +-- > import qualified Data.Set as Set +-- +-- The implementation of 'Set' 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'. Of course, left-biasing can only be observed +-- when equality is an equivalence relation instead of structural +-- equality. +----------------------------------------------------------------------------- + +-- [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 IntSet, 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 x and then look at the resulting code for hedgeInt. + + +-- [Note: Order of constructors] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- The order of constructors of Set 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.Set.Base ( + -- * Set type + Set(..) -- instance Eq,Ord,Show,Read,Data,Typeable + + -- * Operators + , (\\) + + -- * Query + , null + , size + , member + , notMember + , lookupLT + , lookupGT + , lookupLE + , lookupGE + , isSubsetOf + , isProperSubsetOf + + -- * Construction + , empty + , singleton + , insert + , delete + + -- * Combine + , union + , unions + , difference + , intersection + + -- * Filter + , filter + , partition + , split + , splitMember + + -- * Map + , map + , mapMonotonic + + -- * Folds + , foldr + , foldl + -- ** Strict folds + , foldr' + , foldl' + -- ** Legacy folds + , fold + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , maxView + , minView + + -- * Conversion + + -- ** List + , elems + , toList + , fromList + + -- ** Ordered list + , toAscList + , toDescList + , fromAscList + , fromDistinctAscList + + -- * Debugging + , showTree + , showTreeWith + , valid + + -- Internals (for testing) + , bin + , balanced + , join + , merge + ) where + +import Prelude hiding (filter,foldl,foldr,null,map) +import qualified Data.List as List +import Data.Monoid (Monoid(..)) +import qualified Data.Foldable as Foldable +import Data.Typeable +import Control.DeepSeq (NFData(rnf)) + +#if __GLASGOW_HASKELL__ +import GHC.Exts ( build ) +import Text.Read +import Data.Data +#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_1_OF_3(fn) fn arg _ _ | arg `seq` False = undefined + +{-------------------------------------------------------------------- + Operators +--------------------------------------------------------------------} +infixl 9 \\ -- + +-- | /O(n+m)/. See 'difference'. +(\\) :: Ord a => Set a -> Set a -> Set a +m1 \\ m2 = difference m1 m2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE (\\) #-} +#endif + +{-------------------------------------------------------------------- + Sets are size balanced trees +--------------------------------------------------------------------} +-- | A set of values @a@. + +-- See Note: Order of constructors +data Set a = Bin {-# UNPACK #-} !Size !a !(Set a) !(Set a) + | Tip + +type Size = Int + +instance Ord a => Monoid (Set a) where + mempty = empty + mappend = union + mconcat = unions + +instance Foldable.Foldable Set where + fold Tip = mempty + fold (Bin _ k l r) = Foldable.fold l `mappend` k `mappend` Foldable.fold r + foldr = foldr + foldl = foldl + foldMap _ Tip = mempty + foldMap f (Bin _ k l r) = Foldable.foldMap f l `mappend` f k `mappend` Foldable.foldMap f r + +#if __GLASGOW_HASKELL__ + +{-------------------------------------------------------------------- + A Data instance +--------------------------------------------------------------------} + +-- This instance preserves data abstraction at the cost of inefficiency. +-- We omit reflection services for the sake of data abstraction. + +instance (Data a, Ord a) => Data (Set a) where + gfoldl f z set = z fromList `f` (toList set) + toConstr _ = error "toConstr" + gunfold _ _ = error "gunfold" + dataTypeOf _ = mkNoRepType "Data.Set.Set" + dataCast1 f = gcast1 f + +#endif + +{-------------------------------------------------------------------- + Query +--------------------------------------------------------------------} +-- | /O(1)/. Is this the empty set? +null :: Set a -> Bool +null Tip = True +null (Bin {}) = False +{-# INLINE null #-} + +-- | /O(1)/. The number of elements in the set. +size :: Set a -> Int +size Tip = 0 +size (Bin sz _ _ _) = sz +{-# INLINE size #-} + +-- | /O(log n)/. Is the element in the set? +member :: Ord a => a -> Set a -> Bool +member = go + where + STRICT_1_OF_2(go) + 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 +{-# INLINABLE member #-} +#else +{-# INLINE member #-} +#endif + +-- | /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 +{-# INLINABLE notMember #-} +#else +{-# INLINE notMember #-} +#endif + +-- | /O(log n)/. Find largest element smaller than the given one. +-- +-- > lookupLT 3 (fromList [3, 5]) == Nothing +-- > lookupLT 5 (fromList [3, 5]) == Just 3 +lookupLT :: Ord a => a -> Set a -> Maybe a +lookupLT = goNothing + where + STRICT_1_OF_2(goNothing) + 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 x best (Bin _ y l r) | x <= y = goJust x best l + | otherwise = goJust x y r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE lookupLT #-} +#else +{-# INLINE lookupLT #-} +#endif + +-- | /O(log n)/. Find smallest element greater than the given one. +-- +-- > lookupGT 4 (fromList [3, 5]) == Just 5 +-- > lookupGT 5 (fromList [3, 5]) == Nothing +lookupGT :: Ord a => a -> Set a -> Maybe a +lookupGT = goNothing + where + STRICT_1_OF_2(goNothing) + 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 x best (Bin _ y l r) | x < y = goJust x y l + | otherwise = goJust x best r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE lookupGT #-} +#else +{-# INLINE lookupGT #-} +#endif + +-- | /O(log n)/. Find largest element smaller or equal to the given one. +-- +-- > lookupLE 2 (fromList [3, 5]) == Nothing +-- > lookupLE 4 (fromList [3, 5]) == Just 3 +-- > lookupLE 5 (fromList [3, 5]) == Just 5 +lookupLE :: Ord a => a -> Set a -> Maybe a +lookupLE = goNothing + where + STRICT_1_OF_2(goNothing) + 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 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 +{-# INLINABLE lookupLE #-} +#else +{-# INLINE lookupLE #-} +#endif + +-- | /O(log n)/. Find smallest element greater or equal to the given one. +-- +-- > lookupGE 3 (fromList [3, 5]) == Just 3 +-- > lookupGE 4 (fromList [3, 5]) == Just 5 +-- > lookupGE 6 (fromList [3, 5]) == Nothing +lookupGE :: Ord a => a -> Set a -> Maybe a +lookupGE = goNothing + where + STRICT_1_OF_2(goNothing) + 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 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 +{-# INLINABLE lookupGE #-} +#else +{-# INLINE lookupGE #-} +#endif + +{-------------------------------------------------------------------- + Construction +--------------------------------------------------------------------} +-- | /O(1)/. The empty set. +empty :: Set a +empty = Tip +{-# INLINE empty #-} + +-- | /O(1)/. Create a singleton set. +singleton :: a -> Set a +singleton x = Bin 1 x Tip Tip +{-# INLINE singleton #-} + +{-------------------------------------------------------------------- + Insertion, Deletion +--------------------------------------------------------------------} +-- | /O(log n)/. Insert an element in a set. +-- If the set already contains an element equal to the given value, +-- it is replaced with the new value. + +-- See Note: Type of local 'go' function +insert :: Ord a => a -> Set a -> Set a +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 +{-# INLINABLE insert #-} +#else +{-# INLINE insert #-} +#endif + +-- Insert an element to the set only if it is not in the set. +-- Used by `union`. + +-- See Note: Type of local 'go' function +insertR :: Ord a => a -> Set a -> Set a +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) + EQ -> t +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertR #-} +#else +{-# INLINE insertR #-} +#endif + +-- | /O(log n)/. Delete an element from a set. + +-- See Note: Type of local 'go' function +delete :: Ord a => a -> Set a -> Set a +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) + EQ -> glue l r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE delete #-} +#else +{-# INLINE delete #-} +#endif + +{-------------------------------------------------------------------- + Subset +--------------------------------------------------------------------} +-- | /O(n+m)/. Is this a proper subset? (ie. a subset but not equal). +isProperSubsetOf :: Ord a => Set a -> Set a -> Bool +isProperSubsetOf s1 s2 + = (size s1 < size s2) && (isSubsetOf s1 s2) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE isProperSubsetOf #-} +#endif + + +-- | /O(n+m)/. Is this a subset? +-- @(s1 `isSubsetOf` s2)@ tells whether @s1@ is a subset of @s2@. +isSubsetOf :: Ord a => Set a -> Set a -> Bool +isSubsetOf t1 t2 + = (size t1 <= size t2) && (isSubsetOfX t1 t2) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE isSubsetOf #-} +#endif + +isSubsetOfX :: Ord a => Set a -> Set a -> Bool +isSubsetOfX Tip _ = True +isSubsetOfX _ Tip = False +isSubsetOfX (Bin _ x l r) t + = found && isSubsetOfX l lt && isSubsetOfX r gt + where + (lt,found,gt) = splitMember x t +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE isSubsetOfX #-} +#endif + + +{-------------------------------------------------------------------- + Minimal, Maximal +--------------------------------------------------------------------} +-- | /O(log n)/. The minimal element of a set. +findMin :: Set a -> a +findMin (Bin _ x Tip _) = x +findMin (Bin _ _ l _) = findMin l +findMin Tip = error "Set.findMin: empty set has no minimal element" + +-- | /O(log n)/. The maximal element of a set. +findMax :: Set a -> a +findMax (Bin _ x _ Tip) = x +findMax (Bin _ _ _ r) = findMax r +findMax Tip = error "Set.findMax: empty set has no maximal element" + +-- | /O(log n)/. Delete the minimal element. +deleteMin :: Set a -> Set a +deleteMin (Bin _ _ Tip r) = r +deleteMin (Bin _ x l r) = balanceR x (deleteMin l) r +deleteMin Tip = Tip + +-- | /O(log n)/. Delete the maximal element. +deleteMax :: Set a -> Set a +deleteMax (Bin _ _ l Tip) = l +deleteMax (Bin _ x l r) = balanceL x l (deleteMax r) +deleteMax Tip = Tip + +{-------------------------------------------------------------------- + Union. +--------------------------------------------------------------------} +-- | 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 +{-# INLINABLE unions #-} +#endif + +-- | /O(n+m)/. The union of two sets, preferring the first set when +-- equal elements are encountered. +-- The implementation uses the efficient /hedge-union/ algorithm. +-- Hedge-union is more efficient on (bigset `union` smallset). +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 +{-# 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) = join 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 = join 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/. +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 +{-# 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 = join 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. +-- Elements of the result come from the first set, so for example +-- +-- > import qualified Data.Set as S +-- > data AB = A | B deriving Show +-- > instance Ord AB where compare _ _ = EQ +-- > instance Eq AB where _ == _ = True +-- > main = print (S.singleton A `S.intersection` S.singleton B, +-- > S.singleton B `S.intersection` S.singleton A) +-- +-- prints @(fromList [A],fromList [B])@. +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 +{-# 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 join 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 = join x (filter p l) (filter p r) + | otherwise = merge (filter p l) (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. +-- See also 'split'. +partition :: (a -> Bool) -> Set a -> (Set a,Set a) +partition _ Tip = (Tip, Tip) +partition p (Bin _ x l r) = case (partition p l, partition p r) of + ((l1, l2), (r1, r2)) + | p x -> (join x l1 r1, merge l2 r2) + | otherwise -> (merge l1 r1, join x l2 r2) + +{---------------------------------------------------------------------- + Map +----------------------------------------------------------------------} + +-- | /O(n*log n)/. +-- @'map' f s@ is the set obtained by applying @f@ to each element of @s@. +-- +-- It's worth noting that the size of the result may be smaller if, +-- for some @(x,y)@, @x \/= y && f x == f y@ + +map :: (Ord a, Ord b) => (a->b) -> Set a -> Set b +map f = fromList . List.map f . toList +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE map #-} +#endif + +-- | /O(n)/. The +-- +-- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is monotonic. +-- /The precondition is not checked./ +-- Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapMonotonic f s == map f s +-- > where ls = toList s + +mapMonotonic :: (a->b) -> Set a -> Set b +mapMonotonic _ Tip = Tip +mapMonotonic f (Bin sz x l r) = Bin sz (f x) (mapMonotonic f l) (mapMonotonic f r) + +{-------------------------------------------------------------------- + Fold +--------------------------------------------------------------------} +-- | /O(n)/. Fold the elements in the set using the given right-associative +-- binary operator. This function is an equivalent of 'foldr' and is present +-- for compatibility only. +-- +-- /Please note that fold will be deprecated in the future and removed./ +fold :: (a -> b -> b) -> b -> Set a -> b +fold = foldr +{-# INLINE fold #-} + +-- | /O(n)/. Fold the elements in the set using the given right-associative +-- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'toAscList'@. +-- +-- For example, +-- +-- > toAscList set = foldr (:) [] set +foldr :: (a -> b -> b) -> b -> Set 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 -> Set 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 elements in the set using the given left-associative +-- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'toAscList'@. +-- +-- For example, +-- +-- > toDescList set = foldl (flip (:)) [] set +foldl :: (a -> b -> a) -> a -> Set 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 -> Set 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' #-} + +{-------------------------------------------------------------------- + List variations +--------------------------------------------------------------------} +-- | /O(n)/. An alias of 'toAscList'. The elements of a set in ascending order. +-- Subject to list fusion. +elems :: Set a -> [a] +elems = toAscList + +{-------------------------------------------------------------------- + Lists +--------------------------------------------------------------------} +-- | /O(n)/. Convert the set to a list of elements. Subject to list fusion. +toList :: Set a -> [a] +toList = toAscList + +-- | /O(n)/. Convert the set to an ascending list of elements. Subject to list fusion. +toAscList :: Set a -> [a] +toAscList = foldr (:) [] + +-- | /O(n)/. Convert the set to a descending list of elements. Subject to list +-- fusion. +toDescList :: Set a -> [a] +toDescList = foldl (flip (:)) [] + +-- List fusion for the list generating functions. +#if __GLASGOW_HASKELL__ +-- The foldrFB and foldlFB are foldr and foldl equivalents, used for list fusion. +-- They are important to convert unfused to{Asc,Desc}List back, see mapFB in prelude. +foldrFB :: (a -> b -> b) -> b -> Set a -> b +foldrFB = foldr +{-# INLINE[0] foldrFB #-} +foldlFB :: (a -> b -> a) -> a -> Set b -> a +foldlFB = foldl +{-# INLINE[0] foldlFB #-} + +-- Inline elems and toList, so that we need to fuse only toAscList. +{-# INLINE elems #-} +{-# INLINE toList #-} + +-- The fusion is enabled up to phase 2 included. If it does not succeed, +-- convert in phase 1 the expanded to{Asc,Desc}List calls back to +-- 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 to{Asc,Desc}List -- it was forbidden to inline it +-- before phase 0, otherwise the fusion rules would not fire at all. +{-# NOINLINE[0] toAscList #-} +{-# NOINLINE[0] toDescList #-} +{-# RULES "Set.toAscList" [~1] forall s . toAscList s = build (\c n -> foldrFB c n s) #-} +{-# RULES "Set.toAscListBack" [1] foldrFB (:) [] = toAscList #-} +{-# RULES "Set.toDescList" [~1] forall s . toDescList s = build (\c n -> foldlFB (\xs x -> c x xs) n s) #-} +{-# RULES "Set.toDescListBack" [1] foldlFB (\xs x -> x : xs) [] = toDescList #-} +#endif + +-- | /O(n*log n)/. Create a set from a list of elements. +fromList :: Ord a => [a] -> Set a +fromList = foldlStrict ins empty + where + ins t x = insert x t +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromList #-} +#endif + +{-------------------------------------------------------------------- + Building trees from ascending/descending lists can be done in linear time. + + Note that if [xs] is ascending that: + fromAscList xs == fromList xs +--------------------------------------------------------------------} +-- | /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 +{-# INLINABLE fromAscList #-} +#endif + + +-- | /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./ +fromDistinctAscList :: [a] -> Set a +fromDistinctAscList xs + = create const (length xs) xs + where + -- 1) use continutations so that we use heap space instead of stack space. + -- 2) special case for n==5 to create bushier trees. + create c 0 xs' = c Tip xs' + create c 5 xs' = case xs' of + (x1:x2:x3:x4:x5:xx) + -> c (bin x4 (bin x2 (singleton x1) (singleton x3)) (singleton x5)) xx + _ -> error "fromDistinctAscList create 5" + create c n xs' = seq nr $ create (createR nr c) nl xs' + where nl = n `div` 2 + nr = n - nl - 1 + + createR n c l (x:ys) = create (createB l x c) n ys + createR _ _ _ [] = error "fromDistinctAscList createR []" + createB l x c r zs = c (bin x 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 + and it is certainly the simplest :-) +--------------------------------------------------------------------} +instance Eq a => Eq (Set a) where + t1 == t2 = (size t1 == size t2) && (toAscList t1 == toAscList t2) + +{-------------------------------------------------------------------- + Ord +--------------------------------------------------------------------} + +instance Ord a => Ord (Set a) where + compare s1 s2 = compare (toAscList s1) (toAscList s2) + +{-------------------------------------------------------------------- + Show +--------------------------------------------------------------------} +instance Show a => Show (Set a) where + showsPrec p xs = showParen (p > 10) $ + showString "fromList " . shows (toList xs) + +{-------------------------------------------------------------------- + Read +--------------------------------------------------------------------} +instance (Read a, Ord a) => Read (Set 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 + +{-------------------------------------------------------------------- + Typeable/Data +--------------------------------------------------------------------} + +#include "Typeable.h" +INSTANCE_TYPEABLE1(Set,setTc,"Set") + +{-------------------------------------------------------------------- + NFData +--------------------------------------------------------------------} + +instance NFData a => NFData (Set a) where + rnf Tip = () + 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 -> join 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 -> join 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 _ Tip = (Tip,Tip) +split x (Bin _ y l r) + = case compare x y of + LT -> let (lt,gt) = split x l in (lt,join y gt r) + GT -> let (lt,gt) = split x r in (join y l lt,gt) + EQ -> (l,r) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE split #-} +#endif + +-- | /O(log n)/. Performs a 'split' but also returns whether the pivot +-- element was found in the original set. +splitMember :: Ord a => a -> Set a -> (Set a,Bool,Set a) +splitMember _ Tip = (Tip, False, Tip) +splitMember x (Bin _ y l r) + = case compare x y of + LT -> let (lt, found, gt) = splitMember x l in (lt, found, join y gt r) + GT -> let (lt, found, gt) = splitMember x r in (join y l lt, found, gt) + EQ -> (l, True, r) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE splitMember #-} +#endif + +{-------------------------------------------------------------------- + Utility functions that maintain the balance properties of the tree. + All constructors assume that all values in [l] < [x] and all values + in [r] > [x], and that [l] and [r] are valid trees. + + In order of sophistication: + [Bin sz x l r] The type constructor. + [bin x l r] Maintains the correct size, assumes that both [l] + and [r] are balanced with respect to each other. + [balance 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. + [join 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 [join], [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. +--------------------------------------------------------------------} + +{-------------------------------------------------------------------- + Join +--------------------------------------------------------------------} +join :: a -> Set a -> Set a -> Set a +join x Tip r = insertMin x r +join x l Tip = insertMax x l +join x l@(Bin sizeL y ly ry) r@(Bin sizeR z lz rz) + | delta*sizeL < sizeR = balanceL z (join x l lz) rz + | delta*sizeR < sizeL = balanceR y ly (join x ry r) + | otherwise = bin x l r + + +-- insertMin and insertMax don't perform potentially expensive comparisons. +insertMax,insertMin :: a -> Set a -> Set a +insertMax x t + = case t of + Tip -> singleton x + Bin _ y l r + -> balanceR y l (insertMax x r) + +insertMin x t + = case t of + Tip -> singleton x + Bin _ y l r + -> balanceL y (insertMin x l) r + +{-------------------------------------------------------------------- + [merge l r]: merges two trees. +--------------------------------------------------------------------} +merge :: Set a -> Set a -> Set a +merge Tip r = r +merge l Tip = l +merge l@(Bin sizeL x lx rx) r@(Bin sizeR y ly ry) + | delta*sizeL < sizeR = balanceL y (merge l ly) ry + | delta*sizeR < sizeL = balanceR 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 :: Set a -> Set a -> Set a +glue Tip r = r +glue l Tip = l +glue l r + | size l > size r = let (m,l') = deleteFindMax l in balanceR m l' r + | otherwise = let (m,r') = deleteFindMin r in balanceL m l r' + +-- | /O(log n)/. Delete and find the minimal element. +-- +-- > deleteFindMin set = (findMin set, deleteMin set) + +deleteFindMin :: Set a -> (a,Set a) +deleteFindMin t + = case t of + Bin _ x Tip r -> (x,r) + Bin _ x l r -> let (xm,l') = deleteFindMin l in (xm,balanceR x l' r) + Tip -> (error "Set.deleteFindMin: can not return the minimal element of an empty set", Tip) + +-- | /O(log n)/. Delete and find the maximal element. +-- +-- > deleteFindMax set = (findMax set, deleteMax set) +deleteFindMax :: Set a -> (a,Set a) +deleteFindMax t + = case t of + Bin _ x l Tip -> (x,l) + Bin _ x l r -> let (xm,r') = deleteFindMax r in (xm,balanceL x l r') + Tip -> (error "Set.deleteFindMax: can not return the maximal element of an empty set", Tip) + +-- | /O(log n)/. Retrieves the minimal key of the set, and the set +-- stripped of that element, or 'Nothing' if passed an empty set. +minView :: Set a -> Maybe (a, Set a) +minView Tip = Nothing +minView x = Just (deleteFindMin x) + +-- | /O(log n)/. Retrieves the maximal key of the set, and the set +-- stripped of that element, or 'Nothing' if passed an empty set. +maxView :: Set a -> Maybe (a, Set a) +maxView Tip = Nothing +maxView x = Just (deleteFindMax x) + +{-------------------------------------------------------------------- + [balance x l 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 correspondes 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 errorneous: + - 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 :: a -> Set a -> Set a -> Set a +-- balance x l r +-- | sizeL + sizeR <= 1 = Bin sizeX x l r +-- | sizeR > delta*sizeL = rotateL x l r +-- | sizeL > delta*sizeR = rotateR x l r +-- | otherwise = Bin sizeX x l r +-- where +-- sizeL = size l +-- sizeR = size r +-- sizeX = sizeL + sizeR + 1 +-- +-- rotateL :: a -> Set a -> Set a -> Set a +-- rotateL x l r@(Bin _ _ ly ry) | size ly < ratio*size ry = singleL x l r +-- | otherwise = doubleL x l r +-- rotateR :: a -> Set a -> Set a -> Set a +-- rotateR x l@(Bin _ _ ly ry) r | size ry < ratio*size ly = singleR x l r +-- | otherwise = doubleR x l r +-- +-- singleL, singleR :: a -> Set a -> Set a -> Set a +-- singleL x1 t1 (Bin _ x2 t2 t3) = bin x2 (bin x1 t1 t2) t3 +-- singleR x1 (Bin _ x2 t1 t2) t3 = bin x2 t1 (bin x1 t2 t3) +-- +-- doubleL, doubleR :: a -> Set a -> Set a -> Set a +-- doubleL x1 t1 (Bin _ x2 (Bin _ x3 t2 t3) t4) = bin x3 (bin x1 t1 t2) (bin x2 t3 t4) +-- doubleR x1 (Bin _ x2 t1 (Bin _ x3 t2 t3)) t4 = bin x3 (bin x2 t1 t2) (bin x1 t3 t4) +-- +-- It is only written in such a way that every node is pattern-matched only once. +-- +-- Only balanceL and balanceR are needed at the moment, so balance is not here anymore. +-- In case it is needed, it can be found in Data.Map. + +-- 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 :: a -> Set a -> Set a -> Set a +balanceL x l r = case r of + Tip -> case l of + Tip -> Bin 1 x Tip Tip + (Bin _ _ Tip Tip) -> Bin 2 x l Tip + (Bin _ lx Tip (Bin _ lrx _ _)) -> Bin 3 lrx (Bin 1 lx Tip Tip) (Bin 1 x Tip Tip) + (Bin _ lx ll@(Bin _ _ _ _) Tip) -> Bin 3 lx ll (Bin 1 x Tip Tip) + (Bin ls lx ll@(Bin lls _ _ _) lr@(Bin lrs lrx lrl lrr)) + | lrs < ratio*lls -> Bin (1+ls) lx ll (Bin (1+lrs) x lr Tip) + | otherwise -> Bin (1+ls) lrx (Bin (1+lls+size lrl) lx ll lrl) (Bin (1+size lrr) x lrr Tip) + + (Bin rs _ _ _) -> case l of + Tip -> Bin (1+rs) x Tip r + + (Bin ls lx ll lr) + | ls > delta*rs -> case (ll, lr) of + (Bin lls _ _ _, Bin lrs lrx lrl lrr) + | lrs < ratio*lls -> Bin (1+ls+rs) lx ll (Bin (1+rs+lrs) x lr r) + | otherwise -> Bin (1+ls+rs) lrx (Bin (1+lls+size lrl) lx ll lrl) (Bin (1+rs+size lrr) x lrr r) + (_, _) -> error "Failure in Data.Map.balanceL" + | otherwise -> Bin (1+ls+rs) 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 :: a -> Set a -> Set a -> Set a +balanceR x l r = case l of + Tip -> case r of + Tip -> Bin 1 x Tip Tip + (Bin _ _ Tip Tip) -> Bin 2 x Tip r + (Bin _ rx Tip rr@(Bin _ _ _ _)) -> Bin 3 rx (Bin 1 x Tip Tip) rr + (Bin _ rx (Bin _ rlx _ _) Tip) -> Bin 3 rlx (Bin 1 x Tip Tip) (Bin 1 rx Tip Tip) + (Bin rs rx rl@(Bin rls rlx rll rlr) rr@(Bin rrs _ _ _)) + | rls < ratio*rrs -> Bin (1+rs) rx (Bin (1+rls) x Tip rl) rr + | otherwise -> Bin (1+rs) rlx (Bin (1+size rll) x Tip rll) (Bin (1+rrs+size rlr) rx rlr rr) + + (Bin ls _ _ _) -> case r of + Tip -> Bin (1+ls) x l Tip + + (Bin rs rx rl rr) + | rs > delta*ls -> case (rl, rr) of + (Bin rls rlx rll rlr, Bin rrs _ _ _) + | rls < ratio*rrs -> Bin (1+ls+rs) rx (Bin (1+ls+rls) x l rl) rr + | otherwise -> Bin (1+ls+rs) rlx (Bin (1+ls+size rll) x l rll) (Bin (1+rrs+size rlr) rx rlr rr) + (_, _) -> error "Failure in Data.Map.balanceR" + | otherwise -> Bin (1+ls+rs) x l r +{-# NOINLINE balanceR #-} + +{-------------------------------------------------------------------- + The bin constructor maintains the size of the tree +--------------------------------------------------------------------} +bin :: a -> Set a -> Set a -> Set a +bin x l r + = Bin (size l + size r + 1) x l r +{-# INLINE bin #-} + + +{-------------------------------------------------------------------- + Utilities +--------------------------------------------------------------------} +foldlStrict :: (a -> b -> a) -> a -> [b] -> a +foldlStrict f = go + where + go z [] = z + go z (x:xs) = let z' = f z x in z' `seq` go z' xs +{-# INLINE foldlStrict #-} + +{-------------------------------------------------------------------- + Debugging +--------------------------------------------------------------------} +-- | /O(n)/. Show the tree that implements the set. The tree is shown +-- in a compressed, hanging format. +showTree :: Show a => Set a -> String +showTree s + = showTreeWith True False s + + +{- | /O(n)/. The expression (@showTreeWith hang wide map@) shows + the tree that implements the set. 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. + +> Set> putStrLn $ showTreeWith True False $ fromDistinctAscList [1..5] +> 4 +> +--2 +> | +--1 +> | +--3 +> +--5 +> +> Set> putStrLn $ showTreeWith True True $ fromDistinctAscList [1..5] +> 4 +> | +> +--2 +> | | +> | +--1 +> | | +> | +--3 +> | +> +--5 +> +> Set> putStrLn $ showTreeWith False True $ fromDistinctAscList [1..5] +> +--5 +> | +> 4 +> | +> | +--3 +> | | +> +--2 +> | +> +--1 + +-} +showTreeWith :: Show a => Bool -> Bool -> Set a -> String +showTreeWith hang wide t + | hang = (showsTreeHang wide [] t) "" + | otherwise = (showsTree wide [] [] t) "" + +showsTree :: Show a => Bool -> [String] -> [String] -> Set a -> ShowS +showsTree wide lbars rbars t + = case t of + Tip -> showsBars lbars . showString "|\n" + Bin _ x Tip Tip + -> showsBars lbars . shows x . showString "\n" + Bin _ x l r + -> showsTree wide (withBar rbars) (withEmpty rbars) r . + showWide wide rbars . + showsBars lbars . shows x . showString "\n" . + showWide wide lbars . + showsTree wide (withEmpty lbars) (withBar lbars) l + +showsTreeHang :: Show a => Bool -> [String] -> Set a -> ShowS +showsTreeHang wide bars t + = case t of + Tip -> showsBars bars . showString "|\n" + Bin _ x Tip Tip + -> showsBars bars . shows x . showString "\n" + Bin _ x l r + -> showsBars bars . shows x . showString "\n" . + showWide wide bars . + showsTreeHang wide (withBar bars) l . + showWide wide bars . + showsTreeHang 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 + +{-------------------------------------------------------------------- + Assertions +--------------------------------------------------------------------} +-- | /O(n)/. Test if the internal set structure is valid. +valid :: Ord a => Set a -> Bool +valid t + = balanced t && ordered t && validsize t + +ordered :: Ord a => Set a -> Bool +ordered t + = bounded (const True) (const True) t + where + bounded lo hi t' + = case t' of + Tip -> True + Bin _ x l r -> (lo x) && (hi x) && bounded lo (x) hi r + +balanced :: Set a -> Bool +balanced t + = case t of + Tip -> True + Bin _ _ l r -> (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) && + balanced l && balanced r + +validsize :: Set a -> Bool +validsize t + = (realsize t == Just (size t)) + where + realsize t' + = case t' of + Tip -> Just 0 + Bin sz _ l r -> case (realsize l,realsize r) of + (Just n,Just m) | n+m+1 == sz -> Just sz + _ -> Nothing diff --git a/Data/StrictPair.hs b/Data/StrictPair.hs new file mode 100644 --- /dev/null +++ b/Data/StrictPair.hs @@ -0,0 +1,6 @@ +module Data.StrictPair (strictPair) where + +-- | Evaluate both argument to WHNF and create a pair of the result. +strictPair :: a -> b -> (a, b) +strictPair x y = x `seq` y `seq` (x, y) +{-# INLINE strictPair #-} diff --git a/Data/Tree.hs b/Data/Tree.hs --- a/Data/Tree.hs +++ b/Data/Tree.hs @@ -1,3 +1,7 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif #if __GLASGOW_HASKELL__ >= 703 {-# LANGUAGE Safe #-} #endif @@ -6,7 +10,7 @@ -- Module : Data.Tree -- Copyright : (c) The University of Glasgow 2002 -- License : BSD-style (see the file libraries/base/LICENSE) --- +-- -- Maintainer : libraries@haskell.org -- Stability : experimental -- Portability : portable diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,5 +1,5 @@ This library (libraries/containers) is derived from code from several -sources: +sources: * Code from the GHC project which is largely (c) The University of Glasgow, and distributable under a BSD-style license (see below), @@ -19,7 +19,7 @@ The Glasgow Haskell Compiler License -Copyright 2004, The University Court of the University of Glasgow. +Copyright 2004, The University Court of the University of Glasgow. All rights reserved. Redistribution and use in source and binary forms, with or without @@ -27,14 +27,14 @@ - Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. - + - Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. - + - Neither name of the University nor the names of its contributors may be used to endorse or promote products derived from this software without -specific prior written permission. +specific prior written permission. THIS SOFTWARE IS PROVIDED BY THE UNIVERSITY COURT OF THE UNIVERSITY OF GLASGOW AND THE CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, diff --git a/benchmarks/IntMap.hs b/benchmarks/IntMap.hs new file mode 100644 --- /dev/null +++ b/benchmarks/IntMap.hs @@ -0,0 +1,94 @@ +{-# LANGUAGE BangPatterns #-} +module Main where + +import Control.DeepSeq +import Control.Exception (evaluate) +import Control.Monad.Trans (liftIO) +import Criterion.Config +import Criterion.Main +import Data.List (foldl') +import qualified Data.IntMap as M +import Data.Maybe (fromMaybe) +import Prelude hiding (lookup) + +main = do + let m = M.fromAscList elems :: M.IntMap Int + defaultMainWith + defaultConfig + (liftIO . evaluate $ rnf [m]) + [ bench "lookup" $ whnf (lookup keys) m + , bench "insert" $ whnf (ins elems) M.empty + , bench "insertWith empty" $ whnf (insWith elems) M.empty + , bench "insertWith update" $ whnf (insWith elems) m + , bench "insertWith' empty" $ whnf (insWith' elems) M.empty + , bench "insertWith' update" $ whnf (insWith' elems) m + , bench "insertWithKey empty" $ whnf (insWithKey elems) M.empty + , bench "insertWithKey update" $ whnf (insWithKey elems) m + , bench "insertWithKey' empty" $ whnf (insWithKey' elems) M.empty + , bench "insertWithKey' update" $ whnf (insWithKey' elems) m + , bench "insertLookupWithKey empty" $ whnf (insLookupWithKey elems) M.empty + , bench "insertLookupWithKey update" $ whnf (insLookupWithKey elems) m + , 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" $ whnf (del keys) m + , bench "update" $ whnf (upd keys) m + , bench "updateLookupWithKey" $ whnf (upd' keys) m + , bench "alter" $ whnf (alt keys) m + , bench "mapMaybe" $ whnf (M.mapMaybe maybeDel) m + , bench "mapMaybeWithKey" $ whnf (M.mapMaybeWithKey (const maybeDel)) m + ] + where + elems = zip keys values + keys = [1..2^12] + values = [1..2^12] + sum k v1 v2 = k + v1 + v2 + consPair k v xs = (k, v) : xs + +add3 :: Int -> Int -> Int -> Int +add3 x y z = x + y + z +{-# INLINE add3 #-} + +lookup :: [Int] -> M.IntMap Int -> Int +lookup xs m = foldl' (\n k -> fromMaybe n (M.lookup k m)) 0 xs + +ins :: [(Int, Int)] -> M.IntMap Int -> M.IntMap Int +ins xs m = foldl' (\m (k, v) -> M.insert k v m) m xs + +insWith :: [(Int, Int)] -> M.IntMap Int -> M.IntMap Int +insWith xs m = foldl' (\m (k, v) -> M.insertWith (+) k v m) m xs + +insWithKey :: [(Int, Int)] -> M.IntMap Int -> M.IntMap Int +insWithKey xs m = foldl' (\m (k, v) -> M.insertWithKey add3 k v m) m xs + +insWith' :: [(Int, Int)] -> M.IntMap Int -> M.IntMap Int +insWith' xs m = foldl' (\m (k, v) -> M.insertWith' (+) k v m) m xs + +insWithKey' :: [(Int, Int)] -> M.IntMap Int -> M.IntMap Int +insWithKey' xs m = foldl' (\m (k, v) -> M.insertWithKey' add3 k v m) m xs + +data PairS a b = PS !a !b + +insLookupWithKey :: [(Int, Int)] -> M.IntMap Int -> (Int, M.IntMap Int) +insLookupWithKey xs m = let !(PS a b) = foldl' f (PS 0 m) xs in (a, b) + where + f (PS n m) (k, v) = let !(n', m') = M.insertLookupWithKey add3 k v m + in PS (fromMaybe 0 n' + n) m' + +del :: [Int] -> M.IntMap Int -> M.IntMap Int +del xs m = foldl' (\m k -> M.delete k m) m xs + +upd :: [Int] -> M.IntMap Int -> M.IntMap Int +upd xs m = foldl' (\m k -> M.update Just k m) m xs + +upd' :: [Int] -> M.IntMap Int -> M.IntMap Int +upd' xs m = foldl' (\m k -> snd $ M.updateLookupWithKey (\_ a -> Just a) k m) m xs + +alt :: [Int] -> M.IntMap Int -> M.IntMap Int +alt xs m = foldl' (\m k -> M.alter id k m) m xs + +maybeDel :: Int -> Maybe Int +maybeDel n | n `mod` 3 == 0 = Nothing + | otherwise = Just n diff --git a/benchmarks/IntSet.hs b/benchmarks/IntSet.hs new file mode 100644 --- /dev/null +++ b/benchmarks/IntSet.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE BangPatterns #-} + +module Main where + +import Control.DeepSeq +import Control.Exception (evaluate) +import Control.Monad.Trans (liftIO) +import Criterion.Config +import Criterion.Main +import Data.List (foldl') +import qualified Data.IntSet as S + +main = do + let s = S.fromAscList elems :: S.IntSet + s_even = S.fromAscList elems_even :: S.IntSet + s_odd = S.fromAscList elems_odd :: S.IntSet + defaultMainWith + defaultConfig + (liftIO . evaluate $ rnf [s, s_even, s_odd]) + [ bench "member" $ whnf (member elems) s + , bench "insert" $ whnf (ins elems) S.empty + , bench "map" $ whnf (S.map (+ 1)) s + , bench "filter" $ whnf (S.filter ((== 0) . (`mod` 2))) s + , bench "partition" $ whnf (S.partition ((== 0) . (`mod` 2))) s + , bench "fold" $ whnf (S.fold (:) []) s + , bench "delete" $ whnf (del elems) s + , bench "findMin" $ whnf S.findMin s + , bench "findMax" $ whnf S.findMax s + , bench "deleteMin" $ whnf S.deleteMin s + , bench "deleteMax" $ whnf S.deleteMax s + , bench "unions" $ whnf S.unions [s_even, s_odd] + , bench "union" $ whnf (S.union s_even) s_odd + , bench "difference" $ whnf (S.difference s) s_even + , bench "intersection" $ whnf (S.intersection s) s_even + ] + where + elems = [1..2^10] + elems_even = [2,4..2^10] + elems_odd = [1,3..2^10] + +member :: [Int] -> S.IntSet -> Int +member xs s = foldl' (\n x -> if S.member x s then n + 1 else n) 0 xs + +ins :: [Int] -> S.IntSet -> S.IntSet +ins xs s0 = foldl' (\s a -> S.insert a s) s0 xs + +del :: [Int] -> S.IntSet -> S.IntSet +del xs s0 = foldl' (\s k -> S.delete k s) s0 xs diff --git a/benchmarks/LookupGE/IntMap.hs b/benchmarks/LookupGE/IntMap.hs new file mode 100644 --- /dev/null +++ b/benchmarks/LookupGE/IntMap.hs @@ -0,0 +1,51 @@ +{-# LANGUAGE BangPatterns #-} +module Main where + +import Control.DeepSeq +import Control.Exception (evaluate) +import Control.Monad.Trans (liftIO) +import Criterion.Config +import Criterion.Main +import Data.List (foldl') +import qualified Data.IntMap as M +import qualified LookupGE_IntMap as M +import Data.Maybe (fromMaybe) +import Prelude hiding (lookup) + +main :: IO () +main = do + defaultMainWith + defaultConfig + (liftIO . evaluate $ rnf [m_even, m_odd, m_large]) + [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 + m_large = M.fromAscList elems_large :: M.IntMap Int + bound = 2^12 + elems_even = zip evens evens + elems_odd = zip odds odds + elems_large = zip large large + evens = [2,4..bound] + odds = [1,3..bound] + large = [1,100..50*bound] + benches = + [ \(n,fun) -> bench (n++" present") $ nf (fge fun evens) m_even + , \(n,fun) -> bench (n++" absent") $ nf (fge fun evens) m_odd + , \(n,fun) -> bench (n++" far") $ nf (fge fun odds) m_large + , \(n,fun) -> bench (n++" !present") $ nf (fge2 fun evens) m_even + , \(n,fun) -> bench (n++" !absent") $ nf (fge2 fun evens) m_odd + , \(n,fun) -> bench (n++" !far") $ nf (fge2 fun odds) m_large + ] + funs1 = [ ("GE split", M.lookupGE1) + , ("GE Craig", M.lookupGE2) + , ("GE Twan", M.lookupGE3) + , ("GE Milan", M.lookupGE4) ] + +fge :: (Int -> M.IntMap Int -> Maybe (Int,Int)) -> [Int] -> M.IntMap Int -> (Int,Int) +fge fun xs m = foldl' (\n k -> fromMaybe n (fun k m)) (0,0) xs + +-- forcing values inside tuples! +fge2 :: (Int -> M.IntMap Int -> Maybe (Int,Int)) -> [Int] -> M.IntMap Int -> (Int,Int) +fge2 fun xs m = foldl' (\n@(!_, !_) k -> fromMaybe n (fun k m)) (0,0) xs + diff --git a/benchmarks/LookupGE/LookupGE_IntMap.hs b/benchmarks/LookupGE/LookupGE_IntMap.hs new file mode 100644 --- /dev/null +++ b/benchmarks/LookupGE/LookupGE_IntMap.hs @@ -0,0 +1,97 @@ +{-# LANGUAGE CPP #-} +module LookupGE_IntMap where + +import Prelude hiding (null) +import Data.IntMap.Base +#ifdef TESTING +import Test.QuickCheck +#endif + +lookupGE1 :: Key -> IntMap a -> Maybe (Key,a) +lookupGE1 k m = + case splitLookup k m of + (_,Just v,_) -> Just (k,v) + (_,Nothing,r) -> findMinMaybe r + +lookupGE2 :: Key -> IntMap a -> Maybe (Key,a) +lookupGE2 k t = case t of + Bin _ m l r | m < 0 -> if k >= 0 + then go l + else case go r of + Nothing -> Just $ findMin l + justx -> justx + _ -> go t + where + go (Bin p m l r) + | nomatch k p m = if k < p + then Just $ findMin l + else Nothing + | zero k m = case go l of + Nothing -> Just $ findMin r + justx -> justx + | otherwise = go r + go (Tip ky y) + | k > ky = Nothing + | otherwise = Just (ky, y) + go Nil = Nothing + +lookupGE3 :: Key -> IntMap a -> Maybe (Key,a) +lookupGE3 k t = k `seq` case t of + Bin _ m l r | m < 0 -> if k >= 0 + then go Nothing l + else go (Just (findMin l)) r + _ -> go Nothing t + where + go def (Bin p m l r) + | nomatch k p m = if k < p then Just $ findMin l else def + | zero k m = go (Just $ findMin r) l + | otherwise = go def r + go def (Tip ky y) + | k > ky = def + | otherwise = Just (ky, y) + go def Nil = def + +lookupGE4 :: Key -> IntMap a -> Maybe (Key,a) +lookupGE4 k t = k `seq` case t of + Bin _ m l r | m < 0 -> if k >= 0 then go Nil l + else go l r + _ -> go Nil t + where + go def (Bin p m l r) + | nomatch k p m = if k < p then fMin l else fMin def + | zero k m = go r l + | otherwise = go def r + go def (Tip ky y) + | k > ky = fMin def + | otherwise = Just (ky, y) + go def Nil = fMin def + + fMin :: IntMap a -> Maybe (Key, a) + fMin Nil = Nothing + fMin (Tip ky y) = Just (ky, y) + fMin (Bin _ _ l _) = fMin l + +------------------------------------------------------------------------------- +-- Utilities +------------------------------------------------------------------------------- + +-- | /O(log n)/. The minimal key of the map. +findMinMaybe :: IntMap a -> Maybe (Key, a) +findMinMaybe m + | null m = Nothing + | otherwise = Just (findMin m) + +#ifdef TESTING +------------------------------------------------------------------------------- +-- Properties: +------------------------------------------------------------------------------- + +prop_lookupGE12 :: Int -> [Int] -> Bool +prop_lookupGE12 x xs = case fromList $ zip xs xs of m -> lookupGE1 x m == lookupGE2 x m + +prop_lookupGE13 :: Int -> [Int] -> Bool +prop_lookupGE13 x xs = case fromList $ zip xs xs of m -> lookupGE1 x m == lookupGE3 x m + +prop_lookupGE14 :: Int -> [Int] -> Bool +prop_lookupGE14 x xs = case fromList $ zip xs xs of m -> lookupGE1 x m == lookupGE4 x m +#endif diff --git a/benchmarks/LookupGE/LookupGE_Map.hs b/benchmarks/LookupGE/LookupGE_Map.hs new file mode 100644 --- /dev/null +++ b/benchmarks/LookupGE/LookupGE_Map.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE BangPatterns, CPP #-} +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 = + case splitLookup k m of + (_,Just v,_) -> Just (k,v) + (_,Nothing,r) -> findMinMaybe r +{-# INLINABLE lookupGE1 #-} + +lookupGE2 :: Ord k => k -> Map k a -> Maybe (k,a) +lookupGE2 = go + where + go !_ Tip = Nothing + go !k (Bin _ kx x l r) = + case compare k kx of + LT -> case go k l of + Nothing -> Just (kx,x) + ret -> ret + GT -> go k r + EQ -> Just (kx,x) +{-# INLINABLE lookupGE2 #-} + +lookupGE3 :: Ord k => k -> Map k a -> Maybe (k,a) +lookupGE3 = go Nothing + where + go def !_ Tip = def + go def !k (Bin _ kx x l r) = + case compare k kx of + LT -> go (Just (kx,x)) k l + GT -> go def k r + EQ -> Just (kx,x) +{-# INLINABLE lookupGE3 #-} + +lookupGE4 :: Ord k => k -> Map k a -> Maybe (k,a) +lookupGE4 k = k `seq` goNothing + where + goNothing Tip = Nothing + goNothing (Bin _ kx x l r) = case compare k kx of + LT -> goJust kx x l + EQ -> Just (kx, x) + GT -> goNothing r + + goJust ky y Tip = Just (ky, y) + goJust ky y (Bin _ kx x l r) = case compare k kx of + LT -> goJust kx x l + EQ -> Just (kx, x) + GT -> goJust ky y r +{-# INLINABLE lookupGE4 #-} + +------------------------------------------------------------------------------- +-- Utilities +------------------------------------------------------------------------------- + +findMinMaybe :: Map k a -> Maybe (k,a) +findMinMaybe (Bin _ kx x Tip _) = Just (kx,x) +findMinMaybe (Bin _ _ _ l _) = findMinMaybe l +findMinMaybe Tip = Nothing + +#ifdef TESTING +------------------------------------------------------------------------------- +-- Properties: +------------------------------------------------------------------------------- + +prop_lookupGE12 :: Int -> [Int] -> Bool +prop_lookupGE12 x xs = case fromList $ zip xs xs of m -> lookupGE1 x m == lookupGE2 x m + +prop_lookupGE13 :: Int -> [Int] -> Bool +prop_lookupGE13 x xs = case fromList $ zip xs xs of m -> lookupGE1 x m == lookupGE3 x m + +prop_lookupGE14 :: Int -> [Int] -> Bool +prop_lookupGE14 x xs = case fromList $ zip xs xs of m -> lookupGE1 x m == lookupGE4 x m +#endif diff --git a/benchmarks/LookupGE/Makefile b/benchmarks/LookupGE/Makefile new file mode 100644 --- /dev/null +++ b/benchmarks/LookupGE/Makefile @@ -0,0 +1,3 @@ +TOP = .. + +include ../Makefile diff --git a/benchmarks/LookupGE/Map.hs b/benchmarks/LookupGE/Map.hs new file mode 100644 --- /dev/null +++ b/benchmarks/LookupGE/Map.hs @@ -0,0 +1,50 @@ +{-# LANGUAGE BangPatterns #-} +module Main where + +import Control.DeepSeq +import Control.Exception (evaluate) +import Control.Monad.Trans (liftIO) +import Criterion.Config +import Criterion.Main +import Data.List (foldl') +import qualified Data.Map as M +import qualified LookupGE_Map as M +import Data.Maybe (fromMaybe) +import Prelude hiding (lookup) + +main :: IO () +main = do + defaultMainWith + defaultConfig + (liftIO . evaluate $ rnf [m_even, m_odd, m_large]) + [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 + m_large = M.fromAscList elems_large :: M.Map Int Int + bound = 2^10 + elems_even = zip evens evens + elems_odd = zip odds odds + elems_large = zip large large + evens = [2,4..bound] + odds = [1,3..bound] + large = [1,100..50*bound] + benches = + [ \(n,fun) -> bench (n++" present") $ nf (fge fun evens) m_even + , \(n,fun) -> bench (n++" absent") $ nf (fge fun evens) m_odd + , \(n,fun) -> bench (n++" far") $ nf (fge fun odds) m_large + , \(n,fun) -> bench (n++" !present") $ nf (fge2 fun evens) m_even + , \(n,fun) -> bench (n++" !absent") $ nf (fge2 fun evens) m_odd + , \(n,fun) -> bench (n++" !far") $ nf (fge2 fun odds) m_large + ] + funs1 = [ ("GE split", M.lookupGE1) + , ("GE caseof", M.lookupGE2) + , ("GE Twan", M.lookupGE3) + , ("GE Milan", M.lookupGE4) ] + +fge :: (Int -> M.Map Int Int -> Maybe (Int,Int)) -> [Int] -> M.Map Int Int -> (Int,Int) +fge fun xs m = foldl' (\n k -> fromMaybe n (fun k m)) (0,0) xs + +-- forcing values inside tuples! +fge2 :: (Int -> M.Map Int Int -> Maybe (Int,Int)) -> [Int] -> M.Map Int Int -> (Int,Int) +fge2 fun xs m = foldl' (\n@(!_, !_) k -> fromMaybe n (fun k m)) (0,0) xs diff --git a/benchmarks/Makefile b/benchmarks/Makefile new file mode 100644 --- /dev/null +++ b/benchmarks/Makefile @@ -0,0 +1,16 @@ +all: + +bench-%: %.hs force + ghc -O2 -DTESTING $< -i../$(TOP) -o $@ -outputdir tmp -rtsopts + +bench-%.csv: bench-% + ./bench-$* -v -u bench-$*.csv + +.PHONY: force clean veryclean +force: + +clean: + rm -rf tmp $(patsubst %.hs, bench-%, $(wildcard *.hs)) + +veryclean: clean + rm -rf *.csv diff --git a/benchmarks/Map.hs b/benchmarks/Map.hs new file mode 100644 --- /dev/null +++ b/benchmarks/Map.hs @@ -0,0 +1,126 @@ +{-# LANGUAGE BangPatterns #-} +module Main where + +import Control.DeepSeq +import Control.Exception (evaluate) +import Control.Monad.Trans (liftIO) +import Criterion.Config +import Criterion.Main +import Data.List (foldl') +import qualified Data.Map as M +import Data.Maybe (fromMaybe) +import Prelude hiding (lookup) + +main = do + let m = M.fromAscList elems :: M.Map Int Int + m_even = M.fromAscList elems_even :: M.Map Int Int + m_odd = M.fromAscList elems_odd :: M.Map Int Int + defaultMainWith + defaultConfig + (liftIO . evaluate $ rnf [m, m_even, m_odd]) + [ bench "lookup absent" $ whnf (lookup evens) m_odd + , bench "lookup present" $ whnf (lookup evens) m_even + , bench "insert absent" $ whnf (ins elems_even) m_odd + , bench "insert present" $ whnf (ins elems_even) m_even + , 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 + , bench "insertWith' present" $ whnf (insWith' elems_even) m_even + , bench "insertWithKey absent" $ whnf (insWithKey elems_even) m_odd + , bench "insertWithKey present" $ whnf (insWithKey elems_even) m_even + , bench "insertWithKey' absent" $ whnf (insWithKey' elems_even) m_odd + , bench "insertWithKey' present" $ whnf (insWithKey' elems_even) m_even + , bench "insertLookupWithKey absent" $ whnf (insLookupWithKey elems_even) m_odd + , 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 + , bench "union" $ whnf (M.union m_even) m_odd + , bench "difference" $ whnf (M.difference m) m_even + , bench "intersection" $ whnf (M.intersection m) m_even + ] + where + bound = 2^10 + elems = zip keys values + elems_even = zip evens evens + elems_odd = zip odds odds + keys = [1..bound] + evens = [2,4..bound] + odds = [1,3..bound] + values = [1..bound] + sum k v1 v2 = k + v1 + v2 + consPair k v xs = (k, v) : xs + +add3 :: Int -> Int -> Int -> Int +add3 x y z = x + y + z +{-# INLINE add3 #-} + +lookup :: [Int] -> M.Map Int Int -> Int +lookup xs m = foldl' (\n k -> fromMaybe n (M.lookup 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 + +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 + +insWithKey :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int +insWithKey xs m = foldl' (\m (k, v) -> M.insertWithKey add3 k v 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 + +insWithKey' :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int +insWithKey' xs m = foldl' (\m (k, v) -> M.insertWithKey' add3 k v m) m xs + +data PairS a b = PS !a !b + +insLookupWithKey :: [(Int, Int)] -> M.Map Int Int -> (Int, M.Map Int Int) +insLookupWithKey xs m = let !(PS a b) = foldl' f (PS 0 m) xs in (a, b) + where + f (PS n m) (k, v) = let !(n', m') = M.insertLookupWithKey add3 k v m + in PS (fromMaybe 0 n' + n) m' + +insLookupWithKey' :: [(Int, Int)] -> M.Map Int Int -> (Int, M.Map Int Int) +insLookupWithKey' xs m = let !(PS a b) = foldl' f (PS 0 m) xs in (a, b) + where + f (PS n m) (k, v) = let !(n', m') = M.insertLookupWithKey' add3 k v m + in PS (fromMaybe 0 n' + n) m' + +del :: [Int] -> M.Map Int Int -> M.Map Int Int +del xs m = foldl' (\m k -> M.delete 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 + +upd' :: (Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int +upd' f xs m = foldl' (\m k -> snd $ M.updateLookupWithKey (\_ a -> f a) k m) m xs + +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 + +maybeDel :: Int -> Maybe Int +maybeDel n | n `mod` 3 == 0 = Nothing + | otherwise = Just n diff --git a/benchmarks/Sequence.hs b/benchmarks/Sequence.hs new file mode 100644 --- /dev/null +++ b/benchmarks/Sequence.hs @@ -0,0 +1,34 @@ +-- > ghc -DTESTING --make -O2 -fforce-recomp -i.. Sequence.hs +module Main where + +import Control.DeepSeq +import Criterion.Main +import Data.List (foldl') +import qualified Data.Sequence as S +import qualified Data.Foldable +import System.Random + +main = do + let s10 = S.fromList [1..10] :: S.Seq Int + s100 = S.fromList [1..100] :: S.Seq Int + s1000 = S.fromList [1..1000] :: S.Seq Int + rnf [s10, s100, s1000] `seq` return () + let g = mkStdGen 1 + let rlist n = map (`mod` (n+1)) (take 10000 (randoms g)) :: [Int] + r10 = rlist 10 + r100 = rlist 100 + r1000 = rlist 1000 + rnf [r10, r100, r1000] `seq` return () + defaultMain + [ bench "splitAt/append 10" $ nf (shuffle r10) s10 + , bench "splitAt/append 100" $ nf (shuffle r100) s100 + , bench "splitAt/append 1000" $ nf (shuffle r1000) s1000 + ] + +-- 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. +shuffle :: [Int] -> S.Seq Int -> Int +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 diff --git a/benchmarks/Set.hs b/benchmarks/Set.hs new file mode 100644 --- /dev/null +++ b/benchmarks/Set.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE BangPatterns #-} + +-- > ghc -DTESTING --make -O2 -fforce-recomp -i.. Set.hs +module Main where + +import Control.DeepSeq +import Control.Exception (evaluate) +import Control.Monad.Trans (liftIO) +import Criterion.Config +import Criterion.Main +import Data.List (foldl') +import qualified Data.Set as S + +main = do + let s = S.fromAscList elems :: S.Set Int + s_even = S.fromAscList elems_even :: S.Set Int + s_odd = S.fromAscList elems_odd :: S.Set Int + defaultMainWith + defaultConfig + (liftIO . evaluate $ rnf [s, s_even, s_odd]) + [ bench "member" $ whnf (member elems) s + , bench "insert" $ whnf (ins elems) S.empty + , bench "map" $ whnf (S.map (+ 1)) s + , bench "filter" $ whnf (S.filter ((== 0) . (`mod` 2))) s + , bench "partition" $ whnf (S.partition ((== 0) . (`mod` 2))) s + , bench "fold" $ whnf (S.fold (:) []) s + , bench "delete" $ whnf (del elems) s + , bench "findMin" $ whnf S.findMin s + , bench "findMax" $ whnf S.findMax s + , bench "deleteMin" $ whnf S.deleteMin s + , bench "deleteMax" $ whnf S.deleteMax s + , bench "unions" $ whnf S.unions [s_even, s_odd] + , bench "union" $ whnf (S.union s_even) s_odd + , bench "difference" $ whnf (S.difference s) s_even + , bench "intersection" $ whnf (S.intersection s) s_even + ] + where + elems = [1..2^10] + elems_even = [2,4..2^10] + elems_odd = [1,3..2^10] + +member :: [Int] -> S.Set Int -> Int +member xs s = foldl' (\n x -> if S.member x s then n + 1 else n) 0 xs + +ins :: [Int] -> S.Set Int -> S.Set Int +ins xs s0 = foldl' (\s a -> S.insert a s) s0 xs + +del :: [Int] -> S.Set Int -> S.Set Int +del xs s0 = foldl' (\s k -> S.delete k s) s0 xs diff --git a/benchmarks/SetOperations/Makefile b/benchmarks/SetOperations/Makefile new file mode 100644 --- /dev/null +++ b/benchmarks/SetOperations/Makefile @@ -0,0 +1,3 @@ +TOP = .. + +include ../Makefile diff --git a/benchmarks/SetOperations/SetOperations-IntMap.hs b/benchmarks/SetOperations/SetOperations-IntMap.hs new file mode 100644 --- /dev/null +++ b/benchmarks/SetOperations/SetOperations-IntMap.hs @@ -0,0 +1,6 @@ +module Main where + +import Data.IntMap as C +import SetOperations + +main = benchmark (\xs -> fromList [(x, x) | x <- xs]) True [("union", C.union), ("difference", C.difference), ("intersection", C.intersection)] diff --git a/benchmarks/SetOperations/SetOperations-IntSet.hs b/benchmarks/SetOperations/SetOperations-IntSet.hs new file mode 100644 --- /dev/null +++ b/benchmarks/SetOperations/SetOperations-IntSet.hs @@ -0,0 +1,6 @@ +module Main where + +import Data.IntSet as C +import SetOperations + +main = benchmark fromList True [("union", C.union), ("difference", C.difference), ("intersection", C.intersection)] diff --git a/benchmarks/SetOperations/SetOperations-Map.hs b/benchmarks/SetOperations/SetOperations-Map.hs new file mode 100644 --- /dev/null +++ b/benchmarks/SetOperations/SetOperations-Map.hs @@ -0,0 +1,6 @@ +module Main where + +import Data.Map as C +import SetOperations + +main = benchmark (\xs -> fromList [(x, x) | x <- xs]) True [("union", C.union), ("difference", C.difference), ("intersection", C.intersection)] diff --git a/benchmarks/SetOperations/SetOperations-Set.hs b/benchmarks/SetOperations/SetOperations-Set.hs new file mode 100644 --- /dev/null +++ b/benchmarks/SetOperations/SetOperations-Set.hs @@ -0,0 +1,6 @@ +module Main where + +import Data.Set as C +import SetOperations + +main = benchmark fromList True [("union", C.union), ("difference", C.difference), ("intersection", C.intersection)] diff --git a/benchmarks/SetOperations/SetOperations.hs b/benchmarks/SetOperations/SetOperations.hs new file mode 100644 --- /dev/null +++ b/benchmarks/SetOperations/SetOperations.hs @@ -0,0 +1,45 @@ +{-# LANGUAGE BangPatterns #-} + +module SetOperations (benchmark) where + +import Criterion.Main +import Data.List (partition) + +benchmark :: ([Int] -> container) -> Bool -> [(String, container -> container -> container)] -> IO () +benchmark fromList swap methods = do + defaultMain $ [ bench (method_str++"-"++input_str) $ whnf (method input1) input2 | (method_str, method) <- methods, (input_str, input1, input2) <- inputs ] + + where + n, s, t :: Int + n = 100000 + s {-small-} = n `div` 10 + t {-tiny-} = round $ sqrt $ fromIntegral n + + inputs = [ (mode_str, left, right) + | (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) + ] + + , (mode_str, left, right) <- replicate 2 (mode_str, left, right) ++ + replicate (if swap && take 4 mode_str /= "diff" && last mode_str /= last (init mode_str) then 2 else 0) + (init (init mode_str) ++ [last mode_str] ++ [last (init mode_str)], right, left) + ] + + all_n = fromList [1..n] + + !disj_nn = seqPair $ (all_n, fromList [n+1..n+n]) + !disj_ns = seqPair $ (all_n, fromList [n+1..n+s]) + !disj_nt = seqPair $ (all_n, fromList [n+1..n+t]) + !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] + + fromLists (xs, ys) = (fromList xs, fromList ys) + seqPair pair@(xs, ys) = xs `seq` ys `seq` pair diff --git a/benchmarks/bench-cmp.pl b/benchmarks/bench-cmp.pl new file mode 100644 --- /dev/null +++ b/benchmarks/bench-cmp.pl @@ -0,0 +1,24 @@ +#!/usr/bin/perl +use warnings; +use strict; + +@ARGV >= 2 or die "Usage: bench-cmp.pl csv_file_1 csv_file_2"; +open (my $f1, "<", $ARGV[0]) or die "Cannot open file $ARGV[0]"; +open (my $f2, "<", $ARGV[1]) or die "Cannot open file $ARGV[1]"; + +my $l1 = <$f1>; +my $l2 = <$f2>; +$l1 eq $l2 or die "CSV files do not correspond -- $l1 and $l2"; + +while (defined($l1 = <$f1>)) { + $l2 = <$f2>; + + my @parts1 = split /,/, $l1; + my @parts2 = split /,/, $l2; + + $parts1[0] eq $parts2[0] or die "CSV files do not correspond -- $parts1[0] and $parts2[0]"; + printf "%s;%+7.2f%%;%.2e\n", $parts1[0], 100 * $parts2[1] / $parts1[1] - 100, $parts1[1]; +} + +close $f2; +close $f1; diff --git a/benchmarks/bench-cmp.sh b/benchmarks/bench-cmp.sh new file mode 100644 --- /dev/null +++ b/benchmarks/bench-cmp.sh @@ -0,0 +1,3 @@ +#!/bin/sh + +./bench-cmp.pl "$@" | column -nts\; | less -SR diff --git a/containers.cabal b/containers.cabal --- a/containers.cabal +++ b/containers.cabal @@ -1,9 +1,9 @@ name: containers -version: 0.4.2.1 +version: 0.5.0.0 license: BSD3 license-file: LICENSE maintainer: fox@ucw.cz -bug-reports: http://hackage.haskell.org/trac/ghc/newticket?component=libraries%20%28other%29 +bug-reports: https://github.com/haskell/containers/issues synopsis: Assorted concrete container types category: Data Structures description: @@ -12,34 +12,182 @@ each operation is either worst-case or amortized, but remains valid even if structures are shared. build-type: Simple -cabal-version: >=1.6 -extra-source-files: include/Typeable.h +cabal-version: >=1.8 +extra-source-files: + include/Typeable.h + tests/Makefile + tests/*.hs + benchmarks/Makefile + benchmarks/bench-cmp.pl + benchmarks/bench-cmp.sh + benchmarks/*.hs + benchmarks/SetOperations/Makefile + benchmarks/SetOperations/*.hs + benchmarks/LookupGE/Makefile + benchmarks/LookupGE/*.hs source-repository head type: git location: http://github.com/haskell/containers.git -Library { +Library build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.4 - ghc-options: -O2 - if impl(ghc>6.10) - Ghc-Options: -fregs-graph + if impl(ghc>=6.10) + build-depends: ghc-prim + + ghc-options: -O2 -Wall + exposed-modules: Data.IntMap + Data.IntMap.Lazy + Data.IntMap.Strict Data.IntSet Data.Map + Data.Map.Lazy + Data.Map.Strict Data.Set - include-dirs: include - extensions: CPP - if !impl(nhc98) { + if !impl(nhc98) exposed-modules: Data.Graph Data.Sequence Data.Tree - } - if impl(ghc) { - extensions: DeriveDataTypeable, StandaloneDeriving, - MagicHash, Rank2Types - } -} + other-modules: + Data.IntMap.Base + Data.IntSet.Base + Data.Map.Base + Data.Set.Base + Data.StrictPair + include-dirs: include + + if impl(ghc<7.0) + extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types + +------------------- +-- T E S T I N G -- +------------------- + +-- 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.4, ghc-prim + ghc-options: -O2 + extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types + + build-depends: + HUnit, + QuickCheck, + test-framework, + test-framework-hunit, + test-framework-quickcheck2 + +Test-suite map-strict-properties + hs-source-dirs: tests, . + main-is: map-properties.hs + type: exitcode-stdio-1.0 + cpp-options: -DTESTING -DSTRICT + + build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.4, ghc-prim + ghc-options: -O2 + extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types + + build-depends: + HUnit, + QuickCheck, + test-framework, + test-framework-hunit, + test-framework-quickcheck2 + +Test-suite set-properties + hs-source-dirs: tests, . + main-is: set-properties.hs + type: exitcode-stdio-1.0 + cpp-options: -DTESTING + + build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.4, ghc-prim + ghc-options: -O2 + extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types + + build-depends: + HUnit, + QuickCheck, + test-framework, + test-framework-hunit, + test-framework-quickcheck2 + +Test-suite intmap-lazy-properties + hs-source-dirs: tests, . + main-is: intmap-properties.hs + type: exitcode-stdio-1.0 + cpp-options: -DTESTING + + build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.4, ghc-prim + ghc-options: -O2 + extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types + + build-depends: + HUnit, + QuickCheck, + test-framework, + test-framework-hunit, + test-framework-quickcheck2 + +Test-suite intmap-strict-properties + hs-source-dirs: tests, . + main-is: intmap-properties.hs + type: exitcode-stdio-1.0 + cpp-options: -DTESTING -DSTRICT + + build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.4, ghc-prim + ghc-options: -O2 + extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types + + build-depends: + HUnit, + QuickCheck, + test-framework, + test-framework-hunit, + test-framework-quickcheck2 + +Test-suite intset-properties + hs-source-dirs: tests, . + main-is: intset-properties.hs + type: exitcode-stdio-1.0 + cpp-options: -DTESTING + + build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.4, ghc-prim + ghc-options: -O2 + extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types + + build-depends: + HUnit, + QuickCheck, + test-framework, + test-framework-hunit, + test-framework-quickcheck2 + +Test-suite seq-properties + hs-source-dirs: tests, . + main-is: seq-properties.hs + type: exitcode-stdio-1.0 + cpp-options: -DTESTING + + build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.4, ghc-prim + ghc-options: -O2 + extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types + + build-depends: + QuickCheck, + test-framework, + test-framework-quickcheck2 diff --git a/include/Typeable.h b/include/Typeable.h --- a/include/Typeable.h +++ b/include/Typeable.h @@ -3,11 +3,11 @@ // // INSTANCE_TYPEABLEn(tc,tcname,"tc") defines // -// instance Typeable/n/ tc -// instance Typeable a => Typeable/n-1/ (tc a) -// instance (Typeable a, Typeable b) => Typeable/n-2/ (tc a b) -// ... -// instance (Typeable a1, ..., Typeable an) => Typeable (tc a1 ... an) +// instance Typeable/n/ tc +// instance Typeable a => Typeable/n-1/ (tc a) +// instance (Typeable a, Typeable b) => Typeable/n-2/ (tc a b) +// ... +// instance (Typeable a1, ..., Typeable an) => Typeable (tc a1 ... an) // -------------------------------------------------------------------------- -} diff --git a/tests/Makefile b/tests/Makefile new file mode 100644 --- /dev/null +++ b/tests/Makefile @@ -0,0 +1,20 @@ +# The tests should be compiled and run using cabal: +# > cabal configure --enable-tests +# > cabal build +# > cabal test +# +# This Makefile is used by developers to compile the tests manually. + +all: + +%-properties: %-properties.hs force + ghc -O2 -DTESTING $< -i.. -o $@ -outputdir tmp + +%-strict-properties: %-properties.hs force + ghc -O2 -DTESTING -DSTRICT $< -o $@ -i.. -outputdir tmp + +.PHONY: force clean +force: + +clean: + rm -rf tmp $(patsubst %.hs, %, $(wildcard *-properties.hs)) $(patsubst %-properties.hs, %-strict-properties, $(wildcard *-properties.hs)) diff --git a/tests/intmap-properties.hs b/tests/intmap-properties.hs new file mode 100644 --- /dev/null +++ b/tests/intmap-properties.hs @@ -0,0 +1,1041 @@ +{-# LANGUAGE CPP #-} + +#ifdef STRICT +import Data.IntMap.Strict as Data.IntMap +#else +import Data.IntMap.Lazy as Data.IntMap +#endif + +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 Data.List (nub,sort) +import qualified Data.List as List +import qualified Data.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 () + +default (Int) + +main :: IO () +main = defaultMainWithOpts + [ + testCase "index" test_index + , testCase "size" test_size + , testCase "size2" test_size2 + , testCase "member" test_member + , testCase "notMember" test_notMember + , testCase "lookup" test_lookup + , testCase "findWithDefault" test_findWithDefault + , testCase "lookupLT" test_lookupLT + , testCase "lookupGT" test_lookupGT + , testCase "lookupLE" test_lookupLE + , testCase "lookupGE" test_lookupGE + , testCase "empty" test_empty + , testCase "mempty" test_mempty + , testCase "singleton" test_singleton + , testCase "insert" test_insert + , testCase "insertWith" test_insertWith + , testCase "insertWithKey" test_insertWithKey + , testCase "insertLookupWithKey" test_insertLookupWithKey + , testCase "delete" test_delete + , testCase "adjust" test_adjust + , testCase "adjustWithKey" test_adjustWithKey + , testCase "update" test_update + , testCase "updateWithKey" test_updateWithKey + , testCase "updateLookupWithKey" test_updateLookupWithKey + , testCase "alter" test_alter + , testCase "union" test_union + , testCase "mappend" test_mappend + , testCase "unionWith" test_unionWith + , testCase "unionWithKey" test_unionWithKey + , testCase "unions" test_unions + , testCase "mconcat" test_mconcat + , testCase "unionsWith" test_unionsWith + , testCase "difference" test_difference + , testCase "differenceWith" test_differenceWith + , testCase "differenceWithKey" test_differenceWithKey + , testCase "intersection" test_intersection + , testCase "intersectionWith" test_intersectionWith + , testCase "intersectionWithKey" test_intersectionWithKey + , testCase "map" test_map + , testCase "mapWithKey" test_mapWithKey + , testCase "mapAccum" test_mapAccum + , testCase "mapAccumWithKey" test_mapAccumWithKey + , testCase "mapAccumRWithKey" test_mapAccumRWithKey + , testCase "mapKeys" test_mapKeys + , testCase "mapKeysWith" test_mapKeysWith + , testCase "mapKeysMonotonic" test_mapKeysMonotonic + , testCase "elems" test_elems + , testCase "keys" test_keys + , testCase "assocs" test_assocs + , testCase "keysSet" test_keysSet + , testCase "keysSet" test_fromSet + , testCase "toList" test_toList + , testCase "fromList" test_fromList + , testCase "fromListWith" test_fromListWith + , testCase "fromListWithKey" test_fromListWithKey + , testCase "toAscList" test_toAscList + , testCase "toDescList" test_toDescList + , testCase "showTree" test_showTree + , testCase "fromAscList" test_fromAscList + , testCase "fromAscListWith" test_fromAscListWith + , testCase "fromAscListWithKey" test_fromAscListWithKey + , testCase "fromDistinctAscList" test_fromDistinctAscList + , testCase "filter" test_filter + , testCase "filterWithKey" test_filteWithKey + , testCase "partition" test_partition + , testCase "partitionWithKey" test_partitionWithKey + , testCase "mapMaybe" test_mapMaybe + , testCase "mapMaybeWithKey" test_mapMaybeWithKey + , testCase "mapEither" test_mapEither + , testCase "mapEitherWithKey" test_mapEitherWithKey + , testCase "split" test_split + , testCase "splitLookup" test_splitLookup + , testCase "isSubmapOfBy" test_isSubmapOfBy + , testCase "isSubmapOf" test_isSubmapOf + , testCase "isProperSubmapOfBy" test_isProperSubmapOfBy + , testCase "isProperSubmapOf" test_isProperSubmapOf + , testCase "findMin" test_findMin + , testCase "findMax" test_findMax + , testCase "deleteMin" test_deleteMin + , testCase "deleteMax" test_deleteMax + , testCase "deleteFindMin" test_deleteFindMin + , testCase "deleteFindMax" test_deleteFindMax + , testCase "updateMin" test_updateMin + , testCase "updateMax" test_updateMax + , testCase "updateMinWithKey" test_updateMinWithKey + , testCase "updateMaxWithKey" test_updateMaxWithKey + , testCase "minView" test_minView + , testCase "maxView" test_maxView + , testCase "minViewWithKey" test_minViewWithKey + , testCase "maxViewWithKey" test_maxViewWithKey + , testProperty "insert to singleton" prop_singleton + , testProperty "insert then lookup" prop_insertLookup + , testProperty "insert then delete" prop_insertDelete + , testProperty "delete non member" prop_deleteNonMember + , testProperty "union model" prop_unionModel + , testProperty "union singleton" prop_unionSingleton + , testProperty "union associative" prop_unionAssoc + , testProperty "union+unionWith" prop_unionWith + , testProperty "union sum" prop_unionSum + , testProperty "difference model" prop_differenceModel + , testProperty "intersection model" prop_intersectionModel + , testProperty "intersectionWith model" prop_intersectionWithModel + , testProperty "intersectionWithKey model" prop_intersectionWithKeyModel + , testProperty "mergeWithKey model" prop_mergeWithKeyModel + , testProperty "fromAscList" prop_ordered + , testProperty "fromList then toList" prop_list + , testProperty "toDescList" prop_descList + , testProperty "toAscList+toDescList" prop_ascDescList + , testProperty "alter" prop_alter + , testProperty "index" prop_index + , testProperty "null" prop_null + , testProperty "member" prop_member + , testProperty "notmember" prop_notmember + , testProperty "lookup" prop_lookup + , testProperty "find" prop_find + , testProperty "findWithDefault" prop_findWithDefault + , testProperty "lookupLT" prop_lookupLT + , testProperty "lookupGT" prop_lookupGT + , testProperty "lookupLE" prop_lookupLE + , testProperty "lookupGE" prop_lookupGE + , testProperty "findMin" prop_findMin + , testProperty "findMax" prop_findMax + , testProperty "deleteMin" prop_deleteMinModel + , testProperty "deleteMax" prop_deleteMaxModel + , testProperty "filter" prop_filter + , testProperty "partition" prop_partition + , testProperty "map" prop_map + , testProperty "fmap" prop_fmap + , testProperty "mapkeys" prop_mapkeys + , testProperty "split" prop_splitModel + , testProperty "foldr" prop_foldr + , testProperty "foldr'" prop_foldr' + , testProperty "foldl" prop_foldl + , testProperty "foldl'" prop_foldl' + , testProperty "keysSet" prop_keysSet + , testProperty "fromSet" prop_fromSet + ] opts + + where + opts = mempty { ropt_test_options = Just $ mempty { topt_maximum_generated_tests = Just 500 + , topt_maximum_unsuitable_generated_tests = Just 500 + } + } + +{-------------------------------------------------------------------- + Arbitrary, reasonably balanced trees +--------------------------------------------------------------------} + +instance Arbitrary a => Arbitrary (IntMap a) where + arbitrary = do{ ks <- arbitrary + ; xs <- arbitrary + ; return (fromList (zip xs ks)) + } + + +------------------------------------------------------------------------ + +type UMap = IntMap () +type IMap = IntMap Int +type SMap = IntMap String + +---------------------------------------------------------------- + +tests :: [Test] +tests = [ testGroup "Test Case" [ + ] + , testGroup "Property Test" [ + ] + ] + + +---------------------------------------------------------------- +-- Unit tests +---------------------------------------------------------------- + +---------------------------------------------------------------- +-- Operators + +test_index :: Assertion +test_index = fromList [(5,'a'), (3,'b')] ! 5 @?= 'a' + +---------------------------------------------------------------- +-- Query + +test_size :: Assertion +test_size = do + null (empty) @?= True + null (singleton 1 'a') @?= False + +test_size2 :: Assertion +test_size2 = do + size empty @?= 0 + size (singleton 1 'a') @?= 1 + size (fromList([(1,'a'), (2,'c'), (3,'b')])) @?= 3 + +test_member :: Assertion +test_member = do + member 5 (fromList [(5,'a'), (3,'b')]) @?= True + member 1 (fromList [(5,'a'), (3,'b')]) @?= False + +test_notMember :: Assertion +test_notMember = do + notMember 5 (fromList [(5,'a'), (3,'b')]) @?= False + notMember 1 (fromList [(5,'a'), (3,'b')]) @?= True + +test_lookup :: Assertion +test_lookup = do + employeeCurrency 1 @?= Just 1 + employeeCurrency 2 @?= Nothing + where + employeeDept = fromList([(1,2), (3,1)]) + deptCountry = fromList([(1,1), (2,2)]) + countryCurrency = fromList([(1, 2), (2, 1)]) + employeeCurrency :: Int -> Maybe Int + employeeCurrency name = do + dept <- lookup name employeeDept + country <- lookup dept deptCountry + lookup country countryCurrency + +test_findWithDefault :: Assertion +test_findWithDefault = do + findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) @?= 'x' + findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) @?= 'a' + +test_lookupLT :: Assertion +test_lookupLT = do + lookupLT 3 (fromList [(3,'a'), (5,'b')]) @?= Nothing + lookupLT 4 (fromList [(3,'a'), (5,'b')]) @?= Just (3, 'a') + +test_lookupGT :: Assertion +test_lookupGT = do + lookupGT 4 (fromList [(3,'a'), (5,'b')]) @?= Just (5, 'b') + lookupGT 5 (fromList [(3,'a'), (5,'b')]) @?= Nothing + +test_lookupLE :: Assertion +test_lookupLE = do + 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') + +test_lookupGE :: Assertion +test_lookupGE = do + 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 + +---------------------------------------------------------------- +-- Construction + +test_empty :: Assertion +test_empty = do + (empty :: UMap) @?= fromList [] + size empty @?= 0 + +test_mempty :: Assertion +test_mempty = do + (mempty :: UMap) @?= fromList [] + size (mempty :: UMap) @?= 0 + +test_singleton :: Assertion +test_singleton = do + singleton 1 'a' @?= fromList [(1, 'a')] + size (singleton 1 'a') @?= 1 + +test_insert :: Assertion +test_insert = do + 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' + +test_insertWith :: Assertion +test_insertWith = do + 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" + +test_insertWithKey :: Assertion +test_insertWithKey = do + 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" + where + f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value + +test_insertLookupWithKey :: Assertion +test_insertLookupWithKey = do + insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) @?= (Just "a", fromList [(3, "b"), (5, "5:xxx|a")]) + insertLookupWithKey f 2 "xxx" (fromList [(5,"a"), (3,"b")]) @?= (Nothing,fromList [(2,"xxx"),(3,"b"),(5,"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") + where + f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value + +---------------------------------------------------------------- +-- Delete/Update + +test_delete :: Assertion +test_delete = do + 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 :: IMap) + +test_adjust :: Assertion +test_adjust = do + 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 + +test_adjustWithKey :: Assertion +test_adjustWithKey = do + 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 + where + f key x = (show key) ++ ":new " ++ x + +test_update :: Assertion +test_update = do + 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" + where + f x = if x == "a" then Just "new a" else Nothing + +test_updateWithKey :: Assertion +test_updateWithKey = do + 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" + where + f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing + +test_updateLookupWithKey :: Assertion +test_updateLookupWithKey = do + updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) @?= (Just "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") + where + f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing + +test_alter :: Assertion +test_alter = do + alter f 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a")] + alter f 5 (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "b" + alter g 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a"), (7, "c")] + alter g 5 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "c")] + where + f _ = Nothing + g _ = Just "c" + +---------------------------------------------------------------- +-- Combine + +test_union :: Assertion +test_union = union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "a"), (7, "C")] + +test_mappend :: Assertion +test_mappend = mappend (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "a"), (7, "C")] + +test_unionWith :: Assertion +test_unionWith = unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "aA"), (7, "C")] + +test_unionWithKey :: Assertion +test_unionWithKey = unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "5:a|A"), (7, "C")] + where + f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value + +test_unions :: Assertion +test_unions = do + 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")] + +test_mconcat :: Assertion +test_mconcat = do + mconcat [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] + @?= fromList [(3, "b"), (5, "a"), (7, "C")] + mconcat [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])] + @?= fromList [(3, "B3"), (5, "A3"), (7, "C")] + +test_unionsWith :: Assertion +test_unionsWith = unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] + @?= fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] + +test_difference :: Assertion +test_difference = difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 3 "b" + +test_differenceWith :: Assertion +test_differenceWith = differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")]) + @?= singleton 3 "b:B" + where + f al ar = if al== "b" then Just (al ++ ":" ++ ar) else Nothing + +test_differenceWithKey :: Assertion +test_differenceWithKey = differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")]) + @?= singleton 3 "3:b|B" + where + f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing + +test_intersection :: Assertion +test_intersection = intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "a" + + +test_intersectionWith :: Assertion +test_intersectionWith = intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "aA" + +test_intersectionWithKey :: Assertion +test_intersectionWithKey = intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "5:a|A" + where + f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar + +---------------------------------------------------------------- +-- Traversal + +test_map :: Assertion +test_map = map (++ "x") (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "bx"), (5, "ax")] + +test_mapWithKey :: Assertion +test_mapWithKey = mapWithKey f (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "3:b"), (5, "5:a")] + where + f key x = (show key) ++ ":" ++ x + +test_mapAccum :: Assertion +test_mapAccum = mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) @?= ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) + where + f a b = (a ++ b, b ++ "X") + +test_mapAccumWithKey :: Assertion +test_mapAccumWithKey = mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) @?= ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) + where + f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") + +test_mapAccumRWithKey :: Assertion +test_mapAccumRWithKey = mapAccumRWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) @?= ("Everything: 5-a 3-b", fromList [(3, "bX"), (5, "aX")]) + where + f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") + +test_mapKeys :: Assertion +test_mapKeys = do + 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" + +test_mapKeysWith :: Assertion +test_mapKeysWith = do + 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" + +test_mapKeysMonotonic :: Assertion +test_mapKeysMonotonic = do + mapKeysMonotonic (+ 1) (fromList [(5,"a"), (3,"b")]) @?= fromList [(4, "b"), (6, "a")] + mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) @?= fromList [(6, "b"), (10, "a")] + +---------------------------------------------------------------- +-- Conversion + +test_elems :: Assertion +test_elems = do + elems (fromList [(5,"a"), (3,"b")]) @?= ["b","a"] + elems (empty :: UMap) @?= [] + +test_keys :: Assertion +test_keys = do + keys (fromList [(5,"a"), (3,"b")]) @?= [3,5] + keys (empty :: UMap) @?= [] + +test_assocs :: Assertion +test_assocs = do + assocs (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")] + assocs (empty :: UMap) @?= [] + +test_keysSet :: Assertion +test_keysSet = do + keysSet (fromList [(5,"a"), (3,"b")]) @?= Data.IntSet.fromList [3,5] + keysSet (empty :: UMap) @?= Data.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) + +---------------------------------------------------------------- +-- Lists + +test_toList :: Assertion +test_toList = do + toList (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")] + toList (empty :: SMap) @?= [] + +test_fromList :: Assertion +test_fromList = do + fromList [] @?= (empty :: SMap) + fromList [(5,"a"), (3,"b"), (5, "c")] @?= fromList [(5,"c"), (3,"b")] + fromList [(5,"c"), (3,"b"), (5, "a")] @?= fromList [(5,"a"), (3,"b")] + +test_fromListWith :: Assertion +test_fromListWith = do + fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] @?= fromList [(3, "ab"), (5, "aba")] + fromListWith (++) [] @?= (empty :: SMap) + +test_fromListWithKey :: Assertion +test_fromListWithKey = do + fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] @?= fromList [(3, "3ab"), (5, "5a5ba")] + fromListWithKey f [] @?= (empty :: SMap) + where + f k a1 a2 = (show k) ++ a1 ++ a2 + +---------------------------------------------------------------- +-- Ordered lists + +test_toAscList :: Assertion +test_toAscList = toAscList (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")] + +test_toDescList :: Assertion +test_toDescList = toDescList (fromList [(5,"a"), (3,"b")]) @?= [(5,"a"), (3,"b")] + +test_showTree :: Assertion +test_showTree = + (let t = fromDistinctAscList [(x,()) | x <- [1..5]] + in showTree t) @?= "*\n+--*\n| +-- 1:=()\n| +--*\n| +-- 2:=()\n| +-- 3:=()\n+--*\n +-- 4:=()\n +-- 5:=()\n" + +test_fromAscList :: Assertion +test_fromAscList = do + fromAscList [(3,"b"), (5,"a")] @?= fromList [(3, "b"), (5, "a")] + fromAscList [(3,"b"), (5,"a"), (5,"b")] @?= fromList [(3, "b"), (5, "b")] + + +test_fromAscListWith :: Assertion +test_fromAscListWith = do + fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] @?= fromList [(3, "b"), (5, "ba")] + +test_fromAscListWithKey :: Assertion +test_fromAscListWithKey = do + fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] @?= fromList [(3, "b"), (5, "5:b5:ba")] + where + f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 + +test_fromDistinctAscList :: Assertion +test_fromDistinctAscList = do + fromDistinctAscList [(3,"b"), (5,"a")] @?= fromList [(3, "b"), (5, "a")] + +---------------------------------------------------------------- +-- Filter + +test_filter :: Assertion +test_filter = do + 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 + +test_filteWithKey :: Assertion +test_filteWithKey = filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a" + +test_partition :: Assertion +test_partition = do + 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")]) + +test_partitionWithKey :: Assertion +test_partitionWithKey = do + 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")]) + +test_mapMaybe :: Assertion +test_mapMaybe = mapMaybe f (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "new a" + where + f x = if x == "a" then Just "new a" else Nothing + +test_mapMaybeWithKey :: Assertion +test_mapMaybeWithKey = mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "key : 3" + where + f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing + +test_mapEither :: Assertion +test_mapEither = do + 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 :: SMap), fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) + where + f a = if a < "c" then Left a else Right a + +test_mapEitherWithKey :: Assertion +test_mapEitherWithKey = do + 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 :: SMap), fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")]) + where + f k a = if k < 5 then Left (k * 2) else Right (a ++ a) + +test_split :: Assertion +test_split = do + 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) + +test_splitLookup :: Assertion +test_splitLookup = do + 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) + +---------------------------------------------------------------- +-- Submap + +test_isSubmapOfBy :: Assertion +test_isSubmapOfBy = do + isSubmapOfBy (==) (fromList [(fromEnum 'a',1)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= True + isSubmapOfBy (<=) (fromList [(fromEnum 'a',1)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= True + isSubmapOfBy (==) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= True + isSubmapOfBy (==) (fromList [(fromEnum 'a',2)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= False + isSubmapOfBy (<) (fromList [(fromEnum 'a',1)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= False + isSubmapOfBy (==) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) (fromList [(fromEnum 'a',1)]) @?= False + +test_isSubmapOf :: Assertion +test_isSubmapOf = do + isSubmapOf (fromList [(fromEnum 'a',1)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= True + isSubmapOf (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= True + isSubmapOf (fromList [(fromEnum 'a',2)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= False + isSubmapOf (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) (fromList [(fromEnum 'a',1)]) @?= False + +test_isProperSubmapOfBy :: Assertion +test_isProperSubmapOfBy = do + isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= True + isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= True + isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) @?= False + isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)]) @?= False + isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= False + +test_isProperSubmapOf :: Assertion +test_isProperSubmapOf = do + isProperSubmapOf (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= True + isProperSubmapOf (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) @?= False + isProperSubmapOf (fromList [(1,1),(2,2)]) (fromList [(1,1)]) @?= False + +---------------------------------------------------------------- +-- Min/Max + +test_findMin :: Assertion +test_findMin = findMin (fromList [(5,"a"), (3,"b")]) @?= (3,"b") + +test_findMax :: Assertion +test_findMax = findMax (fromList [(5,"a"), (3,"b")]) @?= (5,"a") + +test_deleteMin :: Assertion +test_deleteMin = do + deleteMin (fromList [(5,"a"), (3,"b"), (7,"c")]) @?= fromList [(5,"a"), (7,"c")] + deleteMin (empty :: SMap) @?= empty + +test_deleteMax :: Assertion +test_deleteMax = do + deleteMax (fromList [(5,"a"), (3,"b"), (7,"c")]) @?= fromList [(3,"b"), (5,"a")] + deleteMax (empty :: SMap) @?= empty + +test_deleteFindMin :: Assertion +test_deleteFindMin = deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) @?= ((3,"b"), fromList[(5,"a"), (10,"c")]) + +test_deleteFindMax :: Assertion +test_deleteFindMax = deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) @?= ((10,"c"), fromList [(3,"b"), (5,"a")]) + +test_updateMin :: Assertion +test_updateMin = do + 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" + +test_updateMax :: Assertion +test_updateMax = do + 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" + +test_updateMinWithKey :: Assertion +test_updateMinWithKey = do + 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" + +test_updateMaxWithKey :: Assertion +test_updateMaxWithKey = do + 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" + +test_minView :: Assertion +test_minView = do + minView (fromList [(5,"a"), (3,"b")]) @?= Just ("b", singleton 5 "a") + minView (empty :: SMap) @?= Nothing + +test_maxView :: Assertion +test_maxView = do + maxView (fromList [(5,"a"), (3,"b")]) @?= Just ("a", singleton 3 "b") + maxView (empty :: SMap) @?= Nothing + +test_minViewWithKey :: Assertion +test_minViewWithKey = do + minViewWithKey (fromList [(5,"a"), (3,"b")]) @?= Just ((3,"b"), singleton 5 "a") + minViewWithKey (empty :: SMap) @?= Nothing + +test_maxViewWithKey :: Assertion +test_maxViewWithKey = do + maxViewWithKey (fromList [(5,"a"), (3,"b")]) @?= Just ((5,"a"), singleton 3 "b") + maxViewWithKey (empty :: SMap) @?= Nothing + +---------------------------------------------------------------- +-- QuickCheck +---------------------------------------------------------------- + +prop_singleton :: Int -> Int -> Bool +prop_singleton k x = insert k x empty == singleton k x + +prop_insertLookup :: Int -> UMap -> Bool +prop_insertLookup k t = lookup k (insert k () t) /= Nothing + +prop_insertDelete :: Int -> UMap -> Property +prop_insertDelete k t = (lookup k t == Nothing) ==> (delete k (insert k () t) == t) + +prop_deleteNonMember :: Int -> UMap -> Property +prop_deleteNonMember k t = (lookup k t == Nothing) ==> (delete k t == t) + +---------------------------------------------------------------- + +prop_unionModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_unionModel xs ys + = sort (keys (union (fromList xs) (fromList ys))) + == sort (nub (Prelude.map fst xs ++ Prelude.map fst ys)) + +prop_unionSingleton :: IMap -> Int -> Int -> Bool +prop_unionSingleton t k x = union (singleton k x) t == insert k x t + +prop_unionAssoc :: IMap -> IMap -> IMap -> Bool +prop_unionAssoc t1 t2 t3 = union t1 (union t2 t3) == union (union t1 t2) t3 + +prop_unionWith :: IMap -> IMap -> Bool +prop_unionWith t1 t2 = (union t1 t2 == unionWith (\_ y -> y) t2 t1) + +prop_unionSum :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_unionSum xs ys + = sum (elems (unionWith (+) (fromListWith (+) xs) (fromListWith (+) ys))) + == (sum (Prelude.map snd xs) + sum (Prelude.map snd ys)) + +prop_differenceModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_differenceModel xs ys + = sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys))) + == sort ((List.\\) (nub (Prelude.map fst xs)) (nub (Prelude.map fst ys))) + +prop_intersectionModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_intersectionModel xs ys + = sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys))) + == sort (nub ((List.intersect) (Prelude.map fst xs) (Prelude.map fst ys))) + +prop_intersectionWithModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_intersectionWithModel xs ys + = toList (intersectionWith f (fromList xs') (fromList ys')) + == [(kx, f vx vy ) | (kx, vx) <- List.sort xs', (ky, vy) <- ys', kx == ky] + where xs' = List.nubBy ((==) `on` fst) xs + ys' = List.nubBy ((==) `on` fst) ys + f l r = l + 2 * r + +prop_intersectionWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_intersectionWithKeyModel xs ys + = toList (intersectionWithKey f (fromList xs') (fromList ys')) + == [(kx, f kx vx vy) | (kx, vx) <- List.sort xs', (ky, vy) <- ys', kx == ky] + where xs' = List.nubBy ((==) `on` fst) xs + ys' = List.nubBy ((==) `on` fst) ys + f k l r = k + 2 * l + 3 * r + +prop_mergeWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_mergeWithKeyModel xs ys + = and [ testMergeWithKey f keep_x keep_y + | f <- [ \_k x1 _x2 -> Just x1 + , \_k _x1 x2 -> Just x2 + , \_k _x1 _x2 -> Nothing + , \k x1 x2 -> if k `mod` 2 == 0 then Nothing else Just (2 * x1 + 3 * x2) + ] + , keep_x <- [ True, False ] + , keep_y <- [ True, False ] + ] + + where xs' = List.nubBy ((==) `on` fst) xs + ys' = List.nubBy ((==) `on` fst) ys + + xm = fromList xs' + ym = fromList ys' + + testMergeWithKey f keep_x keep_y + = toList (mergeWithKey f (keep keep_x) (keep keep_y) xm ym) == emulateMergeWithKey f keep_x keep_y + where keep False _ = empty + keep True m = m + + emulateMergeWithKey f keep_x keep_y + = Maybe.mapMaybe combine (sort $ List.union (List.map fst xs') (List.map fst ys')) + where combine k = case (List.lookup k xs', List.lookup k ys') of + (Nothing, Just y) -> if keep_y then Just (k, y) else Nothing + (Just x, Nothing) -> if keep_x then Just (k, x) else Nothing + (Just x, Just y) -> (\v -> (k, v)) `fmap` f k x y + + -- We prevent inlining testMergeWithKey to disable the SpecConstr + -- optimalization. There are too many call patterns here so several + -- warnings are issued if testMergeWithKey gets inlined. + {-# NOINLINE testMergeWithKey #-} + +---------------------------------------------------------------- + +prop_ordered :: Property +prop_ordered + = forAll (choose (5,100)) $ \n -> + let xs = [(x,()) | x <- [0..n::Int]] + in fromAscList 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_ascDescList :: [Int] -> Bool +prop_ascDescList xs = toAscList m == reverse (toDescList m) + where m = fromList $ zip xs $ repeat () + +---------------------------------------------------------------- + +prop_alter :: UMap -> Int -> Bool +prop_alter t k = case lookup k t of + Just _ -> (size t - 1) == size t' && lookup k t' == Nothing + Nothing -> (size t + 1) == size t' && lookup k t' /= Nothing + where + t' = alter f k t + f Nothing = Just () + f (Just ()) = Nothing + +------------------------------------------------------------------------ +-- Compare against the list model (after nub on keys) + +prop_index :: [Int] -> Property +prop_index xs = length xs > 0 ==> + let m = fromList (zip xs xs) + in xs == [ m ! i | i <- xs ] + +prop_null :: IMap -> Bool +prop_null m = null m == (size m == 0) + +prop_member :: [Int] -> Int -> Bool +prop_member xs n = + let m = fromList (zip xs xs) + in all (\k -> k `member` m == (k `elem` xs)) (n : xs) + +prop_notmember :: [Int] -> Int -> Bool +prop_notmember xs n = + let m = fromList (zip xs xs) + in all (\k -> k `notMember` m == (k `notElem` xs)) (n : xs) + +prop_lookup :: [(Int, Int)] -> Int -> Bool +prop_lookup xs n = + let xs' = List.nubBy ((==) `on` fst) xs + m = fromList xs' + in all (\k -> lookup k m == List.lookup k xs') (n : List.map fst xs') + +prop_find :: [(Int, Int)] -> Bool +prop_find xs = + let xs' = List.nubBy ((==) `on` fst) xs + m = fromList xs' + in all (\(k, v) -> m ! k == v) xs' + +prop_findWithDefault :: [(Int, Int)] -> Int -> Int -> Bool +prop_findWithDefault xs n x = + let xs' = List.nubBy ((==) `on` fst) xs + m = fromList xs' + in all (\k -> findWithDefault x k m == maybe x id (List.lookup k xs')) (n : List.map fst xs') + +test_lookupSomething :: (Int -> IntMap Int -> Maybe (Int, Int)) -> (Int -> Int -> Bool) -> [(Int, Int)] -> Bool +test_lookupSomething lookup' cmp xs = + let odd_sorted_xs = filter_odd $ sort $ List.nubBy ((==) `on` fst) xs + t = fromList odd_sorted_xs + test k = case List.filter ((`cmp` k) . fst) odd_sorted_xs of + [] -> lookup' k t == Nothing + cs | 0 `cmp` 1 -> lookup' k t == Just (last cs) -- we want largest such element + | otherwise -> lookup' k t == Just (head cs) -- we want smallest such element + in all test (List.map fst xs) + + where filter_odd [] = [] + filter_odd [_] = [] + filter_odd (_ : o : xs) = o : filter_odd xs + +prop_lookupLT :: [(Int, Int)] -> Bool +prop_lookupLT = test_lookupSomething lookupLT (<) + +prop_lookupGT :: [(Int, Int)] -> Bool +prop_lookupGT = test_lookupSomething lookupGT (>) + +prop_lookupLE :: [(Int, Int)] -> Bool +prop_lookupLE = test_lookupSomething lookupLE (<=) + +prop_lookupGE :: [(Int, Int)] -> Bool +prop_lookupGE = test_lookupSomething lookupGE (>=) + +prop_findMin :: [(Int, Int)] -> Property +prop_findMin ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in findMin m == List.minimumBy (comparing fst) xs + +prop_findMax :: [(Int, Int)] -> Property +prop_findMax ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in findMax m == List.maximumBy (comparing fst) xs + +prop_deleteMinModel :: [(Int, Int)] -> Property +prop_deleteMinModel ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in toAscList (deleteMin m) == tail (sort xs) + +prop_deleteMaxModel :: [(Int, Int)] -> Property +prop_deleteMaxModel ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in toAscList (deleteMax m) == init (sort xs) + +prop_filter :: (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) + +prop_partition :: (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) + +prop_map :: (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 ] + +prop_fmap :: (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 ] + +prop_mapkeys :: (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]) + +prop_splitModel :: Int -> [(Int, Int)] -> Property +prop_splitModel n ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + (l, r) = split n $ fromList xs + in toAscList l == sort [(k, v) | (k,v) <- xs, k < n] && + toAscList r == sort [(k, v) | (k,v) <- xs, k > n] + +prop_foldr :: Int -> [(Int, Int)] -> Property +prop_foldr n ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in foldr (+) n m == List.foldr (+) n (List.map snd xs) && + foldr (:) [] m == List.map snd (List.sort xs) && + foldrWithKey (\_ a b -> a + b) n m == List.foldr (+) n (List.map snd xs) && + foldrWithKey (\k _ b -> k + b) n m == List.foldr (+) n (List.map fst xs) && + foldrWithKey (\k x xs -> (k,x):xs) [] m == List.sort xs + + +prop_foldr' :: Int -> [(Int, Int)] -> Property +prop_foldr' n ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in foldr' (+) n m == List.foldr (+) n (List.map snd xs) && + foldr' (:) [] m == List.map snd (List.sort xs) && + foldrWithKey' (\_ a b -> a + b) n m == List.foldr (+) n (List.map snd xs) && + foldrWithKey' (\k _ b -> k + b) n m == List.foldr (+) n (List.map fst xs) && + foldrWithKey' (\k x xs -> (k,x):xs) [] m == List.sort xs + +prop_foldl :: Int -> [(Int, Int)] -> Property +prop_foldl n ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in foldl (+) n m == List.foldr (+) n (List.map snd xs) && + foldl (flip (:)) [] m == reverse (List.map snd (List.sort xs)) && + foldlWithKey (\b _ a -> a + b) n m == List.foldr (+) n (List.map snd xs) && + foldlWithKey (\b k _ -> k + b) n m == List.foldr (+) n (List.map fst xs) && + foldlWithKey (\xs k x -> (k,x):xs) [] m == reverse (List.sort xs) + +prop_foldl' :: Int -> [(Int, Int)] -> Property +prop_foldl' n ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in foldl' (+) n m == List.foldr (+) n (List.map snd xs) && + foldl' (flip (:)) [] m == reverse (List.map snd (List.sort xs)) && + foldlWithKey' (\b _ a -> a + b) n m == List.foldr (+) n (List.map snd xs) && + foldlWithKey' (\b k _ -> k + b) n m == List.foldr (+) n (List.map fst xs) && + foldlWithKey' (\xs k x -> (k,x):xs) [] m == reverse (List.sort xs) + +prop_keysSet :: [(Int, Int)] -> Bool +prop_keysSet xs = + keysSet (fromList xs) == Data.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 diff --git a/tests/intset-properties.hs b/tests/intset-properties.hs new file mode 100644 --- /dev/null +++ b/tests/intset-properties.hs @@ -0,0 +1,312 @@ +import Data.Bits ((.&.)) +import Data.IntSet +import Data.List (nub,sort) +import qualified Data.List as List +import Data.Monoid (mempty) +import qualified Data.Set as Set +import Prelude hiding (lookup, null, map, filter, foldr, foldl) +import Test.Framework +import Test.Framework.Providers.HUnit +import Test.Framework.Providers.QuickCheck2 +import Test.HUnit hiding (Test, Testable) +import Test.QuickCheck hiding ((.&.)) + +main :: IO () +main = defaultMainWithOpts [ testCase "lookupLT" test_lookupLT + , testCase "lookupGT" test_lookupGT + , testCase "lookupLE" test_lookupLE + , testCase "lookupGE" test_lookupGE + , testCase "split" test_split + , testProperty "prop_Single" prop_Single + , testProperty "prop_Member" prop_Member + , testProperty "prop_NotMember" prop_NotMember + , testProperty "prop_LookupLT" prop_LookupLT + , testProperty "prop_LookupGT" prop_LookupGT + , testProperty "prop_LookupLE" prop_LookupLE + , testProperty "prop_LookupGE" prop_LookupGE + , testProperty "prop_InsertDelete" prop_InsertDelete + , testProperty "prop_MemberFromList" prop_MemberFromList + , testProperty "prop_UnionInsert" prop_UnionInsert + , testProperty "prop_UnionAssoc" prop_UnionAssoc + , testProperty "prop_UnionComm" prop_UnionComm + , testProperty "prop_Diff" prop_Diff + , testProperty "prop_Int" prop_Int + , testProperty "prop_Ordered" prop_Ordered + , testProperty "prop_List" prop_List + , testProperty "prop_DescList" prop_DescList + , testProperty "prop_AscDescList" prop_AscDescList + , testProperty "prop_fromList" prop_fromList + , testProperty "prop_MaskPow2" prop_MaskPow2 + , testProperty "prop_Prefix" prop_Prefix + , testProperty "prop_LeftRight" prop_LeftRight + , testProperty "prop_isProperSubsetOf" prop_isProperSubsetOf + , testProperty "prop_isProperSubsetOf2" prop_isProperSubsetOf2 + , testProperty "prop_isSubsetOf" prop_isSubsetOf + , testProperty "prop_isSubsetOf2" prop_isSubsetOf2 + , testProperty "prop_size" prop_size + , testProperty "prop_findMax" prop_findMax + , testProperty "prop_findMin" prop_findMin + , testProperty "prop_ord" prop_ord + , testProperty "prop_readShow" prop_readShow + , testProperty "prop_foldR" prop_foldR + , testProperty "prop_foldR'" prop_foldR' + , testProperty "prop_foldL" prop_foldL + , testProperty "prop_foldL'" prop_foldL' + , testProperty "prop_map" prop_map + , testProperty "prop_maxView" prop_maxView + , testProperty "prop_minView" prop_minView + , testProperty "prop_split" prop_split + , testProperty "prop_splitMember" prop_splitMember + , testProperty "prop_partition" prop_partition + , testProperty "prop_filter" prop_filter + ] opts + where + opts = mempty { ropt_test_options = Just $ mempty { topt_maximum_generated_tests = Just 500 + , topt_maximum_unsuitable_generated_tests = Just 500 + } + } + +---------------------------------------------------------------- +-- Unit tests +---------------------------------------------------------------- + +test_lookupLT :: Assertion +test_lookupLT = do + lookupLT 3 (fromList [3, 5]) @?= Nothing + lookupLT 5 (fromList [3, 5]) @?= Just 3 + +test_lookupGT :: Assertion +test_lookupGT = do + lookupGT 4 (fromList [3, 5]) @?= Just 5 + lookupGT 5 (fromList [3, 5]) @?= Nothing + +test_lookupLE :: Assertion +test_lookupLE = do + lookupLE 2 (fromList [3, 5]) @?= Nothing + lookupLE 4 (fromList [3, 5]) @?= Just 3 + lookupLE 5 (fromList [3, 5]) @?= Just 5 + +test_lookupGE :: Assertion +test_lookupGE = do + lookupGE 3 (fromList [3, 5]) @?= Just 3 + lookupGE 4 (fromList [3, 5]) @?= Just 5 + lookupGE 6 (fromList [3, 5]) @?= Nothing + +test_split :: Assertion +test_split = do + split 3 (fromList [1..5]) @?= (fromList [1,2], fromList [4,5]) + +{-------------------------------------------------------------------- + Arbitrary, reasonably balanced trees +--------------------------------------------------------------------} +instance Arbitrary IntSet where + arbitrary = do{ xs <- arbitrary + ; return (fromList xs) + } + + +{-------------------------------------------------------------------- + Single, Member, Insert, Delete, Member, FromList +--------------------------------------------------------------------} +prop_Single :: Int -> Bool +prop_Single x + = (insert x empty == singleton x) + +prop_Member :: [Int] -> Int -> Bool +prop_Member xs n = + let m = fromList xs + in all (\k -> k `member` m == (k `elem` xs)) (n : xs) + +prop_NotMember :: [Int] -> Int -> Bool +prop_NotMember xs n = + let m = fromList xs + in all (\k -> k `notMember` m == (k `notElem` xs)) (n : xs) + +test_LookupSomething :: (Int -> IntSet -> Maybe Int) -> (Int -> Int -> Bool) -> [Int] -> Bool +test_LookupSomething lookup' cmp xs = + let odd_sorted_xs = filter_odd $ nub $ sort xs + t = fromList odd_sorted_xs + test x = case List.filter (`cmp` x) odd_sorted_xs of + [] -> lookup' x t == Nothing + cs | 0 `cmp` 1 -> lookup' x t == Just (last cs) -- we want largest such element + | otherwise -> lookup' x t == Just (head cs) -- we want smallest such element + in all test xs + + where filter_odd [] = [] + filter_odd [_] = [] + filter_odd (_ : o : xs) = o : filter_odd xs + +prop_LookupLT :: [Int] -> Bool +prop_LookupLT = test_LookupSomething lookupLT (<) + +prop_LookupGT :: [Int] -> Bool +prop_LookupGT = test_LookupSomething lookupGT (>) + +prop_LookupLE :: [Int] -> Bool +prop_LookupLE = test_LookupSomething lookupLE (<=) + +prop_LookupGE :: [Int] -> Bool +prop_LookupGE = test_LookupSomething lookupGE (>=) + +prop_InsertDelete :: Int -> IntSet -> Property +prop_InsertDelete k t + = not (member k t) ==> delete k (insert k t) == t + +prop_MemberFromList :: [Int] -> Bool +prop_MemberFromList xs + = all (`member` t) abs_xs && all ((`notMember` t) . negate) abs_xs + where abs_xs = [abs x | x <- xs, x /= 0] + t = fromList abs_xs + +{-------------------------------------------------------------------- + Union +--------------------------------------------------------------------} +prop_UnionInsert :: Int -> IntSet -> Bool +prop_UnionInsert x t + = union t (singleton x) == insert x t + +prop_UnionAssoc :: IntSet -> IntSet -> IntSet -> Bool +prop_UnionAssoc t1 t2 t3 + = union t1 (union t2 t3) == union (union t1 t2) t3 + +prop_UnionComm :: IntSet -> IntSet -> Bool +prop_UnionComm t1 t2 + = (union t1 t2 == union t2 t1) + +prop_Diff :: [Int] -> [Int] -> Bool +prop_Diff xs ys + = toAscList (difference (fromList xs) (fromList ys)) + == List.sort ((List.\\) (nub xs) (nub ys)) + +prop_Int :: [Int] -> [Int] -> Bool +prop_Int xs ys + = toAscList (intersection (fromList xs) (fromList ys)) + == List.sort (nub ((List.intersect) (xs) (ys))) + +{-------------------------------------------------------------------- + Lists +--------------------------------------------------------------------} +prop_Ordered + = forAll (choose (5,100)) $ \n -> + let xs = concat [[i-n,i-n]|i<-[0..2*n :: Int]] + in fromAscList xs == fromList xs + +prop_List :: [Int] -> Bool +prop_List xs + = (sort (nub xs) == toAscList (fromList xs)) + +prop_DescList :: [Int] -> Bool +prop_DescList xs = (reverse (sort (nub xs)) == toDescList (fromList xs)) + +prop_AscDescList :: [Int] -> Bool +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 + nub_sort_xs = List.map List.head $ List.group sort_xs + +{-------------------------------------------------------------------- + Bin invariants +--------------------------------------------------------------------} +powersOf2 :: IntSet +powersOf2 = fromList [2^i | i <- [0..63]] + +-- Check the invariant that the mask is a power of 2. +prop_MaskPow2 :: IntSet -> Bool +prop_MaskPow2 (Bin _ msk left right) = member msk powersOf2 && prop_MaskPow2 left && prop_MaskPow2 right +prop_MaskPow2 _ = True + +-- Check that the prefix satisfies its invariant. +prop_Prefix :: IntSet -> Bool +prop_Prefix s@(Bin prefix msk left right) = all (\elem -> match elem prefix msk) (toList s) && prop_Prefix left && prop_Prefix right +prop_Prefix _ = True + +-- Check that the left elements don't have the mask bit set, and the right +-- ones do. +prop_LeftRight :: IntSet -> Bool +prop_LeftRight (Bin _ msk left right) = and [x .&. msk == 0 | x <- toList left] && and [x .&. msk == msk | x <- toList right] +prop_LeftRight _ = True + +{-------------------------------------------------------------------- + IntSet operations are like Set operations +--------------------------------------------------------------------} +toSet :: IntSet -> Set.Set Int +toSet = Set.fromList . toList + +-- Check that IntSet.isProperSubsetOf is the same as Set.isProperSubsetOf. +prop_isProperSubsetOf :: IntSet -> IntSet -> Bool +prop_isProperSubsetOf a b = isProperSubsetOf a b == Set.isProperSubsetOf (toSet a) (toSet 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 :: IntSet -> IntSet -> Bool +prop_isProperSubsetOf2 a b = isProperSubsetOf a c == (a /= c) where + c = union a b + +prop_isSubsetOf :: IntSet -> IntSet -> Bool +prop_isSubsetOf a b = isSubsetOf a b == Set.isSubsetOf (toSet a) (toSet b) + +prop_isSubsetOf2 :: IntSet -> IntSet -> Bool +prop_isSubsetOf2 a b = isSubsetOf a (union a b) + +prop_size :: IntSet -> Bool +prop_size s = size s == List.length (toList s) + +prop_findMax :: IntSet -> Property +prop_findMax s = not (null s) ==> findMax s == maximum (toList s) + +prop_findMin :: IntSet -> Property +prop_findMin s = not (null s) ==> findMin s == minimum (toList s) + +prop_ord :: IntSet -> IntSet -> Bool +prop_ord s1 s2 = s1 `compare` s2 == toList s1 `compare` toList s2 + +prop_readShow :: IntSet -> Bool +prop_readShow s = s == read (show s) + +prop_foldR :: IntSet -> Bool +prop_foldR s = foldr (:) [] s == toList s + +prop_foldR' :: IntSet -> Bool +prop_foldR' s = foldr' (:) [] s == toList s + +prop_foldL :: IntSet -> Bool +prop_foldL s = foldl (flip (:)) [] s == List.foldl (flip (:)) [] (toList s) + +prop_foldL' :: IntSet -> Bool +prop_foldL' s = foldl' (flip (:)) [] s == List.foldl' (flip (:)) [] (toList s) + +prop_map :: IntSet -> Bool +prop_map s = map id s == s + +prop_maxView :: IntSet -> Bool +prop_maxView s = case maxView s of + Nothing -> null s + Just (m,s') -> m == maximum (toList s) && s == insert m s' && m `notMember` s' + +prop_minView :: IntSet -> Bool +prop_minView s = case minView s of + Nothing -> null s + Just (m,s') -> m == minimum (toList s) && s == insert m s' && m `notMember` s' + +prop_split :: IntSet -> Int -> Bool +prop_split s i = case split i s of + (s1,s2) -> all (i) (toList s2) && i `delete` s == union s1 s2 + +prop_splitMember :: IntSet -> Int -> Bool +prop_splitMember s i = case splitMember i s of + (s1,t,s2) -> all (i) (toList s2) && t == i `member` s && i `delete` s == union s1 s2 + +prop_partition :: IntSet -> Int -> Bool +prop_partition s i = case partition odd s of + (s1,s2) -> all odd (toList s1) && all even (toList s2) && s == s1 `union` s2 + +prop_filter :: IntSet -> Int -> Bool +prop_filter s i = partition odd s == (filter odd s, filter even s) diff --git a/tests/map-properties.hs b/tests/map-properties.hs new file mode 100644 --- /dev/null +++ b/tests/map-properties.hs @@ -0,0 +1,1188 @@ +{-# LANGUAGE CPP #-} + +#ifdef STRICT +import Data.Map.Strict as Data.Map +#else +import Data.Map.Lazy as Data.Map +#endif + +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 Data.List (nub,sort) +import qualified Data.List as List +import qualified Data.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 () + +default (Int) + +main :: IO () +main = defaultMainWithOpts + [ testCase "ticket4242" test_ticket4242 + , testCase "index" test_index + , testCase "size" test_size + , testCase "size2" test_size2 + , testCase "member" test_member + , testCase "notMember" test_notMember + , testCase "lookup" test_lookup + , testCase "findWithDefault" test_findWithDefault + , testCase "lookupLT" test_lookupLT + , testCase "lookupGT" test_lookupGT + , testCase "lookupLE" test_lookupLE + , testCase "lookupGE" test_lookupGE + , testCase "empty" test_empty + , testCase "mempty" test_mempty + , testCase "singleton" test_singleton + , testCase "insert" test_insert + , testCase "insertWith" test_insertWith + , testCase "insertWithKey" test_insertWithKey + , testCase "insertLookupWithKey" test_insertLookupWithKey + , testCase "delete" test_delete + , testCase "adjust" test_adjust + , testCase "adjustWithKey" test_adjustWithKey + , testCase "update" test_update + , testCase "updateWithKey" test_updateWithKey + , testCase "updateLookupWithKey" test_updateLookupWithKey + , testCase "alter" test_alter + , testCase "union" test_union + , testCase "mappend" test_mappend + , testCase "unionWith" test_unionWith + , testCase "unionWithKey" test_unionWithKey + , testCase "unions" test_unions + , testCase "mconcat" test_mconcat + , testCase "unionsWith" test_unionsWith + , testCase "difference" test_difference + , testCase "differenceWith" test_differenceWith + , testCase "differenceWithKey" test_differenceWithKey + , testCase "intersection" test_intersection + , testCase "intersectionWith" test_intersectionWith + , testCase "intersectionWithKey" test_intersectionWithKey + , testCase "map" test_map + , testCase "mapWithKey" test_mapWithKey + , testCase "mapAccum" test_mapAccum + , testCase "mapAccumWithKey" test_mapAccumWithKey + , testCase "mapAccumRWithKey" test_mapAccumRWithKey + , testCase "mapKeys" test_mapKeys + , testCase "mapKeysWith" test_mapKeysWith + , testCase "mapKeysMonotonic" test_mapKeysMonotonic + , testCase "elems" test_elems + , testCase "keys" test_keys + , testCase "assocs" test_assocs + , testCase "keysSet" test_keysSet + , testCase "fromSet" test_fromSet + , testCase "toList" test_toList + , testCase "fromList" test_fromList + , testCase "fromListWith" test_fromListWith + , testCase "fromListWithKey" test_fromListWithKey + , testCase "toAscList" test_toAscList + , testCase "toDescList" test_toDescList + , testCase "showTree" test_showTree + , testCase "showTree'" test_showTree' + , testCase "fromAscList" test_fromAscList + , testCase "fromAscListWith" test_fromAscListWith + , testCase "fromAscListWithKey" test_fromAscListWithKey + , testCase "fromDistinctAscList" test_fromDistinctAscList + , testCase "filter" test_filter + , testCase "filterWithKey" test_filteWithKey + , testCase "partition" test_partition + , testCase "partitionWithKey" test_partitionWithKey + , testCase "mapMaybe" test_mapMaybe + , testCase "mapMaybeWithKey" test_mapMaybeWithKey + , testCase "mapEither" test_mapEither + , testCase "mapEitherWithKey" test_mapEitherWithKey + , testCase "split" test_split + , testCase "splitLookup" test_splitLookup + , testCase "isSubmapOfBy" test_isSubmapOfBy + , testCase "isSubmapOf" test_isSubmapOf + , testCase "isProperSubmapOfBy" test_isProperSubmapOfBy + , testCase "isProperSubmapOf" test_isProperSubmapOf + , testCase "lookupIndex" test_lookupIndex + , testCase "findIndex" test_findIndex + , testCase "elemAt" test_elemAt + , testCase "updateAt" test_updateAt + , testCase "deleteAt" test_deleteAt + , testCase "findMin" test_findMin + , testCase "findMax" test_findMax + , testCase "deleteMin" test_deleteMin + , testCase "deleteMax" test_deleteMax + , testCase "deleteFindMin" test_deleteFindMin + , testCase "deleteFindMax" test_deleteFindMax + , testCase "updateMin" test_updateMin + , testCase "updateMax" test_updateMax + , testCase "updateMinWithKey" test_updateMinWithKey + , testCase "updateMaxWithKey" test_updateMaxWithKey + , testCase "minView" test_minView + , testCase "maxView" test_maxView + , testCase "minViewWithKey" test_minViewWithKey + , testCase "maxViewWithKey" test_maxViewWithKey + , testCase "valid" test_valid + , testProperty "fromList" prop_fromList + , testProperty "insert to singleton" prop_singleton + , testProperty "insert" prop_insert + , testProperty "insert then lookup" prop_insertLookup + , testProperty "insert then delete" prop_insertDelete + , testProperty "insert then delete2" prop_insertDelete2 + , testProperty "delete non member" prop_deleteNonMember + , testProperty "deleteMin" prop_deleteMin + , testProperty "deleteMax" prop_deleteMax + , testProperty "split" prop_split + , testProperty "split then join" prop_join + , testProperty "split then merge" prop_merge + , testProperty "union" prop_union + , testProperty "union model" prop_unionModel + , testProperty "union singleton" prop_unionSingleton + , testProperty "union associative" prop_unionAssoc + , testProperty "union+unionWith" prop_unionWith + , testProperty "unionWith" prop_unionWith2 + , testProperty "union sum" prop_unionSum + , testProperty "difference" prop_difference + , testProperty "difference model" prop_differenceModel + , testProperty "intersection" prop_intersection + , testProperty "intersection model" prop_intersectionModel + , testProperty "intersectionWith" prop_intersectionWith + , testProperty "intersectionWithModel" prop_intersectionWithModel + , testProperty "intersectionWithKey" prop_intersectionWithKey + , testProperty "intersectionWithKeyModel" prop_intersectionWithKeyModel + , testProperty "mergeWithKey model" prop_mergeWithKeyModel + , testProperty "fromAscList" prop_ordered + , testProperty "fromList then toList" prop_list + , testProperty "toDescList" prop_descList + , testProperty "toAscList+toDescList" prop_ascDescList + , testProperty "alter" prop_alter + , testProperty "index" prop_index + , testProperty "null" prop_null + , testProperty "member" prop_member + , testProperty "notmember" prop_notmember + , testProperty "lookup" prop_lookup + , testProperty "find" prop_find + , testProperty "findWithDefault" prop_findWithDefault + , testProperty "lookupLT" prop_lookupLT + , testProperty "lookupGT" prop_lookupGT + , testProperty "lookupLE" prop_lookupLE + , testProperty "lookupGE" prop_lookupGE + , testProperty "findIndex" prop_findIndex + , testProperty "lookupIndex" prop_lookupIndex + , testProperty "findMin" prop_findMin + , testProperty "findMax" prop_findMax + , testProperty "deleteMin" prop_deleteMinModel + , testProperty "deleteMax" prop_deleteMaxModel + , testProperty "filter" prop_filter + , testProperty "partition" prop_partition + , testProperty "map" prop_map + , testProperty "fmap" prop_fmap + , testProperty "mapkeys" prop_mapkeys + , testProperty "split" prop_splitModel + , testProperty "foldr" prop_foldr + , testProperty "foldr'" prop_foldr' + , testProperty "foldl" prop_foldl + , testProperty "foldl'" prop_foldl' + , testProperty "keysSet" prop_keysSet + , testProperty "fromSet" prop_fromSet + ] opts + + where + opts = mempty { ropt_test_options = Just $ mempty { topt_maximum_generated_tests = Just 500 + , topt_maximum_unsuitable_generated_tests = Just 500 + } + } + +{-------------------------------------------------------------------- + Arbitrary, reasonably balanced trees +--------------------------------------------------------------------} +instance (Enum k,Arbitrary a) => Arbitrary (Map k a) where + arbitrary = sized (arbtree 0 maxkey) + where maxkey = 10^5 + + arbtree :: (Enum k, Arbitrary a) => Int -> Int -> Int -> Gen (Map k 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{ x <- arbitrary + ; 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) x l r) + } + +------------------------------------------------------------------------ + +type UMap = Map Int () +type IMap = Map Int Int +type SMap = Map Int String + +---------------------------------------------------------------- +-- Unit tests +---------------------------------------------------------------- + +test_ticket4242 :: Assertion +test_ticket4242 = (valid $ deleteMin $ deleteMin $ fromList [ (i, ()) | i <- [0,2,5,1,6,4,8,9,7,11,10,3] :: [Int] ]) @?= True + +---------------------------------------------------------------- +-- Operators + +test_index :: Assertion +test_index = fromList [(5,'a'), (3,'b')] ! 5 @?= 'a' + +---------------------------------------------------------------- +-- Query + +test_size :: Assertion +test_size = do + null (empty) @?= True + null (singleton 1 'a') @?= False + +test_size2 :: Assertion +test_size2 = do + size empty @?= 0 + size (singleton 1 'a') @?= 1 + size (fromList([(1,'a'), (2,'c'), (3,'b')])) @?= 3 + +test_member :: Assertion +test_member = do + member 5 (fromList [(5,'a'), (3,'b')]) @?= True + member 1 (fromList [(5,'a'), (3,'b')]) @?= False + +test_notMember :: Assertion +test_notMember = do + notMember 5 (fromList [(5,'a'), (3,'b')]) @?= False + notMember 1 (fromList [(5,'a'), (3,'b')]) @?= True + +test_lookup :: Assertion +test_lookup = do + employeeCurrency "John" @?= Just "Euro" + employeeCurrency "Pete" @?= Nothing + where + 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 + +test_findWithDefault :: Assertion +test_findWithDefault = do + findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) @?= 'x' + findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) @?= 'a' + +test_lookupLT :: Assertion +test_lookupLT = do + lookupLT 3 (fromList [(3,'a'), (5,'b')]) @?= Nothing + lookupLT 4 (fromList [(3,'a'), (5,'b')]) @?= Just (3, 'a') + +test_lookupGT :: Assertion +test_lookupGT = do + lookupGT 4 (fromList [(3,'a'), (5,'b')]) @?= Just (5, 'b') + lookupGT 5 (fromList [(3,'a'), (5,'b')]) @?= Nothing + +test_lookupLE :: Assertion +test_lookupLE = do + 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') + +test_lookupGE :: Assertion +test_lookupGE = do + 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 + +---------------------------------------------------------------- +-- Construction + +test_empty :: Assertion +test_empty = do + (empty :: UMap) @?= fromList [] + size empty @?= 0 + +test_mempty :: Assertion +test_mempty = do + (mempty :: UMap) @?= fromList [] + size (mempty :: UMap) @?= 0 + +test_singleton :: Assertion +test_singleton = do + singleton 1 'a' @?= fromList [(1, 'a')] + size (singleton 1 'a') @?= 1 + +test_insert :: Assertion +test_insert = do + 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' + +test_insertWith :: Assertion +test_insertWith = do + 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" + +test_insertWithKey :: Assertion +test_insertWithKey = do + 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" + where + f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value + +test_insertLookupWithKey :: Assertion +test_insertLookupWithKey = do + insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) @?= (Just "a", fromList [(3, "b"), (5, "5:xxx|a")]) + insertLookupWithKey f 2 "xxx" (fromList [(5,"a"), (3,"b")]) @?= (Nothing,fromList [(2,"xxx"),(3,"b"),(5,"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") + where + f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value + +---------------------------------------------------------------- +-- Delete/Update + +test_delete :: Assertion +test_delete = do + 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 :: IMap) + +test_adjust :: Assertion +test_adjust = do + 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 + +test_adjustWithKey :: Assertion +test_adjustWithKey = do + 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 + where + f key x = (show key) ++ ":new " ++ x + +test_update :: Assertion +test_update = do + 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" + where + f x = if x == "a" then Just "new a" else Nothing + +test_updateWithKey :: Assertion +test_updateWithKey = do + 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" + where + f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing + +test_updateLookupWithKey :: Assertion +test_updateLookupWithKey = do + 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") + where + f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing + +test_alter :: Assertion +test_alter = do + alter f 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a")] + alter f 5 (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "b" + alter g 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a"), (7, "c")] + alter g 5 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "c")] + where + f _ = Nothing + g _ = Just "c" + +---------------------------------------------------------------- +-- Combine + +test_union :: Assertion +test_union = union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "a"), (7, "C")] + +test_mappend :: Assertion +test_mappend = mappend (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "a"), (7, "C")] + +test_unionWith :: Assertion +test_unionWith = unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "aA"), (7, "C")] + +test_unionWithKey :: Assertion +test_unionWithKey = unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "5:a|A"), (7, "C")] + where + f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value + +test_unions :: Assertion +test_unions = do + 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")] + +test_mconcat :: Assertion +test_mconcat = do + mconcat [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] + @?= fromList [(3, "b"), (5, "a"), (7, "C")] + mconcat [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])] + @?= fromList [(3, "B3"), (5, "A3"), (7, "C")] + +test_unionsWith :: Assertion +test_unionsWith = unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] + @?= fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] + +test_difference :: Assertion +test_difference = difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 3 "b" + +test_differenceWith :: Assertion +test_differenceWith = differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")]) + @?= singleton 3 "b:B" + where + f al ar = if al== "b" then Just (al ++ ":" ++ ar) else Nothing + +test_differenceWithKey :: Assertion +test_differenceWithKey = differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")]) + @?= singleton 3 "3:b|B" + where + f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing + +test_intersection :: Assertion +test_intersection = intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "a" + + +test_intersectionWith :: Assertion +test_intersectionWith = intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "aA" + +test_intersectionWithKey :: Assertion +test_intersectionWithKey = intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "5:a|A" + where + f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar + +---------------------------------------------------------------- +-- Traversal + +test_map :: Assertion +test_map = map (++ "x") (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "bx"), (5, "ax")] + +test_mapWithKey :: Assertion +test_mapWithKey = mapWithKey f (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "3:b"), (5, "5:a")] + where + f key x = (show key) ++ ":" ++ x + +test_mapAccum :: Assertion +test_mapAccum = mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) @?= ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) + where + f a b = (a ++ b, b ++ "X") + +test_mapAccumWithKey :: Assertion +test_mapAccumWithKey = mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) @?= ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) + where + f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") + +test_mapAccumRWithKey :: Assertion +test_mapAccumRWithKey = mapAccumRWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) @?= ("Everything: 5-a 3-b", fromList [(3, "bX"), (5, "aX")]) + where + f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") + +test_mapKeys :: Assertion +test_mapKeys = do + 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" + +test_mapKeysWith :: Assertion +test_mapKeysWith = do + 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" + +test_mapKeysMonotonic :: Assertion +test_mapKeysMonotonic = do + mapKeysMonotonic (+ 1) (fromList [(5,"a"), (3,"b")]) @?= fromList [(4, "b"), (6, "a")] + 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 + +---------------------------------------------------------------- +-- Conversion + +test_elems :: Assertion +test_elems = do + elems (fromList [(5,"a"), (3,"b")]) @?= ["b","a"] + elems (empty :: UMap) @?= [] + +test_keys :: Assertion +test_keys = do + keys (fromList [(5,"a"), (3,"b")]) @?= [3,5] + keys (empty :: UMap) @?= [] + +test_assocs :: Assertion +test_assocs = do + assocs (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")] + assocs (empty :: UMap) @?= [] + +test_keysSet :: Assertion +test_keysSet = do + keysSet (fromList [(5,"a"), (3,"b")]) @?= Data.Set.fromList [3,5] + keysSet (empty :: UMap) @?= Data.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) + +---------------------------------------------------------------- +-- Lists + +test_toList :: Assertion +test_toList = do + toList (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")] + toList (empty :: SMap) @?= [] + +test_fromList :: Assertion +test_fromList = do + fromList [] @?= (empty :: SMap) + fromList [(5,"a"), (3,"b"), (5, "c")] @?= fromList [(5,"c"), (3,"b")] + fromList [(5,"c"), (3,"b"), (5, "a")] @?= fromList [(5,"a"), (3,"b")] + +test_fromListWith :: Assertion +test_fromListWith = do + fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] @?= fromList [(3, "ab"), (5, "aba")] + fromListWith (++) [] @?= (empty :: SMap) + +test_fromListWithKey :: Assertion +test_fromListWithKey = do + fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] @?= fromList [(3, "3ab"), (5, "5a5ba")] + fromListWithKey f [] @?= (empty :: SMap) + where + f k a1 a2 = (show k) ++ a1 ++ a2 + +---------------------------------------------------------------- +-- Ordered lists + +test_toAscList :: Assertion +test_toAscList = toAscList (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")] + +test_toDescList :: Assertion +test_toDescList = toDescList (fromList [(5,"a"), (3,"b")]) @?= [(5,"a"), (3,"b")] + +test_showTree :: Assertion +test_showTree = + (let t = fromDistinctAscList [(x,()) | x <- [1..5]] + in showTree t) @?= "4:=()\n+--2:=()\n| +--1:=()\n| +--3:=()\n+--5:=()\n" + +test_showTree' :: Assertion +test_showTree' = + (let t = fromDistinctAscList [(x,()) | x <- [1..5]] + in s t ) @?= "+--5:=()\n|\n4:=()\n|\n| +--3:=()\n| |\n+--2:=()\n |\n +--1:=()\n" + where + showElem k x = show k ++ ":=" ++ show x + + s = showTreeWith showElem False True + + +test_fromAscList :: Assertion +test_fromAscList = do + 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 + +test_fromAscListWith :: Assertion +test_fromAscListWith = do + 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 + +test_fromAscListWithKey :: Assertion +test_fromAscListWithKey = do + 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 + where + f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 + +test_fromDistinctAscList :: Assertion +test_fromDistinctAscList = do + 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 + +---------------------------------------------------------------- +-- Filter + +test_filter :: Assertion +test_filter = do + 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 + +test_filteWithKey :: Assertion +test_filteWithKey = filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a" + +test_partition :: Assertion +test_partition = do + 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")]) + +test_partitionWithKey :: Assertion +test_partitionWithKey = do + 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")]) + +test_mapMaybe :: Assertion +test_mapMaybe = mapMaybe f (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "new a" + where + f x = if x == "a" then Just "new a" else Nothing + +test_mapMaybeWithKey :: Assertion +test_mapMaybeWithKey = mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "key : 3" + where + f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing + +test_mapEither :: Assertion +test_mapEither = do + 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 :: SMap), fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) + where + f a = if a < "c" then Left a else Right a + +test_mapEitherWithKey :: Assertion +test_mapEitherWithKey = do + 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 :: SMap), fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")]) + where + f k a = if k < 5 then Left (k * 2) else Right (a ++ a) + +test_split :: Assertion +test_split = do + 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) + +test_splitLookup :: Assertion +test_splitLookup = do + 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) + +---------------------------------------------------------------- +-- Submap + +test_isSubmapOfBy :: Assertion +test_isSubmapOfBy = do + isSubmapOfBy (==) (fromList [('a',1)]) (fromList [('a',1),('b',2)]) @?= True + isSubmapOfBy (<=) (fromList [('a',1)]) (fromList [('a',1),('b',2)]) @?= True + isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1),('b',2)]) @?= True + isSubmapOfBy (==) (fromList [('a',2)]) (fromList [('a',1),('b',2)]) @?= False + isSubmapOfBy (<) (fromList [('a',1)]) (fromList [('a',1),('b',2)]) @?= False + isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1)]) @?= False + +test_isSubmapOf :: Assertion +test_isSubmapOf = do + isSubmapOf (fromList [('a',1)]) (fromList [('a',1),('b',2)]) @?= True + isSubmapOf (fromList [('a',1),('b',2)]) (fromList [('a',1),('b',2)]) @?= True + isSubmapOf (fromList [('a',2)]) (fromList [('a',1),('b',2)]) @?= False + isSubmapOf (fromList [('a',1),('b',2)]) (fromList [('a',1)]) @?= False + +test_isProperSubmapOfBy :: Assertion +test_isProperSubmapOfBy = do + isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= True + isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= True + isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) @?= False + isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)]) @?= False + isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= False + +test_isProperSubmapOf :: Assertion +test_isProperSubmapOf = do + isProperSubmapOf (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= True + isProperSubmapOf (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) @?= False + isProperSubmapOf (fromList [(1,1),(2,2)]) (fromList [(1,1)]) @?= False + +---------------------------------------------------------------- +-- Indexed + +test_lookupIndex :: Assertion +test_lookupIndex = do + 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 + +test_findIndex :: Assertion +test_findIndex = do + findIndex 3 (fromList [(5,"a"), (3,"b")]) @?= 0 + findIndex 5 (fromList [(5,"a"), (3,"b")]) @?= 1 + +test_elemAt :: Assertion +test_elemAt = do + elemAt 0 (fromList [(5,"a"), (3,"b")]) @?= (3,"b") + elemAt 1 (fromList [(5,"a"), (3,"b")]) @?= (5, "a") + +test_updateAt :: Assertion +test_updateAt = do + 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 (\_ _ -> Nothing) 0 (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a" + updateAt (\_ _ -> Nothing) 1 (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "b" +-- updateAt (\_ _ -> Nothing) 7 (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "b" + +test_deleteAt :: Assertion +test_deleteAt = do + deleteAt 0 (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a" + deleteAt 1 (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "b" + +---------------------------------------------------------------- +-- Min/Max + +test_findMin :: Assertion +test_findMin = findMin (fromList [(5,"a"), (3,"b")]) @?= (3,"b") + +test_findMax :: Assertion +test_findMax = findMax (fromList [(5,"a"), (3,"b")]) @?= (5,"a") + +test_deleteMin :: Assertion +test_deleteMin = do + deleteMin (fromList [(5,"a"), (3,"b"), (7,"c")]) @?= fromList [(5,"a"), (7,"c")] + deleteMin (empty :: SMap) @?= empty + +test_deleteMax :: Assertion +test_deleteMax = do + deleteMax (fromList [(5,"a"), (3,"b"), (7,"c")]) @?= fromList [(3,"b"), (5,"a")] + deleteMax (empty :: SMap) @?= empty + +test_deleteFindMin :: Assertion +test_deleteFindMin = deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) @?= ((3,"b"), fromList[(5,"a"), (10,"c")]) + +test_deleteFindMax :: Assertion +test_deleteFindMax = deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) @?= ((10,"c"), fromList [(3,"b"), (5,"a")]) + +test_updateMin :: Assertion +test_updateMin = do + 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" + +test_updateMax :: Assertion +test_updateMax = do + 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" + +test_updateMinWithKey :: Assertion +test_updateMinWithKey = do + 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" + +test_updateMaxWithKey :: Assertion +test_updateMaxWithKey = do + 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" + +test_minView :: Assertion +test_minView = do + minView (fromList [(5,"a"), (3,"b")]) @?= Just ("b", singleton 5 "a") + minView (empty :: SMap) @?= Nothing + +test_maxView :: Assertion +test_maxView = do + maxView (fromList [(5,"a"), (3,"b")]) @?= Just ("a", singleton 3 "b") + maxView (empty :: SMap) @?= Nothing + +test_minViewWithKey :: Assertion +test_minViewWithKey = do + minViewWithKey (fromList [(5,"a"), (3,"b")]) @?= Just ((3,"b"), singleton 5 "a") + minViewWithKey (empty :: SMap) @?= Nothing + +test_maxViewWithKey :: Assertion +test_maxViewWithKey = do + maxViewWithKey (fromList [(5,"a"), (3,"b")]) @?= Just ((5,"a"), singleton 3 "b") + maxViewWithKey (empty :: SMap) @?= Nothing + +---------------------------------------------------------------- +-- Debug + +test_valid :: Assertion +test_valid = do + valid (fromAscList [(3,"b"), (5,"a")]) @?= True + valid (fromAscList [(5,"a"), (3,"b")]) @?= False + +---------------------------------------------------------------- +-- QuickCheck +---------------------------------------------------------------- + +prop_fromList :: UMap -> Bool +prop_fromList t = valid t + +prop_singleton :: Int -> Int -> Bool +prop_singleton k x = insert k x empty == singleton k x + +prop_insert :: Int -> UMap -> Bool +prop_insert k t = valid $ insert k () t + +prop_insertLookup :: Int -> UMap -> Bool +prop_insertLookup k t = lookup k (insert k () t) /= Nothing + +prop_insertDelete :: Int -> UMap -> Bool +prop_insertDelete k t = valid $ delete k (insert k () t) + +prop_insertDelete2 :: Int -> UMap -> Property +prop_insertDelete2 k t = (lookup k t == Nothing) ==> (delete k (insert k () t) == t) + +prop_deleteNonMember :: Int -> UMap -> Property +prop_deleteNonMember k t = (lookup k t == Nothing) ==> (delete k t == t) + +prop_deleteMin :: UMap -> Bool +prop_deleteMin t = valid $ deleteMin $ deleteMin t + +prop_deleteMax :: UMap -> Bool +prop_deleteMax t = valid $ deleteMax $ deleteMax t + +---------------------------------------------------------------- + +prop_split :: Int -> UMap -> Bool +prop_split k t = let (r,l) = split k t + in (valid r, valid l) == (True, True) + +prop_join :: Int -> UMap -> Bool +prop_join k t = let (l,r) = split k t + in valid (join k () l r) + +prop_merge :: Int -> UMap -> Bool +prop_merge k t = let (l,r) = split k t + in valid (merge l r) + +---------------------------------------------------------------- + +prop_union :: UMap -> UMap -> Bool +prop_union t1 t2 = valid (union t1 t2) + +prop_unionModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_unionModel xs ys + = sort (keys (union (fromList xs) (fromList ys))) + == sort (nub (Prelude.map fst xs ++ Prelude.map fst ys)) + +prop_unionSingleton :: IMap -> Int -> Int -> Bool +prop_unionSingleton t k x = union (singleton k x) t == insert k x t + +prop_unionAssoc :: IMap -> IMap -> IMap -> Bool +prop_unionAssoc t1 t2 t3 = union t1 (union t2 t3) == union (union t1 t2) t3 + +prop_unionWith :: IMap -> IMap -> Bool +prop_unionWith t1 t2 = (union t1 t2 == unionWith (\_ y -> y) t2 t1) + +prop_unionWith2 :: IMap -> IMap -> Bool +prop_unionWith2 t1 t2 = valid (unionWithKey (\_ x y -> x+y) t1 t2) + +prop_unionSum :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_unionSum xs ys + = sum (elems (unionWith (+) (fromListWith (+) xs) (fromListWith (+) ys))) + == (sum (Prelude.map snd xs) + sum (Prelude.map snd ys)) + +prop_difference :: IMap -> IMap -> Bool +prop_difference t1 t2 = valid (difference t1 t2) + +prop_differenceModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_differenceModel xs ys + = sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys))) + == sort ((List.\\) (nub (Prelude.map fst xs)) (nub (Prelude.map fst ys))) + +prop_intersection :: IMap -> IMap -> Bool +prop_intersection t1 t2 = valid (intersection t1 t2) + +prop_intersectionModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_intersectionModel xs ys + = 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_intersectionWithModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_intersectionWithModel xs ys + = toList (intersectionWith f (fromList xs') (fromList ys')) + == [(kx, f vx vy) | (kx, vx) <- List.sort xs', (ky, vy) <- ys', kx == ky] + where xs' = List.nubBy ((==) `on` fst) xs + 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_intersectionWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_intersectionWithKeyModel xs ys + = toList (intersectionWithKey f (fromList xs') (fromList ys')) + == [(kx, f kx vx vy) | (kx, vx) <- List.sort xs', (ky, vy) <- ys', kx == ky] + where xs' = List.nubBy ((==) `on` fst) xs + ys' = List.nubBy ((==) `on` fst) ys + f k l r = k + 2 * l + 3 * r + +prop_mergeWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_mergeWithKeyModel xs ys + = and [ testMergeWithKey f keep_x keep_y + | f <- [ \_k x1 _x2 -> Just x1 + , \_k _x1 x2 -> Just x2 + , \_k _x1 _x2 -> Nothing + , \k x1 x2 -> if k `mod` 2 == 0 then Nothing else Just (2 * x1 + 3 * x2) + ] + , keep_x <- [ True, False ] + , keep_y <- [ True, False ] + ] + + where xs' = List.nubBy ((==) `on` fst) xs + ys' = List.nubBy ((==) `on` fst) ys + + xm = fromList xs' + ym = fromList ys' + + testMergeWithKey f keep_x keep_y + = toList (mergeWithKey f (keep keep_x) (keep keep_y) xm ym) == emulateMergeWithKey f keep_x keep_y + where keep False _ = empty + keep True m = m + + emulateMergeWithKey f keep_x keep_y + = Maybe.mapMaybe combine (sort $ List.union (List.map fst xs') (List.map fst ys')) + where combine k = case (List.lookup k xs', List.lookup k ys') of + (Nothing, Just y) -> if keep_y then Just (k, y) else Nothing + (Just x, Nothing) -> if keep_x then Just (k, x) else Nothing + (Just x, Just y) -> (\v -> (k, v)) `fmap` f k x y + + -- We prevent inlining testMergeWithKey to disable the SpecConstr + -- optimalization. There are too many call patterns here so several + -- warnings are issued if testMergeWithKey gets inlined. + {-# NOINLINE testMergeWithKey #-} + +---------------------------------------------------------------- + +prop_ordered :: Property +prop_ordered + = forAll (choose (5,100)) $ \n -> + let xs = [(x,()) | x <- [0..n::Int]] + in fromAscList 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_ascDescList :: [Int] -> Bool +prop_ascDescList xs = toAscList m == reverse (toDescList m) + where m = fromList $ zip xs $ repeat () + +---------------------------------------------------------------- + +prop_alter :: UMap -> Int -> Bool +prop_alter t k = balanced t' && case lookup k t of + Just _ -> (size t - 1) == size t' && lookup k t' == Nothing + Nothing -> (size t + 1) == size t' && lookup k t' /= Nothing + where + t' = alter f k t + f Nothing = Just () + f (Just ()) = Nothing + +------------------------------------------------------------------------ +-- Compare against the list model (after nub on keys) + +prop_index :: [Int] -> Property +prop_index xs = length xs > 0 ==> + let m = fromList (zip xs xs) + in xs == [ m ! i | i <- xs ] + +prop_null :: IMap -> Bool +prop_null m = null m == (size m == 0) + +prop_member :: [Int] -> Int -> Bool +prop_member xs n = + let m = fromList (zip xs xs) + in all (\k -> k `member` m == (k `elem` xs)) (n : xs) + +prop_notmember :: [Int] -> Int -> Bool +prop_notmember xs n = + let m = fromList (zip xs xs) + in all (\k -> k `notMember` m == (k `notElem` xs)) (n : xs) + +prop_lookup :: [(Int, Int)] -> Int -> Bool +prop_lookup xs n = + let xs' = List.nubBy ((==) `on` fst) xs + m = fromList xs' + in all (\k -> lookup k m == List.lookup k xs') (n : List.map fst xs') + +prop_find :: [(Int, Int)] -> Bool +prop_find xs = + let xs' = List.nubBy ((==) `on` fst) xs + m = fromList xs' + in all (\(k, v) -> m ! k == v) xs' + +prop_findWithDefault :: [(Int, Int)] -> Int -> Int -> Bool +prop_findWithDefault xs n x = + let xs' = List.nubBy ((==) `on` fst) xs + m = fromList xs' + in all (\k -> findWithDefault x k m == maybe x id (List.lookup k xs')) (n : List.map fst xs') + +test_lookupSomething :: (Int -> Map Int Int -> Maybe (Int, Int)) -> (Int -> Int -> Bool) -> [(Int, Int)] -> Bool +test_lookupSomething lookup' cmp xs = + let odd_sorted_xs = filter_odd $ sort $ List.nubBy ((==) `on` fst) xs + t = fromList odd_sorted_xs + test k = case List.filter ((`cmp` k) . fst) odd_sorted_xs of + [] -> lookup' k t == Nothing + cs | 0 `cmp` 1 -> lookup' k t == Just (last cs) -- we want largest such element + | otherwise -> lookup' k t == Just (head cs) -- we want smallest such element + in all test (List.map fst xs) + + where filter_odd [] = [] + filter_odd [_] = [] + filter_odd (_ : o : xs) = o : filter_odd xs + +prop_lookupLT :: [(Int, Int)] -> Bool +prop_lookupLT = test_lookupSomething lookupLT (<) + +prop_lookupGT :: [(Int, Int)] -> Bool +prop_lookupGT = test_lookupSomething lookupGT (>) + +prop_lookupLE :: [(Int, Int)] -> Bool +prop_lookupLE = test_lookupSomething lookupLE (<=) + +prop_lookupGE :: [(Int, Int)] -> Bool +prop_lookupGE = test_lookupSomething lookupGE (>=) + +prop_findIndex :: [(Int, Int)] -> Property +prop_findIndex ys = length ys > 0 ==> + let m = fromList ys + in findIndex (fst (head ys)) m `seq` True + +prop_lookupIndex :: [(Int, Int)] -> Property +prop_lookupIndex ys = length ys > 0 ==> + let m = fromList ys + in isJust (lookupIndex (fst (head ys)) m) + +prop_findMin :: [(Int, Int)] -> Property +prop_findMin ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in findMin m == List.minimumBy (comparing fst) xs + +prop_findMax :: [(Int, Int)] -> Property +prop_findMax ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in findMax m == List.maximumBy (comparing fst) xs + +prop_deleteMinModel :: [(Int, Int)] -> Property +prop_deleteMinModel ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in toAscList (deleteMin m) == tail (sort xs) + +prop_deleteMaxModel :: [(Int, Int)] -> Property +prop_deleteMaxModel ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in toAscList (deleteMax m) == init (sort xs) + +prop_filter :: (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) + +prop_partition :: (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) + +prop_map :: (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 ] + +prop_fmap :: (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 ] + +prop_mapkeys :: (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]) + +prop_splitModel :: Int -> [(Int, Int)] -> Property +prop_splitModel n ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + (l, r) = split n $ fromList xs + in toAscList l == sort [(k, v) | (k,v) <- xs, k < n] && + toAscList r == sort [(k, v) | (k,v) <- xs, k > n] + +prop_foldr :: Int -> [(Int, Int)] -> Property +prop_foldr n ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in foldr (+) n m == List.foldr (+) n (List.map snd xs) && + foldr (:) [] m == List.map snd (List.sort xs) && + foldrWithKey (\_ a b -> a + b) n m == List.foldr (+) n (List.map snd xs) && + foldrWithKey (\k _ b -> k + b) n m == List.foldr (+) n (List.map fst xs) && + foldrWithKey (\k x xs -> (k,x):xs) [] m == List.sort xs + + +prop_foldr' :: Int -> [(Int, Int)] -> Property +prop_foldr' n ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in foldr' (+) n m == List.foldr (+) n (List.map snd xs) && + foldr' (:) [] m == List.map snd (List.sort xs) && + foldrWithKey' (\_ a b -> a + b) n m == List.foldr (+) n (List.map snd xs) && + foldrWithKey' (\k _ b -> k + b) n m == List.foldr (+) n (List.map fst xs) && + foldrWithKey' (\k x xs -> (k,x):xs) [] m == List.sort xs + +prop_foldl :: Int -> [(Int, Int)] -> Property +prop_foldl n ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in foldl (+) n m == List.foldr (+) n (List.map snd xs) && + foldl (flip (:)) [] m == reverse (List.map snd (List.sort xs)) && + foldlWithKey (\b _ a -> a + b) n m == List.foldr (+) n (List.map snd xs) && + foldlWithKey (\b k _ -> k + b) n m == List.foldr (+) n (List.map fst xs) && + foldlWithKey (\xs k x -> (k,x):xs) [] m == reverse (List.sort xs) + +prop_foldl' :: Int -> [(Int, Int)] -> Property +prop_foldl' n ys = length ys > 0 ==> + let xs = List.nubBy ((==) `on` fst) ys + m = fromList xs + in foldl' (+) n m == List.foldr (+) n (List.map snd xs) && + foldl' (flip (:)) [] m == reverse (List.map snd (List.sort xs)) && + foldlWithKey' (\b _ a -> a + b) n m == List.foldr (+) n (List.map snd xs) && + foldlWithKey' (\b k _ -> k + b) n m == List.foldr (+) n (List.map fst xs) && + foldlWithKey' (\xs k x -> (k,x):xs) [] m == reverse (List.sort xs) + +prop_keysSet :: [(Int, Int)] -> Bool +prop_keysSet xs = + keysSet (fromList xs) == Data.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 diff --git a/tests/seq-properties.hs b/tests/seq-properties.hs new file mode 100644 --- /dev/null +++ b/tests/seq-properties.hs @@ -0,0 +1,598 @@ +import Data.Sequence -- needs to be compiled with -DTESTING for use here + +import Control.Applicative (Applicative(..)) +import Control.Arrow ((***)) +import Data.Foldable (Foldable(..), toList, all, sum) +import Data.Functor ((<$>), (<$)) +import Data.Maybe +import Data.Monoid (Monoid(..)) +import Data.Traversable (Traversable(traverse), sequenceA) +import Prelude hiding ( + null, length, take, drop, splitAt, + foldl, foldl1, foldr, foldr1, scanl, scanl1, scanr, scanr1, + filter, reverse, replicate, zip, zipWith, zip3, zipWith3, + all, sum) +import qualified Prelude +import qualified Data.List +import Test.QuickCheck hiding ((><)) +import Test.QuickCheck.Poly +import Test.Framework +import Test.Framework.Providers.QuickCheck2 + + +main :: IO () +main = defaultMainWithOpts + [ testProperty "fmap" prop_fmap + , testProperty "(<$)" prop_constmap + , testProperty "foldr" prop_foldr + , testProperty "foldr1" prop_foldr1 + , testProperty "foldl" prop_foldl + , testProperty "foldl1" prop_foldl1 + , testProperty "(==)" prop_equals + , testProperty "compare" prop_compare + , testProperty "mappend" prop_mappend + , testProperty "singleton" prop_singleton + , testProperty "(<|)" prop_cons + , testProperty "(|>)" prop_snoc + , testProperty "(><)" prop_append + , testProperty "fromList" prop_fromList + , testProperty "replicate" prop_replicate + , testProperty "replicateA" prop_replicateA + , testProperty "replicateM" prop_replicateM + , testProperty "iterateN" prop_iterateN + , testProperty "unfoldr" prop_unfoldr + , testProperty "unfoldl" prop_unfoldl + , testProperty "null" prop_null + , testProperty "length" prop_length + , testProperty "viewl" prop_viewl + , testProperty "viewr" prop_viewr + , testProperty "scanl" prop_scanl + , testProperty "scanl1" prop_scanl1 + , testProperty "scanr" prop_scanr + , testProperty "scanr1" prop_scanr1 + , testProperty "tails" prop_tails + , testProperty "inits" prop_inits + , testProperty "takeWhileL" prop_takeWhileL + , testProperty "takeWhileR" prop_takeWhileR + , testProperty "dropWhileL" prop_dropWhileL + , testProperty "dropWhileR" prop_dropWhileR + , testProperty "spanl" prop_spanl + , testProperty "spanr" prop_spanr + , testProperty "breakl" prop_breakl + , testProperty "breakr" prop_breakr + , testProperty "partition" prop_partition + , testProperty "filter" prop_filter + , testProperty "sort" prop_sort + , testProperty "sortBy" prop_sortBy + , testProperty "unstableSort" prop_unstableSort + , testProperty "unstableSortBy" prop_unstableSortBy + , testProperty "index" prop_index + , testProperty "adjust" prop_adjust + , testProperty "update" prop_update + , testProperty "take" prop_take + , testProperty "drop" prop_drop + , testProperty "splitAt" prop_splitAt + , testProperty "elemIndexL" prop_elemIndexL + , testProperty "elemIndicesL" prop_elemIndicesL + , testProperty "elemIndexR" prop_elemIndexR + , testProperty "elemIndicesR" prop_elemIndicesR + , testProperty "findIndexL" prop_findIndexL + , testProperty "findIndicesL" prop_findIndicesL + , testProperty "findIndexR" prop_findIndexR + , testProperty "findIndicesR" prop_findIndicesR + , testProperty "foldlWithIndex" prop_foldlWithIndex + , testProperty "foldrWithIndex" prop_foldrWithIndex + , testProperty "mapWithIndex" prop_mapWithIndex + , testProperty "reverse" prop_reverse + , testProperty "zip" prop_zip + , testProperty "zipWith" prop_zipWith + , testProperty "zip3" prop_zip3 + , testProperty "zipWith3" prop_zipWith3 + , testProperty "zip4" prop_zip4 + , testProperty "zipWith4" prop_zipWith4 + ] opts + + where + opts = mempty { ropt_test_options = Just $ mempty { topt_maximum_generated_tests = Just 500 + , topt_maximum_unsuitable_generated_tests = Just 500 + } + } + +------------------------------------------------------------------------ +-- Arbitrary +------------------------------------------------------------------------ + +instance Arbitrary a => Arbitrary (Seq a) where + arbitrary = Seq <$> arbitrary + shrink (Seq x) = map Seq (shrink x) + +instance Arbitrary a => Arbitrary (Elem a) where + arbitrary = Elem <$> arbitrary + +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 1 = Single <$> arbitrary + arb n = deep <$> arbitrary <*> arb (n `div` 2) <*> arbitrary + + shrink (Deep _ (One a) Empty (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 = [] + +instance (Arbitrary a, Sized a) => Arbitrary (Node a) where + arbitrary = oneof [ + node2 <$> arbitrary <*> arbitrary, + node3 <$> arbitrary <*> arbitrary <*> arbitrary] + + shrink (Node2 _ a b) = + [node2 a' b | a' <- shrink a] ++ + [node2 a b' | b' <- shrink b] + shrink (Node3 _ a b c) = + [node2 a b, node2 a c, node2 b c] ++ + [node3 a' b c | a' <- shrink a] ++ + [node3 a b' c | b' <- shrink b] ++ + [node3 a b c' | c' <- shrink c] + +instance Arbitrary a => Arbitrary (Digit a) where + arbitrary = oneof [ + One <$> arbitrary, + Two <$> arbitrary <*> arbitrary, + Three <$> arbitrary <*> arbitrary <*> arbitrary, + Four <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary] + + shrink (One a) = map One (shrink a) + shrink (Two a b) = [One a, One b] + shrink (Three a b c) = [Two a b, Two a c, Two b c] + shrink (Four a b c d) = [Three a b c, Three a b d, Three a c d, Three b c d] + +------------------------------------------------------------------------ +-- Valid trees +------------------------------------------------------------------------ + +class Valid a where + valid :: a -> Bool + +instance Valid (Elem a) where + valid _ = True + +instance Valid (Seq a) where + valid (Seq xs) = valid xs + +instance (Sized a, Valid a) => Valid (FingerTree a) where + valid Empty = 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 + +instance (Sized a, Valid a) => Valid (Node a) where + valid node = size node == sum (fmap size node) && all valid node + +instance Valid a => Valid (Digit a) where + valid = all valid + +{-------------------------------------------------------------------- + The general plan is to compare each function with a list equivalent. + Each operation should produce a valid tree representing the same + sequence as produced by its list counterpart on corresponding inputs. + (The list versions are often lazier, but these properties ignore + strictness.) +--------------------------------------------------------------------} + +-- utilities for partial conversions + +infix 4 ~= + +(~=) :: Eq a => Maybe a -> a -> Bool +(~=) = maybe (const False) (==) + +-- Partial conversion of an output sequence to a list. +toList' :: Seq a -> Maybe [a] +toList' xs + | valid xs = Just (toList xs) + | otherwise = Nothing + +toListList' :: Seq (Seq a) -> Maybe [[a]] +toListList' xss = toList' xss >>= mapM toList' + +toListPair' :: (Seq a, Seq b) -> Maybe ([a], [b]) +toListPair' (xs, ys) = (,) <$> toList' xs <*> toList' ys + +-- instances + +prop_fmap :: Seq Int -> Bool +prop_fmap xs = + toList' (fmap f xs) ~= map f (toList xs) + where f = (+100) + +prop_constmap :: A -> Seq A -> Bool +prop_constmap x xs = + toList' (x <$ xs) ~= map (const x) (toList xs) + +prop_foldr :: Seq A -> Bool +prop_foldr xs = + foldr f z xs == Prelude.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 xs = + foldl f z xs == Prelude.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) + where f = (-) + +prop_equals :: Seq OrdA -> Seq OrdA -> Bool +prop_equals xs ys = + (xs == ys) == (toList xs == toList ys) + +prop_compare :: Seq OrdA -> Seq OrdA -> Bool +prop_compare xs ys = + compare xs ys == compare (toList xs) (toList ys) + +prop_mappend :: Seq A -> Seq A -> Bool +prop_mappend xs ys = + toList' (mappend xs ys) ~= toList xs ++ toList ys + +-- * Construction + +{- + toList' empty ~= [] +-} + +prop_singleton :: A -> Bool +prop_singleton x = + toList' (singleton x) ~= [x] + +prop_cons :: A -> Seq A -> Bool +prop_cons x xs = + toList' (x <| xs) ~= x : toList xs + +prop_snoc :: Seq A -> A -> Bool +prop_snoc xs x = + toList' (xs |> x) ~= toList xs ++ [x] + +prop_append :: Seq A -> Seq A -> Bool +prop_append xs ys = + toList' (xs >< ys) ~= toList xs ++ toList ys + +prop_fromList :: [A] -> Bool +prop_fromList xs = + toList' (fromList xs) ~= xs + +-- ** Repetition + +prop_replicate :: NonNegative Int -> A -> Bool +prop_replicate (NonNegative m) x = + toList' (replicate n x) ~= Prelude.replicate n x + where n = m `mod` 10000 + +prop_replicateA :: NonNegative Int -> Bool +prop_replicateA (NonNegative m) = + traverse toList' (replicateA n a) ~= sequenceA (Prelude.replicate n a) + where + n = m `mod` 10000 + a = Action 1 0 :: M Int + +prop_replicateM :: NonNegative Int -> Bool +prop_replicateM (NonNegative m) = + traverse toList' (replicateM n a) ~= sequence (Prelude.replicate n a) + where + n = m `mod` 10000 + a = Action 1 0 :: M Int + +-- ** Iterative construction + +prop_iterateN :: NonNegative Int -> Int -> Bool +prop_iterateN (NonNegative m) x = + toList' (iterateN n f x) ~= Prelude.take n (Prelude.iterate f x) + where + n = m `mod` 10000 + f = (+1) + +prop_unfoldr :: [A] -> Bool +prop_unfoldr z = + toList' (unfoldr f z) ~= Data.List.unfoldr f z + where + f [] = Nothing + f (x:xs) = Just (x, xs) + +prop_unfoldl :: [A] -> Bool +prop_unfoldl z = + toList' (unfoldl f z) ~= Data.List.reverse (Data.List.unfoldr (fmap swap . f) z) + where + f [] = Nothing + f (x:xs) = Just (xs, x) + swap (x,y) = (y,x) + +-- * Deconstruction + +-- ** Queries + +prop_null :: Seq A -> Bool +prop_null xs = + null xs == Prelude.null (toList xs) + +prop_length :: Seq A -> Bool +prop_length xs = + length xs == Prelude.length (toList xs) + +-- ** Views + +prop_viewl :: Seq A -> Bool +prop_viewl xs = + case viewl xs of + EmptyL -> Prelude.null (toList xs) + x :< xs' -> valid xs' && toList xs == x : toList xs' + +prop_viewr :: Seq A -> Bool +prop_viewr xs = + case viewr xs of + EmptyR -> Prelude.null (toList xs) + xs' :> x -> valid xs' && toList xs == toList xs' ++ [x] + +-- * Scans + +prop_scanl :: [A] -> Seq A -> Bool +prop_scanl z xs = + toList' (scanl f z xs) ~= Data.List.scanl f z (toList xs) + where f = flip (:) + +prop_scanl1 :: Seq Int -> Property +prop_scanl1 xs = + not (null xs) ==> toList' (scanl1 f xs) ~= Data.List.scanl1 f (toList xs) + where f = (-) + +prop_scanr :: [A] -> Seq A -> Bool +prop_scanr z xs = + toList' (scanr f z xs) ~= Data.List.scanr f z (toList xs) + where f = (:) + +prop_scanr1 :: Seq Int -> Property +prop_scanr1 xs = + not (null xs) ==> toList' (scanr1 f xs) ~= Data.List.scanr1 f (toList xs) + where f = (-) + +-- * Sublists + +prop_tails :: Seq A -> Bool +prop_tails xs = + toListList' (tails xs) ~= Data.List.tails (toList xs) + +prop_inits :: Seq A -> Bool +prop_inits xs = + toListList' (inits xs) ~= Data.List.inits (toList xs) + +-- ** Sequential searches +-- We use predicates with varying density. + +prop_takeWhileL :: Positive Int -> Seq Int -> Bool +prop_takeWhileL (Positive n) xs = + toList' (takeWhileL p xs) ~= Prelude.takeWhile p (toList xs) + where p x = x `mod` n == 0 + +prop_takeWhileR :: Positive Int -> Seq Int -> Bool +prop_takeWhileR (Positive n) xs = + toList' (takeWhileR p xs) ~= Prelude.reverse (Prelude.takeWhile p (Prelude.reverse (toList xs))) + where p x = x `mod` n == 0 + +prop_dropWhileL :: Positive Int -> Seq Int -> Bool +prop_dropWhileL (Positive n) xs = + toList' (dropWhileL p xs) ~= Prelude.dropWhile p (toList xs) + where p x = x `mod` n == 0 + +prop_dropWhileR :: Positive Int -> Seq Int -> Bool +prop_dropWhileR (Positive n) xs = + toList' (dropWhileR p xs) ~= Prelude.reverse (Prelude.dropWhile p (Prelude.reverse (toList xs))) + where p x = x `mod` n == 0 + +prop_spanl :: Positive Int -> Seq Int -> Bool +prop_spanl (Positive n) xs = + toListPair' (spanl p xs) ~= Data.List.span p (toList xs) + where p x = x `mod` n == 0 + +prop_spanr :: Positive Int -> Seq Int -> Bool +prop_spanr (Positive n) xs = + toListPair' (spanr p xs) ~= (Prelude.reverse *** Prelude.reverse) (Data.List.span p (Prelude.reverse (toList xs))) + where p x = x `mod` n == 0 + +prop_breakl :: Positive Int -> Seq Int -> Bool +prop_breakl (Positive n) xs = + toListPair' (breakl p xs) ~= Data.List.break p (toList xs) + where p x = x `mod` n == 0 + +prop_breakr :: Positive Int -> Seq Int -> Bool +prop_breakr (Positive n) xs = + toListPair' (breakr p xs) ~= (Prelude.reverse *** Prelude.reverse) (Data.List.break p (Prelude.reverse (toList xs))) + where p x = x `mod` n == 0 + +prop_partition :: Positive Int -> Seq Int -> Bool +prop_partition (Positive n) xs = + toListPair' (partition p xs) ~= Data.List.partition p (toList xs) + where p x = x `mod` n == 0 + +prop_filter :: Positive Int -> Seq Int -> Bool +prop_filter (Positive n) xs = + toList' (filter p xs) ~= Prelude.filter p (toList xs) + where p x = x `mod` n == 0 + +-- * Sorting + +prop_sort :: Seq OrdA -> Bool +prop_sort xs = + toList' (sort xs) ~= Data.List.sort (toList xs) + +prop_sortBy :: Seq (OrdA, B) -> Bool +prop_sortBy xs = + toList' (sortBy f xs) ~= Data.List.sortBy f (toList xs) + where f (x1, _) (x2, _) = compare x1 x2 + +prop_unstableSort :: Seq OrdA -> Bool +prop_unstableSort xs = + toList' (unstableSort xs) ~= Data.List.sort (toList xs) + +prop_unstableSortBy :: Seq OrdA -> Bool +prop_unstableSortBy xs = + toList' (unstableSortBy compare xs) ~= Data.List.sort (toList xs) + +-- * Indexing + +prop_index :: Seq A -> Property +prop_index xs = + not (null xs) ==> forAll (choose (0, length xs-1)) $ \ i -> + index xs i == toList xs !! i + +prop_adjust :: Int -> Int -> Seq Int -> Bool +prop_adjust n i xs = + toList' (adjust f i xs) ~= adjustList f i (toList xs) + where f = (+n) + +prop_update :: Int -> A -> Seq A -> Bool +prop_update i x xs = + toList' (update i x xs) ~= adjustList (const x) i (toList xs) + +prop_take :: Int -> Seq A -> Bool +prop_take n xs = + toList' (take n xs) ~= Prelude.take n (toList xs) + +prop_drop :: Int -> Seq A -> Bool +prop_drop n xs = + toList' (drop n xs) ~= Prelude.drop n (toList xs) + +prop_splitAt :: Int -> Seq A -> Bool +prop_splitAt n xs = + toListPair' (splitAt n xs) ~= Prelude.splitAt n (toList 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] + +-- ** Indexing with predicates +-- The elem* tests have poor coverage, but for find* we use predicates +-- of varying density. + +prop_elemIndexL :: A -> Seq A -> Bool +prop_elemIndexL x xs = + elemIndexL x xs == Data.List.elemIndex x (toList xs) + +prop_elemIndicesL :: A -> Seq A -> Bool +prop_elemIndicesL x xs = + elemIndicesL x xs == Data.List.elemIndices x (toList xs) + +prop_elemIndexR :: A -> Seq A -> Bool +prop_elemIndexR x xs = + elemIndexR x xs == listToMaybe (Prelude.reverse (Data.List.elemIndices x (toList xs))) + +prop_elemIndicesR :: A -> Seq A -> Bool +prop_elemIndicesR x xs = + elemIndicesR x xs == Prelude.reverse (Data.List.elemIndices x (toList xs)) + +prop_findIndexL :: Positive Int -> Seq Int -> Bool +prop_findIndexL (Positive n) xs = + findIndexL p xs == Data.List.findIndex p (toList xs) + where p x = x `mod` n == 0 + +prop_findIndicesL :: Positive Int -> Seq Int -> Bool +prop_findIndicesL (Positive n) xs = + findIndicesL p xs == Data.List.findIndices p (toList xs) + where p x = x `mod` n == 0 + +prop_findIndexR :: Positive Int -> Seq Int -> Bool +prop_findIndexR (Positive n) xs = + findIndexR p xs == listToMaybe (Prelude.reverse (Data.List.findIndices p (toList xs))) + where p x = x `mod` n == 0 + +prop_findIndicesR :: Positive Int -> Seq Int -> Bool +prop_findIndicesR (Positive n) xs = + findIndicesR p xs == Prelude.reverse (Data.List.findIndices p (toList xs)) + where p x = x `mod` n == 0 + +-- * Folds + +prop_foldlWithIndex :: [(Int, A)] -> Seq A -> Bool +prop_foldlWithIndex z xs = + foldlWithIndex f z xs == Data.List.foldl (uncurry . f) z (Data.List.zip [0..] (toList xs)) + where f ys n y = (n,y):ys + +prop_foldrWithIndex :: [(Int, A)] -> Seq A -> Bool +prop_foldrWithIndex z xs = + foldrWithIndex f z xs == Data.List.foldr (uncurry f) z (Data.List.zip [0..] (toList xs)) + where f n y ys = (n,y):ys + +-- * Transformations + +prop_mapWithIndex :: Seq A -> Bool +prop_mapWithIndex xs = + toList' (mapWithIndex f xs) ~= map (uncurry f) (Data.List.zip [0..] (toList xs)) + where f = (,) + +prop_reverse :: Seq A -> Bool +prop_reverse xs = + toList' (reverse xs) ~= Prelude.reverse (toList xs) + +-- ** Zips + +prop_zip :: Seq A -> Seq B -> Bool +prop_zip xs ys = + toList' (zip xs ys) ~= Prelude.zip (toList xs) (toList ys) + +prop_zipWith :: Seq A -> Seq B -> Bool +prop_zipWith xs ys = + toList' (zipWith f xs ys) ~= Prelude.zipWith f (toList xs) (toList ys) + where f = (,) + +prop_zip3 :: Seq A -> Seq B -> Seq C -> Bool +prop_zip3 xs ys zs = + toList' (zip3 xs ys zs) ~= Prelude.zip3 (toList xs) (toList ys) (toList zs) + +prop_zipWith3 :: Seq A -> Seq B -> Seq C -> Bool +prop_zipWith3 xs ys zs = + toList' (zipWith3 f xs ys zs) ~= Prelude.zipWith3 f (toList xs) (toList ys) (toList zs) + where f = (,,) + +prop_zip4 :: Seq A -> Seq B -> Seq C -> Seq Int -> Bool +prop_zip4 xs ys zs ts = + toList' (zip4 xs ys zs ts) ~= Data.List.zip4 (toList xs) (toList ys) (toList zs) (toList ts) + +prop_zipWith4 :: Seq A -> Seq B -> Seq C -> Seq Int -> Bool +prop_zipWith4 xs ys zs ts = + toList' (zipWith4 f xs ys zs ts) ~= Data.List.zipWith4 f (toList xs) (toList ys) (toList zs) (toList ts) + where f = (,,,) + +-- Simple test monad + +data M a = Action Int a + deriving (Eq, Show) + +instance Functor M where + fmap f (Action n x) = Action n (f x) + +instance Applicative M where + pure x = Action 0 x + Action m f <*> Action n x = Action (m+n) (f x) + +instance Monad M where + return x = Action 0 x + Action m x >>= f = let Action n y = f x in Action (m+n) y + +instance Foldable M where + foldMap f (Action _ x) = f x + +instance Traversable M where + traverse f (Action n x) = Action n <$> f x diff --git a/tests/set-properties.hs b/tests/set-properties.hs new file mode 100644 --- /dev/null +++ b/tests/set-properties.hs @@ -0,0 +1,341 @@ +import qualified Data.IntSet as IntSet +import Data.List (nub,sort) +import qualified Data.List as List +import Data.Monoid (mempty) +import Data.Set +import Prelude hiding (lookup, null, map, filter, foldr, foldl) +import Test.Framework +import Test.Framework.Providers.HUnit +import Test.Framework.Providers.QuickCheck2 +import Test.HUnit hiding (Test, Testable) +import Test.QuickCheck + +main :: IO () +main = defaultMainWithOpts [ testCase "lookupLT" test_lookupLT + , testCase "lookupGT" test_lookupGT + , testCase "lookupLE" test_lookupLE + , testCase "lookupGE" test_lookupGE + , testProperty "prop_Valid" prop_Valid + , testProperty "prop_Single" prop_Single + , testProperty "prop_Member" prop_Member + , testProperty "prop_NotMember" prop_NotMember + , testProperty "prop_LookupLT" prop_LookupLT + , testProperty "prop_LookupGT" prop_LookupGT + , testProperty "prop_LookupLE" prop_LookupLE + , testProperty "prop_LookupGE" prop_LookupGE + , testProperty "prop_InsertValid" prop_InsertValid + , testProperty "prop_InsertDelete" prop_InsertDelete + , testProperty "prop_DeleteValid" prop_DeleteValid + , testProperty "prop_Join" prop_Join + , testProperty "prop_Merge" prop_Merge + , testProperty "prop_UnionValid" prop_UnionValid + , testProperty "prop_UnionInsert" prop_UnionInsert + , testProperty "prop_UnionAssoc" prop_UnionAssoc + , testProperty "prop_UnionComm" prop_UnionComm + , testProperty "prop_DiffValid" prop_DiffValid + , testProperty "prop_Diff" prop_Diff + , testProperty "prop_IntValid" prop_IntValid + , testProperty "prop_Int" prop_Int + , testProperty "prop_Ordered" prop_Ordered + , testProperty "prop_List" prop_List + , testProperty "prop_DescList" prop_DescList + , testProperty "prop_AscDescList" prop_AscDescList + , testProperty "prop_fromList" prop_fromList + , testProperty "prop_isProperSubsetOf" prop_isProperSubsetOf + , testProperty "prop_isProperSubsetOf2" prop_isProperSubsetOf2 + , testProperty "prop_isSubsetOf" prop_isSubsetOf + , testProperty "prop_isSubsetOf2" prop_isSubsetOf2 + , testProperty "prop_size" prop_size + , testProperty "prop_findMax" prop_findMax + , testProperty "prop_findMin" prop_findMin + , testProperty "prop_ord" prop_ord + , testProperty "prop_readShow" prop_readShow + , testProperty "prop_foldR" prop_foldR + , testProperty "prop_foldR'" prop_foldR' + , testProperty "prop_foldL" prop_foldL + , testProperty "prop_foldL'" prop_foldL' + , testProperty "prop_map" prop_map + , testProperty "prop_maxView" prop_maxView + , testProperty "prop_minView" prop_minView + , testProperty "prop_split" prop_split + , testProperty "prop_splitMember" prop_splitMember + , testProperty "prop_partition" prop_partition + , testProperty "prop_filter" prop_filter + ] opts + where + opts = mempty { ropt_test_options = Just $ mempty { topt_maximum_generated_tests = Just 500 + , topt_maximum_unsuitable_generated_tests = Just 500 + } + } + +---------------------------------------------------------------- +-- Unit tests +---------------------------------------------------------------- + +test_lookupLT :: Assertion +test_lookupLT = do + lookupLT 3 (fromList [3, 5]) @?= Nothing + lookupLT 5 (fromList [3, 5]) @?= Just 3 + +test_lookupGT :: Assertion +test_lookupGT = do + lookupGT 4 (fromList [3, 5]) @?= Just 5 + lookupGT 5 (fromList [3, 5]) @?= Nothing + +test_lookupLE :: Assertion +test_lookupLE = do + lookupLE 2 (fromList [3, 5]) @?= Nothing + lookupLE 4 (fromList [3, 5]) @?= Just 3 + lookupLE 5 (fromList [3, 5]) @?= Just 5 + +test_lookupGE :: Assertion +test_lookupGE = do + lookupGE 3 (fromList [3, 5]) @?= Just 3 + lookupGE 4 (fromList [3, 5]) @?= Just 5 + lookupGE 6 (fromList [3, 5]) @?= Nothing + +{-------------------------------------------------------------------- + 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) + +{-------------------------------------------------------------------- + Valid tree's +--------------------------------------------------------------------} +forValid :: (Enum a,Show 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 + +forValidUnitTree :: Testable a => (Set Int -> a) -> Property +forValidUnitTree f = forValid f + +prop_Valid :: Property +prop_Valid = forValidUnitTree $ \t -> valid t + +{-------------------------------------------------------------------- + Single, Member, Insert, Delete +--------------------------------------------------------------------} +prop_Single :: Int -> Bool +prop_Single x = (insert x empty == singleton x) + +prop_Member :: [Int] -> Int -> Bool +prop_Member xs n = + let m = fromList xs + in all (\k -> k `member` m == (k `elem` xs)) (n : xs) + +prop_NotMember :: [Int] -> Int -> Bool +prop_NotMember xs n = + let m = fromList xs + in all (\k -> k `notMember` m == (k `notElem` xs)) (n : xs) + +test_LookupSomething :: (Int -> Set Int -> Maybe Int) -> (Int -> Int -> Bool) -> [Int] -> Bool +test_LookupSomething lookup' cmp xs = + let odd_sorted_xs = filter_odd $ nub $ sort xs + t = fromList odd_sorted_xs + test x = case List.filter (`cmp` x) odd_sorted_xs of + [] -> lookup' x t == Nothing + cs | 0 `cmp` 1 -> lookup' x t == Just (last cs) -- we want largest such element + | otherwise -> lookup' x t == Just (head cs) -- we want smallest such element + in all test xs + + where filter_odd [] = [] + filter_odd [_] = [] + filter_odd (_ : o : xs) = o : filter_odd xs + +prop_LookupLT :: [Int] -> Bool +prop_LookupLT = test_LookupSomething lookupLT (<) + +prop_LookupGT :: [Int] -> Bool +prop_LookupGT = test_LookupSomething lookupGT (>) + +prop_LookupLE :: [Int] -> Bool +prop_LookupLE = test_LookupSomething lookupLE (<=) + +prop_LookupGE :: [Int] -> Bool +prop_LookupGE = test_LookupSomething lookupGE (>=) + +prop_InsertValid :: Int -> Property +prop_InsertValid k = forValidUnitTree $ \t -> valid (insert k t) + +prop_InsertDelete :: Int -> Set Int -> Property +prop_InsertDelete k t = not (member k t) ==> delete k (insert k t) == t + +prop_DeleteValid :: Int -> Property +prop_DeleteValid k = forValidUnitTree $ \t -> valid (delete k (insert k t)) + +{-------------------------------------------------------------------- + Balance +--------------------------------------------------------------------} +prop_Join :: Int -> Property +prop_Join x = forValidUnitTree $ \t -> + let (l,r) = split x t + in valid (join x l r) + +prop_Merge :: Int -> Property +prop_Merge x = forValidUnitTree $ \t -> + let (l,r) = split x t + in valid (merge l r) + +{-------------------------------------------------------------------- + Union +--------------------------------------------------------------------} +prop_UnionValid :: Property +prop_UnionValid + = forValidUnitTree $ \t1 -> + forValidUnitTree $ \t2 -> + valid (union t1 t2) + +prop_UnionInsert :: Int -> Set Int -> Bool +prop_UnionInsert x t = union t (singleton x) == insert x t + +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_DiffValid :: Property +prop_DiffValid = forValidUnitTree $ \t1 -> + forValidUnitTree $ \t2 -> + valid (difference t1 t2) + +prop_Diff :: [Int] -> [Int] -> Bool +prop_Diff xs ys = toAscList (difference (fromList xs) (fromList ys)) + == List.sort ((List.\\) (nub xs) (nub ys)) + +prop_IntValid :: Property +prop_IntValid = forValidUnitTree $ \t1 -> + forValidUnitTree $ \t2 -> + valid (intersection t1 t2) + +prop_Int :: [Int] -> [Int] -> Bool +prop_Int xs ys = toAscList (intersection (fromList xs) (fromList ys)) + == List.sort (nub ((List.intersect) (xs) (ys))) + +{-------------------------------------------------------------------- + Lists +--------------------------------------------------------------------} +prop_Ordered :: Property +prop_Ordered = forAll (choose (5,100)) $ \n -> + let xs = [0..n::Int] + in fromAscList xs == fromList xs + +prop_List :: [Int] -> Bool +prop_List xs = (sort (nub xs) == toList (fromList xs)) + +prop_DescList :: [Int] -> Bool +prop_DescList xs = (reverse (sort (nub xs)) == toDescList (fromList xs)) + +prop_AscDescList :: [Int] -> Bool +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 + nub_sort_xs = List.map List.head $ List.group sort_xs + +{-------------------------------------------------------------------- + Set operations are like IntSet operations +--------------------------------------------------------------------} +toIntSet :: Set Int -> IntSet.IntSet +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) + +-- 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 + 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_isSubsetOf2 :: Set Int -> Set Int -> Bool +prop_isSubsetOf2 a b = isSubsetOf a (union a b) + +prop_size :: Set Int -> Bool +prop_size s = size s == List.length (toList s) + +prop_findMax :: Set Int -> Property +prop_findMax s = not (null s) ==> findMax s == maximum (toList s) + +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_readShow :: Set Int -> Bool +prop_readShow s = s == read (show s) + +prop_foldR :: Set Int -> Bool +prop_foldR s = foldr (:) [] s == toList s + +prop_foldR' :: Set Int -> Bool +prop_foldR' s = foldr' (:) [] s == toList s + +prop_foldL :: Set Int -> Bool +prop_foldL s = foldl (flip (:)) [] s == List.foldl (flip (:)) [] (toList s) + +prop_foldL' :: Set Int -> Bool +prop_foldL' s = foldl' (flip (:)) [] s == List.foldl' (flip (:)) [] (toList s) + +prop_map :: Set Int -> Bool +prop_map s = map id s == s + +prop_maxView :: Set Int -> Bool +prop_maxView s = case maxView s of + Nothing -> null s + Just (m,s') -> m == maximum (toList s) && s == insert m s' && m `notMember` s' + +prop_minView :: Set Int -> Bool +prop_minView s = case minView s of + Nothing -> null s + Just (m,s') -> m == minimum (toList s) && s == insert m s' && m `notMember` s' + +prop_split :: Set Int -> Int -> Bool +prop_split s i = case split i s of + (s1,s2) -> all (i) (toList s2) && i `delete` s == union s1 s2 + +prop_splitMember :: Set Int -> Int -> Bool +prop_splitMember s i = case splitMember i s of + (s1,t,s2) -> all (i) (toList s2) && t == i `member` s && i `delete` s == union s1 s2 + +prop_partition :: Set Int -> Int -> Bool +prop_partition s i = case partition odd s of + (s1,s2) -> all odd (toList s1) && all even (toList s2) && s == s1 `union` s2 + +prop_filter :: Set Int -> Int -> Bool +prop_filter s i = partition odd s == (filter odd s, filter even s)